LCORE-693: Added rag_chunks to streaming_query #585
Conversation
Walkthroughstream_end_event now accepts a TurnSummary, derives rag_chunks and referenced_documents via new utils.endpoints helpers, and includes them in the SSE end payload; query endpoint uses response-provided referenced_documents; new utils/endpoints functions normalize, deduplicate, validate, and serialize RAG results; tests updated/added. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant API as streaming_query.py
participant Gen as ResponseGenerator
participant Utils as utils.endpoints
participant Store as TranscriptStore
Note over API,Utils: New helpers compute/serialize rag_chunks & referenced_documents
Client->>API: start streaming request
API->>Gen: generate tokens / partial responses
Gen-->>API: token/event chunks
API-->>Client: SSE partial events
alt successful completion
Gen->>API: summary (TurnSummary)
API->>Utils: create_rag_chunks_dict(summary)
API->>Utils: create_referenced_documents_with_metadata(summary, metadata_map)
API->>Store: store_transcript(..., rag_chunks=summary.rag_chunks)
Store-->>API: ack
API-->>Client: SSE end payload (includes rag_chunks, referenced_documents, token counts)
else error
API-->>Client: SSE error event
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (2)
🧰 Additional context used📓 Path-based instructions (5)src/**/*.py📄 CodeRabbit inference engine (CLAUDE.md)
Files:
src/app/**/*.py📄 CodeRabbit inference engine (CLAUDE.md)
Files:
**/*.py📄 CodeRabbit inference engine (CLAUDE.md)
Files:
src/{app/**/*.py,client.py}📄 CodeRabbit inference engine (CLAUDE.md)
Files:
src/app/endpoints/**/*.py📄 CodeRabbit inference engine (CLAUDE.md)
Files:
🧬 Code graph analysis (2)src/app/endpoints/streaming_query.py (2)
src/utils/endpoints.py (2)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
src/app/endpoints/streaming_query.py (4)
138-153: Document the new parameter and tighten types in signature.Add the summary param to the docstring and use a precise type for metadata_map.
-def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str: +def stream_end_event(metadata_map: dict[str, dict[str, Any]], summary: TurnSummary) -> str: """ Yield the end of the data stream. - Format and return the end event for a streaming response, - including referenced document metadata and placeholder token - counts. + Format and return the end event for a streaming response, + including referenced document metadata, RAG chunks from the turn + summary, and placeholder token counts. Parameters: - metadata_map (dict): A mapping containing metadata about - referenced documents. + metadata_map (dict[str, dict[str, Any]]): Mapping of document_id -> metadata + (expects at least 'title' and 'docs_url' when available). + summary (TurnSummary): Summary of the turn containing tool calls and + extracted rag_chunks. Returns: str: A Server-Sent Events (SSE) formatted string representing the end of the data stream. """
154-163: DRY: factor rag_chunks serialization into a helper.The same comprehension is duplicated. Extract a helper and reuse it.
Apply within these ranges:
- rag_chunks = [ - { - "content": chunk.content, - "source": chunk.source, - "score": chunk.score - } - for chunk in summary.rag_chunks - ] + rag_chunks = serialize_rag_chunks(summary.rag_chunks)- rag_chunks_for_transcript = [ - { - "content": chunk.content, - "source": chunk.source, - "score": chunk.score - } - for chunk in summary.rag_chunks - ] + rag_chunks_for_transcript = serialize_rag_chunks(summary.rag_chunks)Add this helper once at module scope (outside the changed ranges):
def serialize_rag_chunks(chunks: list[RAGChunk]) -> list[dict[str, Any]]: """Serialize RAGChunk models into JSON-safe dicts.""" return [ {"content": c.content, "source": c.source, "score": c.score} for c in chunks ]Also applies to: 712-720
186-198: Update example to reflect new end payload.The SSE example in streaming_query_responses doesn't include rag_chunks. Update it to avoid API doc drift.
679-681: TurnSummary default rag_chunks may be a shared mutable.In utils/types.py, rag_chunks is defined with a bare [] default. Prefer Field(default_factory=list) to avoid shared state across instances.
Do you want a follow-up patch to adjust TurnSummary.rag_chunks to use default_factory=list?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/app/endpoints/streaming_query.py(5 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/streaming_query.py
🧬 Code graph analysis (1)
src/app/endpoints/streaming_query.py (3)
src/models/responses.py (2)
RAGChunk(37-42)ReferencedDocument(53-65)src/utils/types.py (1)
TurnSummary(61-135)src/app/endpoints/query.py (1)
is_transcripts_enabled(90-96)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: e2e_tests
🔇 Additional comments (3)
src/app/endpoints/streaming_query.py (3)
32-32: Confirm necessity of new imports.RAGChunk and ReferencedDocument are imported but not referenced directly in this module. If unused, consider removing to satisfy linters.
707-707: LGTM: summary passed into stream_end_event.Signature usage looks correct.
721-734: store_transcript rag_chunks shape verified
store_transcript’s signature expects a list[dict] and streaming_query builds rag_chunks_for_transcript as dicts with content, source, and score.
bsatapat-jpg
force-pushed
the
stream_query
branch
from
a417236 to
dc04bf4
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
src/app/endpoints/streaming_query.py (2)
147-153: Docstring missingsummaryparameter.The signature now requires
summary, but the docstring still documents onlymetadata_map. Please keep the docstring in sync so future readers don’t miss the new contract.Apply this diff to update the docstring:
Parameters: metadata_map (dict): A mapping containing metadata about referenced documents. + summary (TurnSummary): Summary of the completed turn, including any collected RAG chunks.
736-743: Consider reusing the serialized RAG chunks.We serialize
summary.rag_chunkstwice (here and instream_end_event). Extracting a small helper (or returning the serialized list fromstream_end_event) would keep the representations aligned and avoid drift.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/app/endpoints/streaming_query.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/streaming_query.py
🧬 Code graph analysis (1)
src/app/endpoints/streaming_query.py (2)
src/utils/types.py (1)
TurnSummary(61-135)src/app/endpoints/query.py (1)
is_transcripts_enabled(90-96)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: e2e_tests
🔇 Additional comments (1)
src/app/endpoints/streaming_query.py (1)
172-207: RAG metadata dedupe looks solid.Thanks for enriching from
metadata_mapwhile deduping by both document IDs and URLs—this addresses the duplication edge cases we had flagged earlier.
|
@bsatapat-jpg Most of the logic very similar to what we have in query endpoint. Please check if common logic related to chunk retrieval/ref_doc creation can be moved to an util.. this way it will be easy to maintain. |
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/app/endpoints/streaming_query.py (3)
145-160: Update docstring to document the newsummaryparameter.The function signature now includes a
summary: TurnSummaryparameter, but the docstring's Parameters section does not document it.Apply this diff to update the docstring:
def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str: """ Yield the end of the data stream. Format and return the end event for a streaming response, including referenced document metadata and placeholder token counts. Parameters: metadata_map (dict): A mapping containing metadata about referenced documents. + summary (TurnSummary): The turn summary containing LLM response, + tool calls, and RAG chunks. Returns: str: A Server-Sent Events (SSE) formatted string representing the end of the data stream. """
742-760: Eliminate duplicate RAG chunk serialization logic.The conversion of
summary.rag_chunksto a list of dicts (lines 743-750) duplicates the identical logic from lines 162-169. This violates the DRY principle.Consider extracting this transformation into a helper function or reusing the
rag_chunksvariable created earlier:Option 1: Reuse the existing variable (simpler)
Since
rag_chunksis already created at line 162-169 withinstream_end_event, you could return it or pass it through. However, this requires refactoring the function signature.Option 2: Extract to helper function (cleaner)
def _serialize_rag_chunks(chunks: list[RAGChunk]) -> list[dict[str, Any]]: """Convert RAG chunks to serializable format. Args: chunks: List of RAG chunks to serialize. Returns: List of dictionaries containing chunk content, source, and score. """ return [ { "content": chunk.content, "source": chunk.source, "score": chunk.score } for chunk in chunks ]Then use it in both places:
- rag_chunks = [ - { - "content": chunk.content, - "source": chunk.source, - "score": chunk.score - } - for chunk in summary.rag_chunks - ] + rag_chunks = _serialize_rag_chunks(summary.rag_chunks)And:
- # Convert RAG chunks to serializable format for store_transcript - rag_chunks_for_transcript = [ - { - "content": chunk.content, - "source": chunk.source, - "score": chunk.score - } - for chunk in summary.rag_chunks - ] + # Convert RAG chunks to serializable format for store_transcript + rag_chunks_for_transcript = _serialize_rag_chunks(summary.rag_chunks)
1-1: Fix Black formatting violations.The pipeline failure indicates that this file has formatting issues that violate Black's style requirements.
Run the following command to fix the formatting:
uv tool run black src/app/endpoints/streaming_query.pyTo check what would change without modifying the file:
uv tool run black --check --diff src/app/endpoints/streaming_query.py
🧹 Nitpick comments (2)
src/app/endpoints/streaming_query.py (2)
171-214: Align streaming_query referenced_docs with query endpoint and DRY extraction logic
- In
streaming_query.py, wrap HTTPdoc_urlvalues withAnyUrl(chunk.source)(as inquery.pylines 371–374) for consistent URL validation.- Extract the duplicated referenced‐document parsing (from both
query.pyandstreaming_query.py) into a sharedparse_referenced_documentsutility.
161-214: Extract RAG processing into a shared utilityThe inline chunk-serialization and referenced-documents logic in streaming_query.py largely mirrors patterns in query.py (e.g. serializing
summary.rag_chunks, deduping/enriching sources). Introduce asrc/utils/rag_processing.pywith:
serialize_rag_chunks(chunks: Sequence[RAGChunk]) -> list[dict[str, Any]]build_referenced_documents_from_chunks(chunks: Sequence[RAGChunk], metadata_map: dict[str, Any]) -> list[dict[str, str | None]]Then replace the manual loops in both streaming_query.py and query.py (alongside, or in place of, the existing
parse_referenced_documents) with calls to these helpers to eliminate duplication and centralize future updates.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/app/endpoints/streaming_query.py(4 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/streaming_query.py
🧬 Code graph analysis (1)
src/app/endpoints/streaming_query.py (2)
src/utils/types.py (1)
TurnSummary(89-163)src/app/endpoints/query.py (1)
is_transcripts_enabled(93-99)
🪛 GitHub Actions: Black
src/app/endpoints/streaming_query.py
[error] 1-1: Black formatting check failed. 1 file would be reformatted. Run 'uv tool run black --check .' to see details or 'uv tool run black .' to reformat.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
🔇 Additional comments (2)
src/app/endpoints/streaming_query.py (2)
220-221: LGTM! RAG chunks successfully added to streaming response.The end event now correctly includes both
rag_chunksand the enrichedreferenced_documents, fulfilling the PR objective.
737-737: LGTM! Correct invocation of updated function.The call to
stream_end_eventcorrectly passes bothmetadata_mapandsummaryparameters.
bsatapat-jpg
force-pushed
the
stream_query
branch
from
73dbba8 to
e658012
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
src/app/endpoints/streaming_query.py (2)
146-161: Update the docstring to document the new parameter.The function signature now accepts a
summary: TurnSummaryparameter, but the docstring doesn't document it.Apply this diff to update the docstring:
def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str: """ Yield the end of the data stream. Format and return the end event for a streaming response, including referenced document metadata and placeholder token counts. Parameters: metadata_map (dict): A mapping containing metadata about referenced documents. + summary (TurnSummary): The turn summary containing LLM response, + tool calls, and RAG chunks. Returns: str: A Server-Sent Events (SSE) formatted string representing the end of the data stream. """
1-1: Run Black formatter to fix formatting issues.The pipeline failure indicates that Black formatting check failed. After addressing the line length issues, run
black src/app/endpoints/streaming_query.pyto ensure proper formatting.black src/app/endpoints/streaming_query.py
🧹 Nitpick comments (1)
src/app/endpoints/streaming_query.py (1)
743-761: Consider extracting the RAG chunk transformation into a helper function.The RAG chunk transformation at lines 744-751 is duplicated from lines 163-170. This is a simple transformation, but extracting it into a helper would improve maintainability and align with the PR feedback about moving common logic into utilities.
For example, add a helper function at module level:
def _format_rag_chunks(rag_chunks: list) -> list[dict[str, Any]]: """Convert RAG chunks to serializable dictionary format. Args: rag_chunks: List of RAG chunk objects. Returns: List of dictionaries with content, source, and score fields. """ return [ { "content": chunk.content, "source": chunk.source, "score": chunk.score } for chunk in rag_chunks ]Then use it in both locations:
# In stream_end_event rag_chunks = _format_rag_chunks(summary.rag_chunks) # In response_generator rag_chunks_for_transcript = _format_rag_chunks(summary.rag_chunks)Based on PR feedback from asamal4 about moving common logic into utilities.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/app/endpoints/streaming_query.py(5 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/streaming_query.py
🧬 Code graph analysis (1)
src/app/endpoints/streaming_query.py (2)
src/utils/types.py (1)
TurnSummary(89-163)src/app/endpoints/query.py (1)
is_transcripts_enabled(93-99)
🪛 GitHub Actions: Black
src/app/endpoints/streaming_query.py
[error] 1-1: Black formatting check failed. 1 file would be reformatted. Run 'black' to fix formatting.
🪛 GitHub Actions: Python linter
src/app/endpoints/streaming_query.py
[error] 204-204: pylint: Line too long (104/100) (C0301) in streaming_query.py
[error] 215-215: pylint: Line too long (129/100) (C0301) in streaming_query.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
🔇 Additional comments (3)
src/app/endpoints/streaming_query.py (3)
11-11: LGTM!The AnyUrl import is necessary for URL validation in the referenced documents construction, aligning with the query endpoint pattern.
163-170: LGTM!The RAG chunks are correctly extracted from the summary and formatted for the response.
738-738: LGTM!The call to
stream_end_eventcorrectly passes both themetadata_mapandsummaryparameters.
bsatapat-jpg
force-pushed
the
stream_query
branch
2 times, most recently
from
824f112 to
09d33f6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/app/endpoints/query.py (2)
25-25: Remove unused import.The
AnyUrlimport is no longer used after refactoring to usecreate_referenced_documents_from_chunks. This is causing a pipeline failure.-from pydantic import AnyUrl
364-374: Fix variable name mismatch - refactoring is incomplete.Line 365 assigns to
referenced_docsbut line 373 usesreferenced_documents(from theretrieve_responsereturn value at line 284). The new utility function call is dead code, and the old logic inretrieve_responseis still being used. This makes the refactoring incomplete.Apply this diff to complete the refactoring:
logger.info("Extracting referenced documents...") - referenced_docs = create_referenced_documents_from_chunks(summary.rag_chunks) + referenced_documents = create_referenced_documents_from_chunks(summary.rag_chunks) logger.info("Building final response...") response = QueryResponse( conversation_id=conversation_id, response=summary.llm_response, rag_chunks=summary.rag_chunks if summary.rag_chunks else [], tool_calls=tool_calls if tool_calls else None, referenced_documents=referenced_documents, )Also consider removing the
referenced_documentsreturn value fromretrieve_responsesince it's no longer needed with the new utility.src/app/endpoints/streaming_query.py (1)
147-162: Update docstring to document new parameter.The function signature now accepts a
summary: TurnSummaryparameter, but the docstring doesn't document it. Update the Parameters section to include this new argument.def stream_end_event(metadata_map: dict, summary: TurnSummary) -> str: """ Yield the end of the data stream. Format and return the end event for a streaming response, including referenced document metadata and placeholder token counts. Parameters: metadata_map (dict): A mapping containing metadata about referenced documents. + summary (TurnSummary): Summary of the turn containing RAG chunks + and tool calls. Returns: str: A Server-Sent Events (SSE) formatted string representing the end of the data stream. """
🧹 Nitpick comments (2)
src/utils/endpoints.py (1)
445-481: Evaluate consolidating duplicate logic.Both
create_referenced_documents_from_chunksandcreate_referenced_documents_with_metadatashare overlapping logic for processing RAG chunks, handling HTTP URLs vs document IDs, and deduplication. Consider extracting common patterns into a shared helper to improve maintainability.Based on PR objectives.
Also applies to: 365-442
src/app/endpoints/streaming_query.py (1)
696-703: Consider reusing the utility function.The manual construction of
rag_chunks_for_transcriptduplicates the logic increate_rag_chunks_dict. You could simplify this by reusing the utility:- # Convert RAG chunks to serializable format for store_transcript - rag_chunks_for_transcript = [ - { - "content": chunk.content, - "source": chunk.source, - "score": chunk.score, - } - for chunk in summary.rag_chunks - ] + # Convert RAG chunks to serializable format for store_transcript + rag_chunks_for_transcript = create_rag_chunks_dict(summary)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(5 hunks)src/utils/endpoints.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/utils/endpoints.pysrc/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/utils/endpoints.pysrc/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
🧬 Code graph analysis (3)
src/utils/endpoints.py (2)
src/models/responses.py (1)
ReferencedDocument(55-67)src/utils/types.py (1)
TurnSummary(89-163)
src/app/endpoints/query.py (1)
src/utils/endpoints.py (1)
create_referenced_documents_from_chunks(445-481)
src/app/endpoints/streaming_query.py (3)
src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(365-442)create_rag_chunks_dict(349-362)src/utils/types.py (1)
TurnSummary(89-163)src/app/endpoints/query.py (1)
is_transcripts_enabled(94-100)
🪛 GitHub Actions: Ruff
src/app/endpoints/query.py
[error] 25-25: F401: 'pydantic.AnyUrl' imported but unused. Remove unused import.
[error] 365-365: F841: Local variable 'referenced_docs' assigned to but never used. Remove assignment or use the variable.
🔇 Additional comments (5)
src/utils/endpoints.py (2)
4-4: LGTM! Imports are properly structured.The new imports follow the absolute import pattern and all are used by the new utility functions.
Also applies to: 8-8, 13-13, 19-19
349-362: LGTM! Clean utility function.The function is well-documented and provides a clear transformation from RAGChunk objects to dictionaries for streaming responses.
src/app/endpoints/query.py (1)
48-48: LGTM! Utility import is correctly added.The import follows the absolute import pattern and is used in the refactored logic.
src/app/endpoints/streaming_query.py (2)
50-51: LGTM! Utility imports are properly added.The new utility functions are correctly imported and used in
stream_end_event.
690-690: LGTM! Correct integration of summary parameter.The summary is now properly passed to
stream_end_event, enabling RAG chunk and referenced document extraction.
bsatapat-jpg
force-pushed
the
stream_query
branch
4 times, most recently
from
b971bf8 to
eae275a
Compare
are-ces
left a comment
are-ces
left a comment
There was a problem hiding this comment.
Any reason why we need two separate functions to return referenced docs? Could we unify them into a single function for both streaming and query endpoints? I see some duplicated logic here (as @asamal4 already pointed out).
Also, please add unit tests to cover the new functionality. Thank you!
bsatapat-jpg
changed the title
bsatapat-jpg
force-pushed
the
stream_query
branch
from
eae275a to
17fc0d3
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/app/endpoints/streaming_query.py (1)
147-176: Update docstring to cover newsummaryparameter.
stream_end_eventnow acceptssummary, but the docstring still documents onlymetadata_map. Please add the missingArgsentry (including type and purpose) to keep docs aligned with the function signature.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(5 hunks)src/utils/endpoints.py(2 hunks)tests/unit/utils/test_endpoints.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- tests/unit/utils/test_endpoints.py
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/utils/endpoints.pysrc/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/utils/endpoints.pysrc/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
🧬 Code graph analysis (3)
src/utils/endpoints.py (2)
src/models/responses.py (1)
ReferencedDocument(55-67)src/utils/types.py (1)
TurnSummary(89-163)
src/app/endpoints/query.py (1)
src/utils/endpoints.py (1)
create_referenced_documents_from_chunks(503-511)
src/app/endpoints/streaming_query.py (2)
src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(490-500)create_rag_chunks_dict(349-362)src/utils/types.py (1)
TurnSummary(89-163)
🪛 GitHub Actions: Black
src/utils/endpoints.py
[error] 1-1: Black formatting required. 3 files would be reformatted by running 'black'.
src/app/endpoints/streaming_query.py
[error] 1-1: Black formatting required. 3 files would be reformatted by running 'black'.
🪛 GitHub Actions: Pyright
src/utils/endpoints.py
[error] 419-419: pyright: Argument of type "dict[str, AnyUrl | str | Unknown | None]" cannot be assigned to parameter "object" of type "ReferencedDocument" in function "append" (reportArgumentType)
[error] 425-425: pyright: Argument of type "AnyUrl | str | None" cannot be assigned to parameter "doc_url" of type "AnyUrl | None" in function "init"
[error] 454-454: pyright: Argument of type "dict[str, AnyUrl | str | Any | Unknown | None]" cannot be assigned to parameter "object" of type "ReferencedDocument" in function "append" (reportArgumentType)
[error] 460-460: pyright: Argument of type "AnyUrl | str | None" cannot be assigned to parameter "doc_url" of type "AnyUrl | None" in function "init"
[error] 480-480: pyright: Argument of type "dict[str, AnyUrl | str | Any | Unknown | None]" cannot be assigned to parameter "object" of type "ReferencedDocument" in function "append" (reportArgumentType)
[error] 486-486: pyright: Argument of type "AnyUrl | str | None" cannot be assigned to parameter "doc_url" of type "AnyUrl | None" in function "init"
[error] 502-502: pyright: Type "list[ReferencedDocument] | list[dict[str, str | None]]" is not assignable to return type "list[ReferencedDocument]"
[error] 515-515: pyright: Type "list[ReferencedDocument] | list[dict[str, str | None]]" is not assignable to return type "list[ReferencedDocument]"
🪛 GitHub Actions: Python linter
src/utils/endpoints.py
[error] 384-384: Line too long (112/100) (line-too-long)
[error] 392-392: Trailing whitespace (trailing-whitespace)
[error] 408-408: Trailing whitespace (trailing-whitespace)
[error] 417-417: Trailing whitespace (trailing-whitespace)
[error] 433-433: Trailing whitespace (trailing-whitespace)
[error] 437-437: Trailing whitespace (trailing-whitespace)
[error] 442-442: Trailing whitespace (trailing-whitespace)
[error] 446-446: Line too long (107/100) (line-too-long)
[error] 450-450: Trailing whitespace (trailing-whitespace)
[error] 452-452: Trailing whitespace (trailing-whitespace)
[error] 474-474: Line too long (107/100) (line-too-long)
[error] 478-478: Trailing whitespace (trailing-whitespace)
[error] 513-513: Line too long (101/100) (line-too-long)
[error] 369-369: Too many statements (63/50) (too-many-statements)
src/app/endpoints/streaming_query.py
[error] 168-168: Trailing whitespace (trailing-whitespace)
🪛 GitHub Actions: Type checks
src/utils/endpoints.py
[error] 420-420: Mypy error: Dict entry 0 has incompatible type "str": "AnyUrl | str | None"; expected "str": "str | None" [dict-item]
[error] 455-455: Mypy error: Dict entry 0 has incompatible type "str": "AnyUrl | str | None"; expected "str": "str | None" [dict-item]
[error] 481-481: Mypy error: Dict entry 0 has incompatible type "str": "AnyUrl | str | None"; expected "str": "str | None" [dict-item]
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
bsatapat-jpg
force-pushed
the
stream_query
branch
4 times, most recently
from
59eac14 to
b79d444
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/utils/endpoints.py (1)
382-417: Consider simplifying type checking with TypeGuard or stricter typing.The runtime type checks on lines 398-401 suggest that
metadata_mapvalues might have unpredictable types. While defensive, this could indicate a need for stricter typing upstream or usingTypeGuardto narrow types more elegantly.If metadata structure is predictable, consider defining a TypedDict:
from typing import TypedDict class DocumentMetadata(TypedDict, total=False): docs_url: str title: str # other fields... # Then in function signature: def _process_document_id( src: str, doc_ids: set[str], doc_urls: set[str], metas_by_id: dict[str, DocumentMetadata], metadata_map: dict[str, Any] | None, ) -> tuple[AnyUrl | None, str] | None:This would catch type issues earlier and eliminate runtime checks.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
src/app/endpoints/query.py(1 hunks)src/app/endpoints/streaming_query.py(5 hunks)src/utils/endpoints.py(2 hunks)tests/unit/utils/test_endpoints.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/query.pysrc/utils/endpoints.pysrc/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/query.pysrc/utils/endpoints.pytests/unit/utils/test_endpoints.pysrc/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
tests/{unit,integration}/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/{unit,integration}/**/*.py: Use pytest for all unit and integration tests
Do not use unittest in tests; pytest is the standard
Files:
tests/unit/utils/test_endpoints.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Use pytest-mock to create AsyncMock objects for async interactions in tests
Use the shared auth mock constant: MOCK_AUTH = ("mock_user_id", "mock_username", False, "mock_token") in tests
Files:
tests/unit/utils/test_endpoints.py
🧬 Code graph analysis (3)
src/utils/endpoints.py (2)
src/models/responses.py (1)
ReferencedDocument(55-67)src/utils/types.py (1)
TurnSummary(89-163)
tests/unit/utils/test_endpoints.py (1)
src/utils/endpoints.py (1)
create_referenced_documents(491-525)
src/app/endpoints/streaming_query.py (2)
src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(529-543)create_rag_chunks_dict(349-362)src/utils/types.py (1)
TurnSummary(89-163)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
🔇 Additional comments (8)
tests/unit/utils/test_endpoints.py (1)
847-948: LGTM! Comprehensive test coverage for the unified document creation workflow.The test suite thoroughly covers all key scenarios: empty input, HTTP URLs, metadata enrichment, filtering tool names, empty source handling, deduplication, and invalid URL handling. The tests properly verify the behavior of
create_referenced_documentsand will catch regressions.src/app/endpoints/query.py (1)
362-370: LGTM! Clean refactor to use response-provided documents.The change correctly delegates document creation to
retrieve_response, which now returnsreferenced_documentsas part of its tuple (line 282). This eliminates duplicate document construction logic in the query handler and aligns with the unified approach introduced in this PR.src/app/endpoints/streaming_query.py (1)
50-51: LGTM! Proper use of utility functions for consistent document processing.The streaming endpoint now uses
create_rag_chunks_dictandcreate_referenced_documents_with_metadatato derive rag_chunks and referenced_documents, ensuring consistency with the query endpoint. The conversion to dictionary format (lines 170-176) is necessary for JSON serialization in the SSE response.Also applies to: 163-176
src/utils/endpoints.py (5)
349-362: LGTM! Simple and clear utility function.
create_rag_chunks_dictcorrectly extracts content, source, and score from each RAG chunk and returns a list of dictionaries suitable for JSON serialization in streaming responses.
365-379: LGTM! Proper URL validation with error handling.
_process_http_sourcecorrectly validates HTTP URLs using AnyUrl within a try/except block, logging warnings for invalid URLs and gracefully continuing processing. This prevents crashes from malformed URLs while maintaining useful diagnostics.
448-488: LGTM! Well-structured core processing pipeline.
_process_rag_chunks_for_documentsimplements a clean processing flow:
- Filters out tool names and empty sources
- Processes HTTP URLs and document IDs separately
- Enriches with metadata when available
- Deduplicates by both doc_urls and doc_ids
- Adds additional metadata entries
The separation into helper functions improves readability and testability.
491-525: LGTM! Clean unified API with flexible return types.
create_referenced_documentsprovides a single entry point with:
- Optional metadata enrichment via
metadata_map- Flexible return format via
return_dict_formatflag- Proper delegation to core processing logic
- Clear type annotations
The design allows both endpoints (query and streaming) to use the same processing pipeline while getting their preferred return type.
529-559: LGTM! Backward compatibility wrappers maintain existing APIs.The wrapper functions
create_referenced_documents_with_metadataandcreate_referenced_documents_from_chunksmaintain backward compatibility while delegating to the unified processing function. This allows gradual migration without breaking existing code.
bsatapat-jpg
force-pushed
the
stream_query
branch
from
b79d444 to
e42019a
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/app/endpoints/query.py (1)
362-371: Harden URL parsing in parse_metadata_from_text_item to avoid crashesReferencedDocument(doc_url=...) relies on AnyUrl validation via Pydantic; malformed URLs in tool metadata will raise ValidationError and 500 the request. Wrap validation and degrade to None.
Apply:
+from pydantic import AnyUrl, ValidationError @@ - if url and title: - return ReferencedDocument(doc_url=url, doc_title=title) + if url and title: + try: + return ReferencedDocument(doc_url=AnyUrl(url), doc_title=title) + except ValidationError: + logger.warning("Invalid URL in metadata block: %s", url) + return ReferencedDocument(doc_url=None, doc_title=title)Optional cleanup outside this hunk:
- Remove the unused variable docs in parse_metadata_from_text_item.
- Update the Returns docstring to Optional[ReferencedDocument].
🧹 Nitpick comments (5)
src/utils/endpoints.py (2)
382-388: Simplify helper signature by dropping metadata_map parametermetadata_map is only used as a boolean guard; metas_by_id already captures filtered metadata. Remove the extra parameter and branch to reduce complexity and call-site noise.
Apply this minimal change:
-def _process_document_id( - src: str, - doc_ids: set[str], - doc_urls: set[str], - metas_by_id: dict[str, dict[str, Any]], - metadata_map: dict[str, Any] | None, -) -> tuple[AnyUrl | None, str] | None: +def _process_document_id( + src: str, + doc_ids: set[str], + doc_urls: set[str], + metas_by_id: dict[str, dict[str, Any]], +) -> tuple[AnyUrl | None, str] | None: @@ - meta = metas_by_id.get(src, {}) if metadata_map else {} + meta = metas_by_id.get(src, {})And adjust the call site:
- entry = _process_document_id( - src, doc_ids, doc_urls, metas_by_id, metadata_map - ) + entry = _process_document_id(src, doc_ids, doc_urls, metas_by_id)Also applies to: 477-479
347-347: Docstring return tuple mismatchget_temp_agent returns (agent, session_id, conversation_id) but the docstring says only (agent, session_id). Update the Returns section for accuracy.
Outside this hunk, update the docstring:
Returns: tuple[AsyncAgent, str, str]: (agent, session_id, conversation_id).tests/unit/utils/test_endpoints.py (1)
847-949: Add coverage for dict-return branchcreate_referenced_documents has a return_dict_format=True mode that isn’t tested. Add a test to assert dict shape and URL stringification.
You can append:
def test_create_referenced_documents_dict_format(): mock_chunk = type("MockChunk", (), {"source": "https://example.com/doc"})() res = endpoints.create_referenced_documents([mock_chunk], return_dict_format=True) assert isinstance(res, list) and isinstance(res[0], dict) assert res[0]["doc_url"] == str(AnyUrl("https://example.com/doc")) assert res[0]["doc_title"] == "doc"tests/unit/app/endpoints/test_streaming_query.py (1)
341-347: Optionally assert rag_chunks presence in SSE end eventYou already validate referenced_documents; consider asserting d["data"]["rag_chunks"] too to lock the new contract.
src/app/endpoints/streaming_query.py (1)
147-176: Update docstring to include summary parameter and new fieldsstream_end_event now depends on summary and emits rag_chunks and referenced_documents; docstring should list summary in Args and mention both fields in the payload description.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
src/app/endpoints/query.py(1 hunks)src/app/endpoints/streaming_query.py(5 hunks)src/utils/endpoints.py(2 hunks)tests/unit/app/endpoints/test_streaming_query.py(1 hunks)tests/unit/utils/test_endpoints.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (7)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
tests/unit/utils/test_endpoints.pytests/unit/app/endpoints/test_streaming_query.pysrc/app/endpoints/streaming_query.pysrc/app/endpoints/query.pysrc/utils/endpoints.py
tests/{unit,integration}/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/{unit,integration}/**/*.py: Use pytest for all unit and integration tests
Do not use unittest in tests; pytest is the standard
Files:
tests/unit/utils/test_endpoints.pytests/unit/app/endpoints/test_streaming_query.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Use pytest-mock to create AsyncMock objects for async interactions in tests
Use the shared auth mock constant: MOCK_AUTH = ("mock_user_id", "mock_username", False, "mock_token") in tests
Files:
tests/unit/utils/test_endpoints.pytests/unit/app/endpoints/test_streaming_query.py
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/streaming_query.pysrc/app/endpoints/query.pysrc/utils/endpoints.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/streaming_query.pysrc/app/endpoints/query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/streaming_query.pysrc/app/endpoints/query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/streaming_query.pysrc/app/endpoints/query.py
🧬 Code graph analysis (3)
tests/unit/utils/test_endpoints.py (1)
src/utils/endpoints.py (1)
create_referenced_documents(491-525)
src/app/endpoints/streaming_query.py (2)
src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(529-543)create_rag_chunks_dict(349-362)src/utils/types.py (1)
TurnSummary(89-163)
src/utils/endpoints.py (2)
src/models/responses.py (1)
ReferencedDocument(55-67)src/utils/types.py (1)
TurnSummary(89-163)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
🔇 Additional comments (3)
src/utils/endpoints.py (1)
349-362: Good addition: single place to serialize RAG chunks for SSE/transcriptsCentralizing to create_rag_chunks_dict keeps streaming and transcripts consistent. LGTM.
tests/unit/app/endpoints/test_streaming_query.py (1)
377-383: Nice: transcripts now get rag_chunks explicitlyPassing create_rag_chunks_dict(summary) ensures persistence of chunks. LGTM.
src/app/endpoints/streaming_query.py (1)
699-714: Good: end-event and transcript use shared serializersUsing create_rag_chunks_dict and create_referenced_documents_with_metadata ensures consistency and dedupe logic identical to query path. LGTM.
|
What has been done: |
bsatapat-jpg
force-pushed
the
stream_query
branch
from
e42019a to
dfad849
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/app/endpoints/streaming_query.py (1)
147-190: Implementation correctly uses new utility functions.The updated
stream_end_eventproperly:
- Accepts
summaryto extract RAG chunks- Uses
create_rag_chunks_dictandcreate_referenced_documents_with_metadatato centralize logic- Converts
ReferencedDocumentobjects to dicts for JSON serializationThe manual dict conversion (lines 170-176) explicitly handles
AnyUrlto string conversion, which is necessary for JSON serialization.Optional refinement: Consider using
doc.model_dump()if it handlesAnyUrlserialization correctly, which could simplify the conversion logic.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
src/app/endpoints/query.py(1 hunks)src/app/endpoints/streaming_query.py(6 hunks)src/utils/endpoints.py(2 hunks)tests/unit/app/endpoints/test_streaming_query.py(2 hunks)tests/unit/utils/test_endpoints.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/utils/test_endpoints.py
- tests/unit/app/endpoints/test_streaming_query.py
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/app/endpoints/query.pysrc/utils/endpoints.pysrc/app/endpoints/streaming_query.py
src/app/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use standard FastAPI imports (from fastapi import APIRouter, HTTPException, Request, status, Depends) in FastAPI app code
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: All modules start with descriptive module-level docstrings explaining purpose
Use logger = logging.getLogger(name) for module logging after import logging
Define type aliases at module level for clarity
All functions require docstrings with brief descriptions
Provide complete type annotations for all function parameters and return types
Use typing_extensions.Self in model validators where appropriate
Use modern union syntax (str | int) and Optional[T] or T | None consistently
Function names use snake_case with descriptive, action-oriented prefixes (get_, validate_, check_)
Avoid in-place parameter modification; return new data structures instead of mutating arguments
Use appropriate logging levels: debug, info, warning, error with clear messages
All classes require descriptive docstrings explaining purpose
Class names use PascalCase with conventional suffixes (Configuration, Error/Exception, Resolver, Interface)
Abstract base classes should use abc.ABC and @AbstractMethod for interfaces
Provide complete type annotations for all class attributes
Follow Google Python docstring style for modules, classes, and functions, including Args, Returns, Raises, Attributes sections as needed
Files:
src/app/endpoints/query.pysrc/utils/endpoints.pysrc/app/endpoints/streaming_query.py
src/{app/**/*.py,client.py}
📄 CodeRabbit inference engine (CLAUDE.md)
Use async def for I/O-bound operations and external API calls
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
src/app/endpoints/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
In API endpoints, raise FastAPI HTTPException with appropriate status codes for error handling
Files:
src/app/endpoints/query.pysrc/app/endpoints/streaming_query.py
🧬 Code graph analysis (2)
src/utils/endpoints.py (2)
src/models/responses.py (1)
ReferencedDocument(55-67)src/utils/types.py (1)
TurnSummary(89-163)
src/app/endpoints/streaming_query.py (2)
src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(529-543)create_rag_chunks_dict(349-362)src/utils/types.py (1)
TurnSummary(89-163)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build-pr
- GitHub Check: e2e_tests
🔇 Additional comments (13)
src/app/endpoints/query.py (1)
362-362: Log message accurately reflects the fixed behavior.The updated log message correctly describes that the endpoint now uses
referenced_documentsfrom the response (viaparse_referenced_documents(response)at line 702) rather than reconstructing them locally. This aligns with the bug fix confirmed in past review comments.src/app/endpoints/streaming_query.py (3)
50-51: LGTM!The new utility imports are correctly placed and properly used throughout the file to create rag_chunks and referenced_documents.
679-680: Good defensive check for None events.The
if chunk.event is None: continueguard preventsAttributeErrorwhen accessingchunk.event.payloaddownstream, making the streaming response more resilient to unexpected chunk formats.
715-715: Correctly passes actual RAG chunks to transcript storage.This change fixes the issue flagged in past reviews where an empty list was passed. Now
create_rag_chunks_dict(summary)properly provides the RAG chunks for transcript storage.src/utils/endpoints.py (9)
4-4: LGTM!The new imports are properly used throughout the file for URL validation, type hints, and the new utility functions.
Also applies to: 8-8, 13-13, 19-19
349-362: LGTM!The
create_rag_chunks_dictfunction cleanly transformsRAGChunkobjects into dictionaries suitable for JSON serialization. The list comprehension is efficient and maintains all necessary fields.
365-379: Robust HTTP source processing with error handling.The function properly:
- Validates URLs using
AnyUrlwith try/except- Logs warnings for invalid URLs without crashing
- Deduplicates using the
doc_urlsset- Extracts a meaningful title from the URL path
382-417: Well-structured document ID processing with metadata enrichment.The function correctly:
- Deduplicates using
doc_idsanddoc_urlssets- Enriches document IDs with metadata when available
- Type-checks
doc_urlandtitleto handle malformed metadata gracefully- Validates URLs with proper error handling
- Constructs meaningful titles with fallback logic
420-445: LGTM!The function properly adds additional documents from
metadata_mapwith:
- Type checking to handle malformed metadata
- URL validation with error handling
- Deduplication to prevent duplicate entries
448-488: Well-designed core processing function.The function properly orchestrates document processing by:
- Filtering out empty sources and
DEFAULT_RAG_TOOL(which is not a document source)- Dispatching to appropriate handlers based on source type (HTTP vs document ID)
- Adding additional documents from metadata
- Maintaining deduplication across all processing paths
491-525: Excellent unified function addressing code duplication.This function effectively addresses the code duplication concern raised by @asamal4 by:
- Centralizing RAG chunk processing logic
- Supporting both dictionary and
ReferencedDocumentoutput formats viareturn_dict_format- Delegating core processing to
_process_rag_chunks_for_documentsThe design allows both streaming and query endpoints to use the same underlying logic while maintaining backward compatibility through the format parameter.
529-543: LGTM!The backward compatibility wrapper correctly:
- Maintains the existing API for the streaming endpoint
- Extracts
rag_chunksfromsummary.rag_chunks- Returns properly typed
ReferencedDocumentobjects
546-559: LGTM!The backward compatibility wrapper correctly:
- Maintains the existing API for chunk-based document creation
- Delegates to the core processing function
- Returns properly typed
ReferencedDocumentobjects
…mon structure for both query and streaming_query
bsatapat-jpg
force-pushed
the
stream_query
branch
from
dfad849 to
54a5858
Compare
|
LGTM |
6 participants
Description
Type of change
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
Bug Fixes
Tests