[RHDHPAI-1143] Implement referenced_documents caching #643
Conversation
maysunfaisal
changed the title
WalkthroughUnifies caching around a CacheEntry that can include referenced_documents, moves ConversationData into models.responses, updates cache backends and endpoint handlers to build/pass CacheEntry objects, adjusts assistant message shaping to optionally include referenced_documents, and makes ReferencedDocument.doc_title nullable in OpenAPI. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant API as API Endpoint
participant LLM as LLM Provider
participant Cache as Cache Backend
rect rgb(245,248,255)
note right of API: Query/Streaming — build unified CacheEntry (may include referenced_documents)
User->>API: Send query
API->>LLM: Request completion / stream
LLM-->>API: Response (+metadata)
API->>API: create ReferencedDocument[] from metadata (if any)
API->>API: construct CacheEntry {query,response,provider,model,started_at,completed_at,referenced_documents}
API->>Cache: insert_or_append(CacheEntry) %% single-object cache API %%
Cache-->>API: Ack
API-->>User: Return response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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 |
This comment was marked as outdated.
This comment was marked as outdated.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 (2)
src/cache/postgres_cache.py (1)
41-81: Handle schema migration foradditional_kwargscolumn.In existing deployments the
cachetable already exists without this column, so the new SELECT/INSERT paths will immediately raisecolumn "additional_kwargs" does not exist. Please extendinitialize_cache()toALTER TABLE(or otherwise migrate) before we read/write the column.logger.info("Initializing table for cache") cursor.execute(PostgresCache.CREATE_CACHE_TABLE) + logger.info("Ensuring additional_kwargs column exists") + cursor.execute( + "ALTER TABLE cache ADD COLUMN IF NOT EXISTS additional_kwargs jsonb" + ) +src/cache/sqlite_cache.py (1)
45-86: Add migration for SQLiteadditional_kwargscolumn.Production environments already have a
cachetable without this column; the new SELECT will crash withOperationalError: no such column: additional_kwargs. Please teachinitialize_cache()to add the column when missing (e.g., checkPRAGMA table_info('cache')and issueALTER TABLE cache ADD COLUMN additional_kwargs TEXT).logger.info("Initializing table for cache") cursor.execute(SQLiteCache.CREATE_CACHE_TABLE) + logger.info("Ensuring additional_kwargs column exists") + existing_cols = { + row[1] for row in cursor.execute("PRAGMA table_info('cache')") + } + if "additional_kwargs" not in existing_cols: + cursor.execute("ALTER TABLE cache ADD COLUMN additional_kwargs TEXT") +
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
docs/openapi.json(6 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(3 hunks)src/app/endpoints/streaming_query.py(3 hunks)src/cache/cache.py(1 hunks)src/cache/in_memory_cache.py(1 hunks)src/cache/noop_cache.py(1 hunks)src/cache/postgres_cache.py(6 hunks)src/cache/sqlite_cache.py(6 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)src/utils/endpoints.py(1 hunks)tests/unit/app/endpoints/test_conversations_v2.py(2 hunks)tests/unit/app/endpoints/test_query.py(4 hunks)tests/unit/app/endpoints/test_streaming_query.py(3 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (9)
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/conversations_v2.pysrc/utils/endpoints.pysrc/cache/noop_cache.pysrc/app/endpoints/query.pysrc/cache/in_memory_cache.pysrc/cache/cache.pysrc/cache/sqlite_cache.pysrc/models/cache_entry.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/models/responses.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/conversations_v2.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/app/endpoints/conversations_v2.pysrc/utils/endpoints.pysrc/cache/noop_cache.pytests/unit/cache/test_postgres_cache.pysrc/app/endpoints/query.pysrc/cache/in_memory_cache.pysrc/cache/cache.pysrc/cache/sqlite_cache.pytests/unit/app/endpoints/test_streaming_query.pysrc/models/cache_entry.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/app/endpoints/test_query.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/models/responses.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/conversations_v2.pysrc/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/conversations_v2.pysrc/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/cache/test_postgres_cache.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/app/endpoints/test_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/cache/test_postgres_cache.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/app/endpoints/test_query.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/cache_entry.pysrc/models/responses.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/cache_entry.pysrc/models/responses.py
🧠 Learnings (1)
📚 Learning: 2025-09-18T16:46:33.353Z
Learnt from: CR
PR: lightspeed-core/lightspeed-stack#0
File: CLAUDE.md:0-0
Timestamp: 2025-09-18T16:46:33.353Z
Learning: Applies to src/client.py : Use Llama Stack client import: from llama_stack_client import AsyncLlamaStackClient
Applied to files:
src/app/endpoints/query.py
🧬 Code graph analysis (14)
src/utils/endpoints.py (1)
src/models/cache_entry.py (1)
CacheEntry(12-29)
src/cache/noop_cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(12-29)src/models/responses.py (1)
ConversationData(101-112)
tests/unit/cache/test_postgres_cache.py (3)
src/cache/postgres_cache.py (3)
PostgresCache(16-395)insert_or_append(235-285)get(191-232)src/models/cache_entry.py (2)
AdditionalKwargs(7-9)CacheEntry(12-29)src/models/responses.py (2)
ReferencedDocument(114-128)ConversationData(101-112)
src/app/endpoints/query.py (2)
src/models/cache_entry.py (2)
CacheEntry(12-29)AdditionalKwargs(7-9)src/utils/endpoints.py (1)
store_conversation_into_cache(188-208)
src/cache/in_memory_cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(12-29)src/models/responses.py (1)
ConversationData(101-112)
src/cache/cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(12-29)src/models/responses.py (1)
ConversationData(101-112)
src/cache/sqlite_cache.py (4)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (2)
CacheEntry(12-29)AdditionalKwargs(7-9)src/models/config.py (2)
config(139-145)SQLiteDatabaseConfiguration(74-77)src/models/responses.py (1)
ConversationData(101-112)
tests/unit/app/endpoints/test_streaming_query.py (1)
src/models/cache_entry.py (1)
CacheEntry(12-29)
src/models/cache_entry.py (1)
src/models/responses.py (1)
ReferencedDocument(114-128)
tests/unit/cache/test_sqlite_cache.py (4)
src/models/cache_entry.py (2)
AdditionalKwargs(7-9)CacheEntry(12-29)src/models/responses.py (2)
ReferencedDocument(114-128)ConversationData(101-112)tests/unit/cache/test_postgres_cache.py (2)
test_insert_and_get_with_additional_kwargs(385-430)test_insert_and_get_without_additional_kwargs(433-470)src/cache/sqlite_cache.py (2)
insert_or_append(233-283)get(189-230)
tests/unit/app/endpoints/test_conversations_v2.py (3)
src/app/endpoints/conversations_v2.py (1)
transform_chat_message(241-262)src/models/cache_entry.py (2)
AdditionalKwargs(7-9)CacheEntry(12-29)src/models/responses.py (1)
ReferencedDocument(114-128)
tests/unit/app/endpoints/test_query.py (2)
src/models/cache_entry.py (1)
CacheEntry(12-29)src/models/responses.py (1)
ReferencedDocument(114-128)
src/app/endpoints/streaming_query.py (3)
src/models/cache_entry.py (2)
CacheEntry(12-29)AdditionalKwargs(7-9)src/models/responses.py (1)
ReferencedDocument(114-128)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(520-534)store_conversation_into_cache(188-208)
src/cache/postgres_cache.py (2)
src/models/cache_entry.py (2)
CacheEntry(12-29)AdditionalKwargs(7-9)src/models/responses.py (1)
ConversationData(101-112)
🔇 Additional comments (11)
src/cache/cache.py (1)
5-6: LGTM! Import refactoring aligns with architectural improvements.The import path change correctly reflects that
ConversationDatais now defined inmodels.responsesas a public response model, improving code organization.src/cache/in_memory_cache.py (1)
4-6: LGTM! Consistent import refactoring.Import changes align with the architectural refactoring to move
ConversationDatato the responses module.src/models/responses.py (2)
101-112: LGTM! Well-structured response model.The
ConversationDatamodel is properly documented with clear field descriptions and appropriate type annotations.
126-128: No unsafe.doc_titleusages detected; making it optional is safe.src/app/endpoints/conversations_v2.py (1)
243-254: LGTM! Improved readability and proper additional_kwargs handling.The refactoring enhances code clarity by:
- Using explicit
user_messageandassistant_messagevariables instead of inline dictionaries- Conditionally including
additional_kwargsviamodel_dump(), which properly serializes the Pydantic model- Maintaining type safety with the
dict[str, Any]annotation onassistant_messagesrc/cache/noop_cache.py (1)
4-5: LGTM! Consistent with cache module refactoring.Import changes align with the architectural refactoring across all cache implementations.
tests/unit/app/endpoints/test_query.py (2)
181-206: LGTM! Test setup properly reflects new caching architecture.The test updates correctly:
- Mock
store_conversation_into_cacheto expect a singleCacheEntryobject- Create mock
ReferencedDocumentobjects with proper structure- Update
retrieve_responsemock to return referenced_documents as the third element
236-245: LGTM! Comprehensive validation of cached entry structure.The assertions thoroughly validate:
- The cached entry is a
CacheEntryinstance- The
additional_kwargsfield is populated correctly- The nested
referenced_documentsstructure contains expected datatests/unit/app/endpoints/test_conversations_v2.py (1)
50-79: LGTM! Thorough test coverage for additional_kwargs transformation.The new test effectively validates:
- Construction of
CacheEntrywithAdditionalKwargsandReferencedDocumentobjects- Proper inclusion of
additional_kwargsin the assistant message after transformation- Correct structure and content of nested
referenced_documentstests/unit/cache/test_sqlite_cache.py (2)
365-392: LGTM! Validates round-trip serialization of additional_kwargs.The test properly verifies:
CacheEntrywithadditional_kwargscan be persisted to SQLite- JSON serialization/deserialization preserves the structure
- Retrieved entry matches the original, including nested
referenced_documents
394-411: LGTM! Ensures backward compatibility.The test confirms that
CacheEntryobjects withoutadditional_kwargscontinue to work correctly, with the field properly stored and retrieved asNone.
Signed-off-by: Maysun J Faisal <[email protected]>
Signed-off-by: Maysun J Faisal <[email protected]>
Signed-off-by: Maysun J Faisal <[email protected]>
maysunfaisal
force-pushed
the
rag-caching-2
branch
from
856f886 to
ef0b934
Compare
|
Here is an example of the /v2/conversations output 
|
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (1)
src/cache/sqlite_cache.py (1)
214-230: Consider using named access for database rows.The current implementation uses hardcoded numeric indexes (e.g.,
conversation_entry[6]) to access columns, which is fragile—if the SELECT column order changes, all indexes must be updated manually. Consider usingcursor.row_factory = sqlite3.Rowor constructing named tuples/dicts for safer, self-documenting access.Example using
sqlite3.Row:# In get() method, after creating cursor: cursor = self.connection.cursor() cursor.row_factory = sqlite3.Row # Enable column access by name cursor.execute( self.SELECT_CONVERSATION_HISTORY_STATEMENT, (user_id, conversation_id) ) conversation_entries = cursor.fetchall() cursor.close() result = [] for row in conversation_entries: docs_json_str = row["referenced_documents"] docs_obj = None if docs_json_str: try: docs_data = json.loads(docs_json_str) docs_obj = [ ReferencedDocument.model_validate(doc) for doc in docs_data ] except (json.JSONDecodeError, ValueError) as e: logger.warning( "Failed to deserialize referenced_documents for " "conversation %s: %s", conversation_id, e, ) cache_entry = CacheEntry( query=row["query"], response=row["response"], provider=row["provider"], model=row["model"], started_at=row["started_at"], completed_at=row["completed_at"], referenced_documents=docs_obj, ) result.append(cache_entry)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
docs/openapi.json(2 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(3 hunks)src/app/endpoints/streaming_query.py(3 hunks)src/cache/cache.py(1 hunks)src/cache/in_memory_cache.py(1 hunks)src/cache/noop_cache.py(1 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)src/utils/endpoints.py(1 hunks)tests/unit/app/endpoints/test_conversations_v2.py(3 hunks)tests/unit/app/endpoints/test_query.py(4 hunks)tests/unit/app/endpoints/test_streaming_query.py(4 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (8)
- src/cache/cache.py
- src/cache/noop_cache.py
- src/models/responses.py
- tests/unit/app/endpoints/test_query.py
- src/app/endpoints/query.py
- tests/unit/app/endpoints/test_conversations_v2.py
- tests/unit/cache/test_postgres_cache.py
- tests/unit/app/endpoints/test_streaming_query.py
🧰 Additional context used
📓 Path-based instructions (9)
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/conversations_v2.pysrc/utils/endpoints.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/cache/in_memory_cache.pysrc/cache/sqlite_cache.pysrc/models/cache_entry.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/conversations_v2.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/conversations_v2.pysrc/utils/endpoints.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/cache/in_memory_cache.pysrc/cache/sqlite_cache.pysrc/models/cache_entry.pytests/unit/cache/test_sqlite_cache.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/conversations_v2.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/conversations_v2.pysrc/app/endpoints/streaming_query.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/cache_entry.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/cache_entry.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/cache/test_sqlite_cache.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/cache/test_sqlite_cache.py
🧬 Code graph analysis (7)
src/utils/endpoints.py (1)
src/models/cache_entry.py (1)
CacheEntry(8-25)
src/app/endpoints/streaming_query.py (3)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (1)
ReferencedDocument(177-191)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(563-577)store_conversation_into_cache(231-251)
src/cache/postgres_cache.py (4)
src/cache/cache.py (1)
Cache(10-129)src/cache/cache_error.py (1)
CacheError(4-5)src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ConversationData(164-175)ReferencedDocument(177-191)
src/cache/in_memory_cache.py (3)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/config.py (2)
config(140-146)InMemoryCacheConfig(81-84)src/models/responses.py (1)
ConversationData(164-175)
src/cache/sqlite_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ConversationData(164-175)ReferencedDocument(177-191)
src/models/cache_entry.py (1)
src/models/responses.py (1)
ReferencedDocument(177-191)
tests/unit/cache/test_sqlite_cache.py (5)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ReferencedDocument(177-191)ConversationData(164-175)tests/unit/cache/test_postgres_cache.py (2)
test_insert_and_get_with_referenced_documents(385-429)test_insert_and_get_without_referenced_documents(432-468)src/cache/postgres_cache.py (2)
insert_or_append(237-288)get(193-234)src/cache/sqlite_cache.py (2)
insert_or_append(235-286)get(190-232)
🪛 Biome (2.1.2)
docs/openapi.json
[error] 3976-3976: The key UnauthorizedResponse was already declared.
This where a duplicated key was declared again.
If a key is defined multiple times, only the last definition takes effect. Previous definitions are ignored.
(lint/suspicious/noDuplicateObjectKeys)
🪛 GitHub Actions: Python linter
src/app/endpoints/streaming_query.py
[warning] 881-881: C0303: Trailing whitespace (trailing-whitespace)
[warning] 24-24: W0611: Unused AnyUrl imported from pydantic (unused-import)
[warning] 51-51: W0611: Unused ReferencedDocument imported from models.responses (unused-import)
src/cache/postgres_cache.py
[warning] 260-260: C0301: Line too long (105/100) (line-too-long)
src/cache/sqlite_cache.py
[warning] 260-260: C0301: Line too long (101/100) (line-too-long)
src/models/cache_entry.py
[warning] 4-4: C0411: standard import "typing.List" should be placed before third party import "pydantic.BaseModel" (wrong-import-order)
[warning] 3-3: W0611: Unused Field imported from pydantic (unused-import)
tests/unit/cache/test_sqlite_cache.py
[warning] 399-399: C0303: Trailing whitespace (trailing-whitespace)
[warning] 410-410: C0304: Final newline missing (missing-final-newline)
🪛 GitHub Actions: Ruff
src/app/endpoints/streaming_query.py
[error] 24-24: F401 [*] pydantic.AnyUrl imported but unused
[error] 51-51: F401 [*] models.responses.ReferencedDocument imported but unused
src/models/cache_entry.py
[error] 3-3: F401 [*] pydantic.Field imported but unused
⏰ 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). (3)
- GitHub Check: build-pr
- GitHub Check: e2e_tests (azure)
- GitHub Check: e2e_tests (ci)
🔇 Additional comments (4)
docs/openapi.json (1)
3694-3701: Nullable doc_title supports optional referenced_documents in cached responses.Making
doc_titlenullable (anyOf [string, null]) aligns well with the caching feature, allowing flexibility when document titles may not always be available in cached referenced_documents.src/cache/in_memory_cache.py (1)
6-6: LGTM! Import refactoring aligns with model reorganization.The import of
ConversationDatahas been correctly updated to referencemodels.responses, consistent with the broader refactoring in this PR that relocates model definitions.src/utils/endpoints.py (1)
235-247: Store API now takes CacheEntry — LGTMSignature and usage align with new CacheEntry model. No issues.
src/app/endpoints/conversations_v2.py (1)
317-336: The review comment is incorrect and should be ignored.The ConversationResponse model (src/models/responses.py, lines 613–680) defines only
conversation_idandchat_historyfields with noadditional_kwargsfield. The current implementation intransform_chat_messagecorrectly returns a dictionary matching this schema. The suggested change would add a field that the response model does not define and is not present in any documented example. Additionally, CacheEntry has noadditional_kwargsfield (the docstring mentioning it is misleading).The current code—which embeds
referenced_documentsin the assistant message—is consistent with the actual API contract.Likely an incorrect or invalid review comment.
maysunfaisal
force-pushed
the
rag-caching-2
branch
from
f6a27ca to
50ce0c6
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
tests/unit/app/endpoints/test_conversations_v2.py (1)
6-6: Avoid direct AnyUrl(...) construction in testsConstructing AnyUrl("http://example.com") can raise TypeError across pydantic versions. Pass a plain string and let pydantic coerce the field; also drop the unused AnyUrl import.
-from pydantic import AnyUrl @@ - docs = [ - ReferencedDocument( - doc_title="Test Doc", doc_url=AnyUrl("http://example.com") - ) - ] + docs = [ReferencedDocument(doc_title="Test Doc", doc_url="http://example.com")]Also applies to: 87-90
tests/unit/app/endpoints/test_query.py (1)
193-199: Use string for doc_url; let pydantic coerceAvoid direct AnyUrl(...) in test data. Use a plain string for stability.
- mock_referenced_documents = [ - ReferencedDocument( - doc_title="Test Doc 1", doc_url=AnyUrl("http://example.com/1") - ) - ] + mock_referenced_documents = [ + ReferencedDocument(doc_title="Test Doc 1", doc_url="http://example.com/1") + ]tests/unit/cache/test_postgres_cache.py (1)
396-409: Avoid AnyUrl(...) when constructing ReferencedDocument in testsUse string URLs; pydantic will coerce and normalize.
- docs = [ - ReferencedDocument(doc_title="Test Doc", doc_url=AnyUrl("http://example.com")) - ] + docs = [ReferencedDocument(doc_title="Test Doc", doc_url="http://example.com")]
🧹 Nitpick comments (4)
tests/unit/app/endpoints/test_streaming_query.py (1)
365-381: CacheEntry assertions are solid; consider parsing SSE more robustlyThe CacheEntry checks are great. For the SSE payload extraction earlier, prefer replace/partition over slicing index 5 to strip "data: " to avoid fragility if the prefix changes.
Example:
- d = json.loads(streaming_content[6][5:]) + payload = streaming_content[6] + d = json.loads(payload.replace("data: ", "", 1))tests/unit/cache/test_postgres_cache.py (1)
437-440: Prefer value comparison over object identity for BaseModel equalityDepending on pydantic version, BaseModel equality may not be structural. Compare model_dump() results to avoid flakiness.
- assert retrieved_entries[0] == entry_with_docs + assert retrieved_entries[0].model_dump() == entry_with_docs.model_dump()- assert retrieved_entries[0] == entry_without_docs + assert retrieved_entries[0].model_dump() == entry_without_docs.model_dump()Also applies to: 478-480
src/app/endpoints/streaming_query.py (1)
65-66: Align logger with guidelinesPer coding guidelines, prefer logger = logging.getLogger(name) to inherit module name and standardize logging.
-logger = logging.getLogger("app.endpoints.handlers") +logger = logging.getLogger(__name__)src/cache/postgres_cache.py (1)
218-225: Defensive decode for referenced_documents from DBIf psycopg2 returns jsonb as a string in some environments, iterating docs_data will yield characters and break model validation. Guard for str and decode.
- docs_data = conversation_entry[6] + docs_data = conversation_entry[6] docs_obj = None - if docs_data: - docs_obj = [ - ReferencedDocument.model_validate(doc) for doc in docs_data - ] + if docs_data: + if isinstance(docs_data, str): + try: + docs_data = json.loads(docs_data) + except Exception: + logger.warning("Invalid JSON in referenced_documents") + docs_data = None + if isinstance(docs_data, list): + docs_obj = [ + ReferencedDocument.model_validate(doc) for doc in docs_data + ]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
docs/openapi.json(3 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(2 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)tests/unit/app/endpoints/test_conversations_v2.py(3 hunks)tests/unit/app/endpoints/test_query.py(5 hunks)tests/unit/app/endpoints/test_streaming_query.py(5 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- src/app/endpoints/conversations_v2.py
- src/app/endpoints/query.py
- src/models/cache_entry.py
- src/cache/sqlite_cache.py
🧰 Additional context used
📓 Path-based instructions (9)
**/*.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/app/endpoints/test_query.pysrc/app/endpoints/streaming_query.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pysrc/models/responses.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/cache/test_postgres_cache.pysrc/cache/postgres_cache.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/app/endpoints/test_query.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/cache/test_postgres_cache.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/app/endpoints/test_query.pytests/unit/cache/test_sqlite_cache.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/cache/test_postgres_cache.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/models/responses.pysrc/cache/postgres_cache.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
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
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/responses.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/responses.py
🧬 Code graph analysis (7)
tests/unit/app/endpoints/test_query.py (3)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (1)
ReferencedDocument(179-191)src/utils/token_counter.py (1)
TokenCounter(18-41)
src/app/endpoints/streaming_query.py (2)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(563-577)store_conversation_into_cache(231-251)
tests/unit/cache/test_sqlite_cache.py (5)
src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ReferencedDocument(179-191)ConversationData(165-176)tests/unit/cache/test_postgres_cache.py (2)
test_insert_and_get_with_referenced_documents(385-439)test_insert_and_get_without_referenced_documents(442-480)src/cache/postgres_cache.py (2)
insert_or_append(239-293)get(193-236)src/cache/sqlite_cache.py (2)
insert_or_append(235-288)get(190-232)
tests/unit/app/endpoints/test_conversations_v2.py (3)
src/models/responses.py (1)
ReferencedDocument(179-191)src/models/cache_entry.py (1)
CacheEntry(8-25)src/app/endpoints/conversations_v2.py (1)
transform_chat_message(315-332)
tests/unit/app/endpoints/test_streaming_query.py (1)
src/models/cache_entry.py (1)
CacheEntry(8-25)
tests/unit/cache/test_postgres_cache.py (5)
src/cache/postgres_cache.py (3)
PostgresCache(17-403)insert_or_append(239-293)get(193-236)src/models/config.py (2)
config(140-146)PostgreSQLDatabaseConfiguration(87-105)src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ReferencedDocument(179-191)ConversationData(165-176)src/cache/sqlite_cache.py (2)
insert_or_append(235-288)get(190-232)
src/cache/postgres_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(8-25)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
⏰ 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). (3)
- GitHub Check: e2e_tests (azure)
- GitHub Check: e2e_tests (ci)
- GitHub Check: build-pr
🔇 Additional comments (5)
tests/unit/app/endpoints/test_streaming_query.py (2)
288-289: LGTM: explicit tool_calls in final messageIncluding tool_calls=[] makes end-of-stream payload deterministic. Good change.
301-303: LGTM: cache write hooked for verificationPatching store_conversation_into_cache is appropriate to assert a single CacheEntry is passed.
tests/unit/app/endpoints/test_query.py (1)
251-263: LGTM: Verifying CacheEntry persisted shapeGood assertions validating response and referenced_documents within the CacheEntry passed to the cache layer.
src/app/endpoints/streaming_query.py (1)
869-883: LGTM: CacheEntry + referenced_documents persisted at end-of-streamConstructing CacheEntry and passing it to store_conversation_into_cache is correct and aligns with the new cache contract. The referenced_documents handling matches the metadata_map extraction.
Also applies to: 885-892
src/cache/postgres_cache.py (1)
259-267: LGTM: Safe serialization for referenced_documentsSerializing referenced_documents via model_dump + json.dumps before INSERT is clear and consistent.
maysunfaisal
force-pushed
the
rag-caching-2
branch
2 times, most recently
from
b033429 to
044f2ff
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
docs/openapi.json (1)
142-154: Fix boolean/null literals in/v1/toolsexample.The example still shows
"required": "True/False"and"default": "null", so schema consumers copying this block get strings instead of proper JSON booleans/null. Please switch these to actual literals (true/false,null).
🧹 Nitpick comments (2)
src/cache/postgres_cache.py (2)
218-234: Consider adding error handling for JSONB deserialization.While PostgreSQL's JSONB type returns Python data structures directly (no
json.loadsneeded), themodel_validate()call at line 223 can still fail if the data doesn't match theReferencedDocumentschema. This could happen due to:
- Schema drift (e.g., missing fields after a model update)
- Manual database modifications
- Data corruption
The SQLite implementation (lines 219-230 in
sqlite_cache.py) wraps deserialization in a try/except block and logs warnings on failure. Consider adding similar error handling here for consistency and robustness:docs_data = conversation_entry[6] docs_obj = None if docs_data: try: docs_obj = [ ReferencedDocument.model_validate(doc) for doc in docs_data ] except (TypeError, ValueError) as e: logger.warning( "Failed to deserialize referenced_documents for " "conversation %s: %s", conversation_id, e, )
260-266: Consider adding error handling for JSON serialization.The serialization at lines 262-266 lacks error handling. If
model_dump()orjson.dumps()fails (e.g., due to a custom field that's not JSON-serializable), the entireinsert_or_append()operation will raise an exception and fail.The SQLite implementation (lines 270-282 in
sqlite_cache.py) wraps serialization in a try/except block. Consider adding similar defensive handling here for consistency:referenced_documents_json = None if cache_entry.referenced_documents: try: docs_as_dicts = [ doc.model_dump(mode="json") for doc in cache_entry.referenced_documents ] referenced_documents_json = json.dumps(docs_as_dicts) except (TypeError, ValueError) as e: logger.warning( "Failed to serialize referenced_documents for " "conversation %s: %s", conversation_id, e, )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
docs/openapi.json(3 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(2 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)tests/unit/app/endpoints/test_conversations_v2.py(2 hunks)tests/unit/app/endpoints/test_query.py(5 hunks)tests/unit/app/endpoints/test_streaming_query.py(5 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/app/endpoints/test_query.py
- src/app/endpoints/streaming_query.py
🧰 Additional context used
📓 Path-based instructions (9)
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/conversations_v2.pysrc/app/endpoints/query.pysrc/models/cache_entry.pysrc/models/responses.pysrc/cache/sqlite_cache.pysrc/cache/postgres_cache.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/conversations_v2.pysrc/app/endpoints/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/conversations_v2.pytests/unit/app/endpoints/test_streaming_query.pysrc/app/endpoints/query.pysrc/models/cache_entry.pytests/unit/app/endpoints/test_conversations_v2.pysrc/models/responses.pytests/unit/cache/test_sqlite_cache.pysrc/cache/sqlite_cache.pytests/unit/cache/test_postgres_cache.pysrc/cache/postgres_cache.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/conversations_v2.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/conversations_v2.pysrc/app/endpoints/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/app/endpoints/test_streaming_query.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/cache/test_sqlite_cache.pytests/unit/cache/test_postgres_cache.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/app/endpoints/test_streaming_query.pytests/unit/app/endpoints/test_conversations_v2.pytests/unit/cache/test_sqlite_cache.pytests/unit/cache/test_postgres_cache.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/cache_entry.pysrc/models/responses.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/cache_entry.pysrc/models/responses.py
🧬 Code graph analysis (8)
tests/unit/app/endpoints/test_streaming_query.py (1)
src/models/cache_entry.py (1)
CacheEntry(7-24)
src/app/endpoints/query.py (3)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/utils/endpoints.py (1)
store_conversation_into_cache(231-251)src/configuration.py (1)
configuration(65-69)
src/models/cache_entry.py (1)
src/models/responses.py (1)
ReferencedDocument(179-191)
tests/unit/app/endpoints/test_conversations_v2.py (3)
src/models/responses.py (1)
ReferencedDocument(179-191)src/models/cache_entry.py (1)
CacheEntry(7-24)src/app/endpoints/conversations_v2.py (1)
transform_chat_message(315-332)
tests/unit/cache/test_sqlite_cache.py (4)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)src/cache/postgres_cache.py (2)
insert_or_append(239-293)get(193-236)src/cache/sqlite_cache.py (2)
insert_or_append(245-307)get(190-242)
src/cache/sqlite_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
tests/unit/cache/test_postgres_cache.py (4)
src/cache/postgres_cache.py (3)
PostgresCache(17-403)insert_or_append(239-293)get(193-236)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)tests/unit/cache/test_sqlite_cache.py (4)
test_insert_and_get_with_referenced_documents(363-392)cursor(58-60)execute(46-48)test_insert_and_get_without_referenced_documents(395-412)
src/cache/postgres_cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
🪛 GitHub Actions: Black
tests/unit/app/endpoints/test_conversations_v2.py
[error] 1-1: Black formatting check would reformat this file.
tests/unit/cache/test_sqlite_cache.py
[error] 1-1: Black formatting check would reformat this file.
tests/unit/cache/test_postgres_cache.py
[error] 1-1: Black formatting check would reformat this file.
⏰ 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). (3)
- GitHub Check: build-pr
- GitHub Check: Pylinter
- GitHub Check: e2e_tests (azure)
🔇 Additional comments (19)
tests/unit/app/endpoints/test_streaming_query.py (2)
54-54: LGTM! Imports and test setup align with the new CacheEntry-based caching flow.The import of
CacheEntryfrommodels.cache_entryis correct, the pylint disable is appropriate for the complex test function, and the explicittool_calls=[]in the mock response aligns with the expected end-of-stream payload structure.Also applies to: 212-213, 288-288
301-381: LGTM! Test correctly validates CacheEntry construction and caching.The test properly mocks
store_conversation_into_cache, extracts theCacheEntryfrom the correct argument position (index 3), and validates that:
- The CacheEntry instance is constructed correctly
- The response field matches "LLM answer"
- The referenced_documents list contains exactly 2 items
- Document titles and URLs are correctly populated
This aligns with the updated caching flow where a single
CacheEntryobject encapsulates all cache metadata.tests/unit/cache/test_sqlite_cache.py (3)
10-11: LGTM! Imports correctly resolved per past review feedback.The imports now correctly pull
CacheEntryfrommodels.cache_entryandReferencedDocumentfrommodels.responses, aligning with the canonical model structure.
363-393: LGTM! Test validates end-to-end referenced_documents round-trip.The test correctly:
- Creates a
ReferencedDocumentwith a plain string URL (Pydantic will coerce it)- Stores a
CacheEntrywithreferenced_documentsviainsert_or_append- Retrieves and validates that the document metadata survives serialization/deserialization
This provides good coverage for the JSON serialization path in the SQLite cache.
Note: The file has a Black formatting failure in the pipeline. Please run Black to reformat before merging.
395-412: LGTM! Test validates handling of absent referenced_documents.The test ensures that when
referenced_documentsisNone(the default), the cache correctly:
- Stores
None(or omits the field)- Retrieves
Noneon readThis provides necessary coverage for the optional field behavior.
tests/unit/app/endpoints/test_conversations_v2.py (2)
17-17: LGTM! Import correctly references ReferencedDocument from models.responses.
62-126: LGTM! Comprehensive test coverage for referenced_documents in message transformation.The test class validates three scenarios:
- Without documents (None): The
referenced_documentskey is omitted from the assistant message.- With documents: The key is present and documents are serialized with
doc_titleanddoc_url.- With empty list: The key is present but the list is empty.
This aligns with the
transform_chat_messagelogic insrc/app/endpoints/conversations_v2.pylines 314-331, wherereferenced_documentsis only added whenentry.referenced_documents is not None.Note: The file has a Black formatting failure in the pipeline. Please run Black to reformat before merging.
src/app/endpoints/query.py (3)
34-34: LGTM! CacheEntry import is correct.
336-344: LGTM! CacheEntry construction includes all required fields.The
CacheEntryobject is properly constructed with all metadata fields, and the conditional on line 343 ensures that emptyreferenced_documentslists are normalized toNone, matching the optional field semantics.
346-353: LGTM! Function call correctly passes CacheEntry as a single parameter.The call to
store_conversation_into_cachenow passes thecache_entryobject as the 4th argument, replacing the previous multi-parameter approach. This aligns with the updated utility function signature insrc/utils/endpoints.pylines 230-250.tests/unit/cache/test_postgres_cache.py (3)
3-3: LGTM! Imports correctly resolved per past review feedback.The imports now correctly pull
ReferencedDocumentfrommodels.responses(notmodels.cache_entry), andjsonis imported to parse serialized data in the tests.Also applies to: 12-13
383-443: LGTM! Test validates PostgreSQL JSON serialization round-trip.The test correctly:
- Finds the INSERT call by SQL pattern (per past review feedback)
- Validates that
referenced_documentsis serialized to JSON with the expected structure- Mocks the DB return to simulate retrieval
- Validates that the deserialized
CacheEntrymatches the originalThis provides good coverage for the PostgreSQL JSONB storage path.
Note: The file has a Black formatting failure in the pipeline. Please run Black to reformat before merging.
445-487: LGTM! Test validates handling of absent referenced_documents in PostgreSQL.The test ensures that when
referenced_documentsisNone:
- The INSERT statement passes
Noneas the last parameter (no JSON serialization)- The DB returns
Nonefor the field- The retrieved
CacheEntryhasreferenced_documents=NoneThis provides necessary coverage for the optional field behavior in PostgreSQL.
src/cache/sqlite_cache.py (4)
6-6: LGTM! Imports correctly updated to support referenced_documents handling.The
jsonmodule is imported for serialization/deserialization, and the model imports are correct withReferencedDocumentcoming frommodels.responses.Also applies to: 10-12
25-36: LGTM! Schema and SQL statements correctly extended for referenced_documents.The SQLite schema adds a
referenced_documentstext column to store JSON-serialized document metadata, and the SELECT/INSERT statements are updated to include the new field. This aligns with the PostgreSQL implementation which uses a JSONB column.Also applies to: 46-59, 76-87
216-240: LGTM! Deserialization includes proper error handling per past review.The
get()method now:
- Extracts the JSON string from the DB row
- Uses a try/except block to handle malformed JSON or validation errors
- Logs a warning if deserialization fails
- Continues processing other entries without crashing
This defensive approach prevents cache corruption from breaking the entire retrieval operation.
268-296: LGTM! Serialization includes proper error handling per past review.The
insert_or_append()method now:
- Uses a try/except block around serialization
- Formats the list comprehension across multiple lines for readability
- Logs a warning if serialization fails
- Safely falls back to
Noneon errorThis prevents serialization failures from crashing the insert operation.
src/cache/postgres_cache.py (2)
3-3: LGTM! Imports correctly updated for PostgreSQL referenced_documents handling.The imports align with the SQLite implementation and the canonical data model structure.
Also applies to: 8-10
23-34: LGTM! PostgreSQL schema correctly uses JSONB for referenced_documents.The schema adds a
referenced_documentsJSONB column, which is the appropriate type for storing JSON data in PostgreSQL. The SELECT/INSERT statements are updated to include the new field, maintaining parity with the SQLite implementation.Also applies to: 42-55, 72-82
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
src/app/endpoints/conversations_v2.py (1)
317-325: Avoid emitting empty referenced_documents arraysUse a truthy check so we don’t include an empty list in responses; keeps payloads cleaner and stable.
- if entry.referenced_documents is not None: + if entry.referenced_documents: assistant_message["referenced_documents"] = [ doc.model_dump(mode="json") for doc in entry.referenced_documents ]Also applies to: 329-329
src/models/responses.py (1)
191-191: Consistency: Optional styledoc_title switched to str | None while this module largely uses Optional[T]. Consider unifying style across this file (prefer T | None or Optional[T] consistently).
src/cache/postgres_cache.py (1)
260-266: Consider simplifying by leveraging psycopg2's native jsonb support.The current code manually serializes to a JSON string via
json.dumps(), but psycopg2 can handle Python lists/dicts directly for jsonb columns, making the explicitjson.dumps()call unnecessary.If you prefer a more idiomatic approach, apply this diff:
referenced_documents_json = None if cache_entry.referenced_documents: docs_as_dicts = [ doc.model_dump(mode="json") for doc in cache_entry.referenced_documents ] - referenced_documents_json = json.dumps(docs_as_dicts) + referenced_documents_json = docs_as_dictsThen, if you remove the manual
json.dumps()call, you can also remove theimport jsonstatement on line 3.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
docs/openapi.json(3 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(2 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)tests/unit/app/endpoints/test_conversations_v2.py(2 hunks)tests/unit/app/endpoints/test_query.py(5 hunks)tests/unit/app/endpoints/test_streaming_query.py(5 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/unit/app/endpoints/test_streaming_query.py
- tests/unit/cache/test_sqlite_cache.py
- docs/openapi.json
- src/app/endpoints/query.py
- src/models/cache_entry.py
- tests/unit/app/endpoints/test_query.py
- tests/unit/app/endpoints/test_conversations_v2.py
🧰 Additional context used
📓 Path-based instructions (9)
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/conversations_v2.pysrc/cache/sqlite_cache.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/models/responses.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/conversations_v2.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/conversations_v2.pysrc/cache/sqlite_cache.pytests/unit/cache/test_postgres_cache.pysrc/app/endpoints/streaming_query.pysrc/cache/postgres_cache.pysrc/models/responses.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/conversations_v2.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/conversations_v2.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/cache/test_postgres_cache.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/cache/test_postgres_cache.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/responses.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/responses.py
🧬 Code graph analysis (4)
src/cache/sqlite_cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
tests/unit/cache/test_postgres_cache.py (4)
src/cache/postgres_cache.py (3)
PostgresCache(17-403)insert_or_append(239-293)get(193-236)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (1)
ReferencedDocument(179-191)tests/unit/cache/test_sqlite_cache.py (3)
test_insert_and_get_with_referenced_documents(363-390)cursor(58-60)execute(46-48)
src/app/endpoints/streaming_query.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(563-577)store_conversation_into_cache(231-251)
src/cache/postgres_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
⏰ 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 (ci)
🔇 Additional comments (8)
src/models/responses.py (1)
165-177: ConversationData looks good; verify example/schema alignmentFields and docs look correct. One ask: the conversations_v2 endpoint examples still show an ISO string for last_message_timestamp, but ConversationData uses float. Please align example payloads or switch the type if the API should return ISO strings.
src/app/endpoints/streaming_query.py (2)
869-883: LGTM: referenced_documents extraction and CacheEntry constructionGood handling of None vs empty lists and timestamping. Matches the new CacheEntry contract.
885-892: LGTM: single call site using CacheEntryPassing CacheEntry to store_conversation_into_cache simplifies the API and reduces arg drift.
src/cache/sqlite_cache.py (2)
216-230: Nice defensive read-path handlingGood JSON decode and model validation guards; avoids hard failures on corrupted rows.
268-283: Good write-path error handlingSerialization wrapped with try/except and logged with context; sensible fallback to None.
src/cache/postgres_cache.py (3)
3-3: LGTM: Imports align with the model reorganization.The imports correctly reflect the migration of
ConversationDataandReferencedDocumenttomodels.responses, and the addition ofjsonsupports the serialization logic ininsert_or_append().Also applies to: 8-8, 10-10
23-38: LGTM: Schema documentation is accurate.The table schema documentation correctly reflects the addition of the
referenced_documentsjsonb column.
41-55: LGTM: SQL statements correctly implement referenced_documents storage.The table definition, SELECT, and INSERT statements properly handle the new
referenced_documentsjsonb column with parameterized queries.Also applies to: 73-73, 79-83
| CREATE_CACHE_TABLE = """ | ||
| CREATE TABLE IF NOT EXISTS cache ( | ||
| user_id text NOT NULL, | ||
| conversation_id text NOT NULL, | ||
| created_at int NOT NULL, | ||
| started_at text, | ||
| completed_at text, | ||
| query text, | ||
| response text, | ||
| provider text, | ||
| model text, | ||
| user_id text NOT NULL, | ||
| conversation_id text NOT NULL, | ||
| created_at int NOT NULL, | ||
| started_at text, | ||
| completed_at text, | ||
| query text, | ||
| response text, | ||
| provider text, | ||
| model text, | ||
| referenced_documents text, | ||
| PRIMARY KEY(user_id, conversation_id, created_at) | ||
| ); | ||
| """ |
There was a problem hiding this comment.
Add SQLite migration for new referenced_documents column
CREATE TABLE IF NOT EXISTS won’t add the new column to existing DBs. On upgraded installs, INSERT/SELECT will fail because the column is missing.
Add a lightweight migration in initialize_cache:
def initialize_cache(self) -> None:
@@
- logger.info("Initializing table for cache")
+ logger.info("Initializing table for cache")
cursor.execute(SQLiteCache.CREATE_CACHE_TABLE)
@@
- logger.info("Initializing table for conversations")
+ # Ensure referenced_documents column exists on upgrades
+ try:
+ cursor.execute("PRAGMA table_info(cache)")
+ cols = {row[1] for row in cursor.fetchall()}
+ if "referenced_documents" not in cols:
+ logger.info("Adding missing 'referenced_documents' column to cache")
+ cursor.execute(
+ "ALTER TABLE cache ADD COLUMN referenced_documents text"
+ )
+ except Exception as e: # pylint: disable=broad-exception-caught
+ logger.error("Failed to ensure referenced_documents column: %s", e)
+ raise
+
+ logger.info("Initializing table for conversations")
cursor.execute(SQLiteCache.CREATE_CONVERSATIONS_TABLE)This preserves backward compatibility without requiring manual DB resets.
Also applies to: 76-87
🤖 Prompt for AI Agents
In src/cache/sqlite_cache.py around lines 45-59 (and similarly for lines 76-87),
the CREATE TABLE change won’t add referenced_documents to existing DBs; add a
lightweight migration in initialize_cache: open a DB connection, run "PRAGMA
table_info('cache')" to check for a column named referenced_documents, and if
missing execute "ALTER TABLE cache ADD COLUMN referenced_documents text" (wrap
in a transaction/try-except and commit). Ensure this runs before any
INSERT/SELECT against the table so upgraded installs get the new column without
manual resets.
There was a problem hiding this comment.
Do we want to do this programatically?
There was a problem hiding this comment.
@maysunfaisal Yes, I'd recommend doing this programmatically for a few reasons:
Benefits of programmatic migration:
- Zero user intervention - Existing deployments upgrade seamlessly without manual DB changes or data loss
- Matches existing pattern - You're already using
CREATE TABLE IF NOT EXISTSininitialize_cache(), so adding column checks is consistent - Low risk - A single nullable column addition is one of the safest schema changes
- Better UX - Users won't hit cryptic "table cache has no column referenced_documents" errors after pulling the update
Alternative approaches:
- Manual migration: Document that users must run
ALTER TABLEthemselves (error-prone, poor UX) - Drop and recreate: Require users to delete their SQLite DB file (data loss)
- Separate script: Provide a one-time migration script (extra step users must remember)
The suggested code is minimal:
- Checks if column exists via
PRAGMA table_info - Adds it only if missing
- Runs once on startup in the existing
initialize_cache()flow - Wrapped in error handling
However, if your team prefers a different migration strategy (e.g., versioned migrations with a framework) or if SQLite DBs are considered ephemeral in your deployment model, then manual reset might be acceptable. What's your deployment/upgrade pattern for SQLite instances?
maysunfaisal
force-pushed
the
rag-caching-2
branch
from
044f2ff to
c871aaa
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 (3)
src/app/endpoints/query.py (2)
240-241: Do not log full Llama Stack configuration (risk of API key leakage)configuration.llama_stack_configuration may include secrets (e.g., api_key). Log only non-sensitive fields or a sanitized view.
- logger.info("Llama stack config: %s", configuration.llama_stack_configuration) + safe_cfg = { + "url": getattr(configuration.llama_stack_configuration, "url", "<redacted>"), + "use_as_library_client": getattr( + configuration.llama_stack_configuration, "use_as_library_client", None + ), + } + logger.info("Llama stack config (sanitized): %s", safe_cfg)
526-557: Return-type bug and fragile parsing in parse_metadata_from_text_item
- On non-TextContentItem, it returns [], but the signature is Optional[ReferencedDocument].
- ast.literal_eval is brittle for JSON (e.g., true/false) and can silently drop valid blocks.
Fix to return None on non-matching inputs and use json.loads.
def parse_metadata_from_text_item( text_item: TextContentItem, ) -> Optional[ReferencedDocument]: @@ - docs: list[ReferencedDocument] = [] - if not isinstance(text_item, TextContentItem): - return docs + if not isinstance(text_item, TextContentItem): + return None @@ - try: - data = ast.literal_eval(block) + try: + data = json.loads(block) url = data.get("docs_url") title = data.get("title") if url and title: return ReferencedDocument(doc_url=url, doc_title=title) logger.debug("Invalid metadata block (missing url or title): %s", block) except (ValueError, SyntaxError) as e: logger.debug("Failed to parse metadata block: %s | Error: %s", block, e) return Nonesrc/cache/sqlite_cache.py (1)
169-187: Add migration for the new referenced_documents column.
CREATE TABLE IF NOT EXISTSwon't add the new column to existing SQLite databases. On upgraded installations, INSERT/SELECT operations will fail with "no such column: referenced_documents" errors.Add migration logic in
initialize_cache():def initialize_cache(self) -> None: """Initialize cache - clean it up etc.""" if self.connection is None: logger.error("Cache is disconnected") raise CacheError("Initialize_cache: cache is disconnected") cursor = self.connection.cursor() logger.info("Initializing table for cache") cursor.execute(SQLiteCache.CREATE_CACHE_TABLE) + + # Ensure referenced_documents column exists on upgrades + try: + cursor.execute("PRAGMA table_info(cache)") + cols = {row[1] for row in cursor.fetchall()} + if "referenced_documents" not in cols: + logger.info("Adding missing 'referenced_documents' column to cache") + cursor.execute( + "ALTER TABLE cache ADD COLUMN referenced_documents text" + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error("Failed to ensure referenced_documents column: %s", e) + raise logger.info("Initializing table for conversations") cursor.execute(SQLiteCache.CREATE_CONVERSATIONS_TABLE)
🧹 Nitpick comments (7)
src/models/cache_entry.py (1)
15-16: Clarify time types or document ISO formatstarted_at/completed_at are str. Consider datetime serialization (or document ISO-8601 explicitly) to avoid downstream parsing bugs. Non-blocking.
Also applies to: 24-24
tests/unit/app/endpoints/test_conversations_v2.py (1)
62-123: Transformation tests cover None vs [] vs populated — goodSolid coverage of presence/absence semantics. Tiny nit: asserting exact trailing slash normalization can be brittle across pydantic versions; consider startswith("http://example.com") to reduce flakiness.
tests/unit/app/endpoints/test_query.py (1)
251-263: Assertion on positional arg index is brittleAccessing the 4th positional arg assumes a stable signature. Prefer asserting by type scanning or kwargs to reduce test churn.
- call_args = mock_store_in_cache.call_args[0] - cached_entry = call_args[3] + args, kwargs = mock_store_in_cache.call_args + cached_entry = next((a for a in args if isinstance(a, CacheEntry)), kwargs.get("cache_entry")) + assert cached_entry is not Nonesrc/models/responses.py (1)
165-176: ConversationData addition looks good; consider Field metadata for OpenAPIAdding Field(..., description=...) improves schema docs consistency with other models. Non-blocking.
src/app/endpoints/query.py (3)
336-345: CacheEntry construction OK; consider preserving empty list semanticsYou coerce empty referenced_documents to None, which drops the distinction between “none” and “empty”. If you want v2/conversations to surface an explicit empty list, pass [] through.
- referenced_documents=referenced_documents if referenced_documents else None, + referenced_documents=referenced_documents,
625-642: Avoid duplicate shields.list() callsYou await client.shields.list() twice back-to-back. Call once and split results.
- available_input_shields = [ - shield.identifier - for shield in filter(is_input_shield, await client.shields.list()) - ] - available_output_shields = [ - shield.identifier - for shield in filter(is_output_shield, await client.shields.list()) - ] + shields = await client.shields.list() + available_input_shields = [s.identifier for s in filter(is_input_shield, shields)] + available_output_shields = [s.identifier for s in filter(is_output_shield, shields)]
573-586: Defensive iteration over response.stepsIf steps can be None, iterating will raise. Guard for None to harden parsing.
- docs = [] - for step in response.steps: + docs: list[ReferencedDocument] = [] + for step in (response.steps or []): if not isinstance(step, ToolExecutionStep): continue for tool_response in step.tool_responses: if tool_response.tool_name != constants.DEFAULT_RAG_TOOL: continue for text_item in tool_response.content: if not isinstance(text_item, TextContentItem): continue doc = parse_metadata_from_text_item(text_item) if doc: docs.append(doc) return docs
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
docs/openapi.json(3 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(2 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)tests/unit/app/endpoints/test_conversations_v2.py(2 hunks)tests/unit/app/endpoints/test_query.py(5 hunks)tests/unit/app/endpoints/test_streaming_query.py(5 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (5)
- tests/unit/cache/test_sqlite_cache.py
- src/app/endpoints/conversations_v2.py
- docs/openapi.json
- tests/unit/app/endpoints/test_streaming_query.py
- tests/unit/cache/test_postgres_cache.py
🧰 Additional context used
📓 Path-based instructions (9)
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/app/endpoints/streaming_query.pysrc/models/responses.pysrc/cache/sqlite_cache.pysrc/models/cache_entry.pysrc/cache/postgres_cache.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.pytests/unit/app/endpoints/test_query.pysrc/app/endpoints/streaming_query.pysrc/models/responses.pytests/unit/app/endpoints/test_conversations_v2.pysrc/cache/sqlite_cache.pysrc/models/cache_entry.pysrc/cache/postgres_cache.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/app/endpoints/test_query.pytests/unit/app/endpoints/test_conversations_v2.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/app/endpoints/test_query.pytests/unit/app/endpoints/test_conversations_v2.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/responses.pysrc/models/cache_entry.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/responses.pysrc/models/cache_entry.py
🧬 Code graph analysis (7)
src/app/endpoints/query.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/utils/endpoints.py (1)
store_conversation_into_cache(231-251)
tests/unit/app/endpoints/test_query.py (3)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (1)
ReferencedDocument(179-191)src/utils/token_counter.py (1)
TokenCounter(18-41)
src/app/endpoints/streaming_query.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(563-577)store_conversation_into_cache(231-251)
tests/unit/app/endpoints/test_conversations_v2.py (3)
src/models/responses.py (1)
ReferencedDocument(179-191)src/models/cache_entry.py (1)
CacheEntry(7-24)src/app/endpoints/conversations_v2.py (1)
transform_chat_message(315-332)
src/cache/sqlite_cache.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
src/models/cache_entry.py (1)
src/models/responses.py (1)
ReferencedDocument(179-191)
src/cache/postgres_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
⏰ 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). (3)
- GitHub Check: build-pr
- GitHub Check: e2e_tests (azure)
- GitHub Check: e2e_tests (ci)
🔇 Additional comments (6)
tests/unit/app/endpoints/test_conversations_v2.py (1)
17-17: Import location is correct and consistentUsing ReferencedDocument from models.responses matches the refactor. LGTM.
tests/unit/app/endpoints/test_query.py (2)
33-33: Importing CacheEntry here is accurate with new caching flowAlignment with endpoint refactor looks good.
189-199: Good: exercises RAG path via mock referenced_documentsConstruction with ReferencedDocument and AnyUrl is correct and mirrors production behavior.
src/models/responses.py (1)
191-191: doc_title made optional — matches endpoint behaviorThis aligns with scenarios where only URL is known. LGTM.
src/app/endpoints/streaming_query.py (1)
869-892: LGTM! CacheEntry construction and caching logic is correct.The implementation properly:
- Derives
referenced_documentsfrom summary and metadata- Constructs
CacheEntrywith all required fields- Handles the empty list case by converting to
None- Passes the unified
cache_entryobject to the cache storage functionsrc/cache/sqlite_cache.py (1)
216-230: LGTM! Error handling for serialization/deserialization is well-implemented.Both the deserialization path in
get()and the serialization path ininsert_or_append()properly:
- Catch appropriate exception types
- Log warnings with useful context (conversation_id)
- Gracefully degrade by setting values to
Noneon failure- Prevent crashes from corrupted data or schema mismatches
Also applies to: 268-282
| CREATE_CACHE_TABLE = """ | ||
| CREATE TABLE IF NOT EXISTS cache ( | ||
| user_id text NOT NULL, | ||
| conversation_id text NOT NULL, | ||
| created_at timestamp NOT NULL, | ||
| started_at text, | ||
| completed_at text, | ||
| query text, | ||
| response text, | ||
| provider text, | ||
| model text, | ||
| user_id text NOT NULL, | ||
| conversation_id text NOT NULL, | ||
| created_at timestamp NOT NULL, | ||
| started_at text, | ||
| completed_at text, | ||
| query text, | ||
| response text, | ||
| provider text, | ||
| model text, | ||
| referenced_documents jsonb, | ||
| PRIMARY KEY(user_id, conversation_id, created_at) | ||
| ); | ||
| """ |
There was a problem hiding this comment.
Add migration for the new referenced_documents column.
CREATE TABLE IF NOT EXISTS won't add the new column to existing PostgreSQL databases. On upgraded installations, INSERT/SELECT operations will fail with "column does not exist" errors.
Add migration logic in initialize_cache():
def initialize_cache(self) -> None:
"""Initialize cache - clean it up etc."""
if self.connection is None:
logger.error("Cache is disconnected")
raise CacheError("Initialize_cache: cache is disconnected")
# cursor as context manager is not used there on purpose
# any CREATE statement can raise it's own exception
# and it should not interfere with other statements
cursor = self.connection.cursor()
logger.info("Initializing table for cache")
cursor.execute(PostgresCache.CREATE_CACHE_TABLE)
+
+ # Ensure referenced_documents column exists on upgrades
+ try:
+ cursor.execute("""
+ SELECT column_name
+ FROM information_schema.columns
+ WHERE table_name='cache' AND column_name='referenced_documents'
+ """)
+ if not cursor.fetchone():
+ logger.info("Adding missing 'referenced_documents' column to cache")
+ cursor.execute(
+ "ALTER TABLE cache ADD COLUMN referenced_documents jsonb"
+ )
+ except Exception as e:
+ logger.error("Failed to ensure referenced_documents column: %s", e)
+ raise
logger.info("Initializing table for conversations")
cursor.execute(PostgresCache.CREATE_CONVERSATIONS_TABLE)Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In src/cache/postgres_cache.py around lines 41 to 55, the CREATE TABLE SQL adds
referenced_documents but existing DBs won't get the new column; update
initialize_cache() to run a migration that alters the existing table to add the
referenced_documents jsonb column if it does not exist before any INSERT/SELECT.
Implement an idempotent ALTER TABLE ... ADD COLUMN IF NOT EXISTS
referenced_documents jsonb (or equivalent conditional check), run it inside the
same connection/transaction used for initialization, handle/log any errors, and
keep the CREATE TABLE statement for fresh installs.
Signed-off-by: Maysun J Faisal <[email protected]>
maysunfaisal
force-pushed
the
rag-caching-2
branch
from
c871aaa to
54c5c5a
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
src/cache/postgres_cache.py (1)
52-83: Add migration forreferenced_documentsExisting deployments already have a
cachetable without this column. As soon as this ships, both the updatedSELECT … referenced_documentsand the INSERT with that column list will fail withpsycopg2.errors.UndefinedColumn: column cache.referenced_documents does not exist. Please extendinitialize_cache()to run an idempotent migration (e.g.,ALTER TABLE cache ADD COLUMN IF NOT EXISTS referenced_documents jsonb) before executing the new queries so upgrades succeed without manual intervention.logger.info("Initializing table for cache") cursor.execute(PostgresCache.CREATE_CACHE_TABLE) + logger.info("Ensuring referenced_documents column exists") + try: + cursor.execute( + "ALTER TABLE cache " + "ADD COLUMN IF NOT EXISTS referenced_documents jsonb" + ) + except Exception as e: + logger.error( + "Failed to ensure referenced_documents column exists: %s", e + ) + raiseAlso applies to: 180-188
src/cache/sqlite_cache.py (1)
45-87: Backfill referenced_documents column during initializationSQLite databases created before this change won’t have the
referenced_documentscolumn. Once this deploys, the updatedSELECTandINSERTstatements will fail withsqlite3.OperationalError: no such column: referenced_documents. Please updateinitialize_cache()to probe the schema (e.g., viaPRAGMA table_info('cache')) and run anALTER TABLE cache ADD COLUMN referenced_documents textwhen the column is missing so upgrades work without manual steps.logger.info("Initializing table for cache") cursor.execute(SQLiteCache.CREATE_CACHE_TABLE) + try: + cursor.execute("PRAGMA table_info(cache)") + existing_columns = {row[1] for row in cursor.fetchall()} + if "referenced_documents" not in existing_columns: + logger.info("Adding missing 'referenced_documents' column to cache") + cursor.execute( + "ALTER TABLE cache ADD COLUMN referenced_documents text" + ) + except Exception as e: # pylint: disable=broad-exception-caught + logger.error( + "Failed to ensure referenced_documents column exists: %s", e + ) + raise + logger.info("Initializing table for conversations") cursor.execute(SQLiteCache.CREATE_CONVERSATIONS_TABLE)Also applies to: 169-188
🧹 Nitpick comments (1)
tests/unit/app/endpoints/test_query.py (1)
189-198: Consider using plain strings fordoc_urlin test fixtures.Based on learnings from other test files in this PR, directly instantiating
AnyUrlcan cause issues. Consider passing a plain string and letting Pydantic coerce it:- mock_referenced_documents = [ - ReferencedDocument( - doc_title="Test Doc 1", doc_url=AnyUrl("http://example.com/1") - ) - ] + mock_referenced_documents = [ + ReferencedDocument( + doc_title="Test Doc 1", doc_url="http://example.com/1" + ) + ]This approach is more resilient across Pydantic versions.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
docs/openapi.json(3 hunks)src/app/endpoints/conversations_v2.py(1 hunks)src/app/endpoints/query.py(2 hunks)src/app/endpoints/streaming_query.py(2 hunks)src/cache/postgres_cache.py(7 hunks)src/cache/sqlite_cache.py(7 hunks)src/models/cache_entry.py(3 hunks)src/models/responses.py(2 hunks)tests/unit/app/endpoints/test_conversations_v2.py(2 hunks)tests/unit/app/endpoints/test_query.py(5 hunks)tests/unit/app/endpoints/test_streaming_query.py(5 hunks)tests/unit/cache/test_postgres_cache.py(2 hunks)tests/unit/cache/test_sqlite_cache.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/unit/cache/test_sqlite_cache.py
- src/app/endpoints/conversations_v2.py
- src/app/endpoints/query.py
- src/models/cache_entry.py
- tests/unit/app/endpoints/test_streaming_query.py
- tests/unit/app/endpoints/test_conversations_v2.py
🧰 Additional context used
📓 Path-based instructions (9)
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use absolute imports for internal modules (e.g., from auth import get_auth_dependency)
Files:
src/models/responses.pysrc/cache/postgres_cache.pysrc/cache/sqlite_cache.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/models/responses.pysrc/cache/postgres_cache.pytests/unit/app/endpoints/test_query.pysrc/cache/sqlite_cache.pysrc/app/endpoints/streaming_query.pytests/unit/cache/test_postgres_cache.py
src/{models/**/*.py,configuration.py}
📄 CodeRabbit inference engine (CLAUDE.md)
src/{models/**/*.py,configuration.py}: Use @field_validator and @model_validator for custom validation in Pydantic models
Use precise type hints in configuration (e.g., Optional[FilePath], PositiveInt, SecretStr)
Files:
src/models/responses.py
src/models/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/models/**/*.py: Pydantic models: use BaseModel for data models and extend ConfigurationBase for configuration
Use @model_validator and @field_validator for Pydantic model validation
Files:
src/models/responses.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/app/endpoints/test_query.pytests/unit/cache/test_postgres_cache.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/app/endpoints/test_query.pytests/unit/cache/test_postgres_cache.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
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 (5)
src/cache/postgres_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
tests/unit/app/endpoints/test_query.py (3)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (1)
ReferencedDocument(179-191)src/utils/token_counter.py (1)
TokenCounter(18-41)
src/cache/sqlite_cache.py (3)
src/cache/cache.py (1)
Cache(10-129)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
src/app/endpoints/streaming_query.py (2)
src/models/cache_entry.py (1)
CacheEntry(7-24)src/utils/endpoints.py (2)
create_referenced_documents_with_metadata(563-577)store_conversation_into_cache(231-251)
tests/unit/cache/test_postgres_cache.py (3)
src/cache/postgres_cache.py (3)
PostgresCache(17-419)insert_or_append(247-309)get(193-244)src/models/cache_entry.py (1)
CacheEntry(7-24)src/models/responses.py (2)
ConversationData(165-176)ReferencedDocument(179-191)
⏰ 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). (3)
- GitHub Check: build-pr
- GitHub Check: e2e_tests (azure)
- GitHub Check: e2e_tests (ci)
🔇 Additional comments (8)
src/models/responses.py (2)
165-177: LGTM! Well-structured model for conversation metadata.The
ConversationDataclass is properly documented and uses appropriate type annotations. The design aligns with the PR's goal of providing conversation data from cache list operations.
191-191: No issues found—doc_title nullability change is safe.Verification confirms that all code paths creating
ReferencedDocumentobjects guaranteedoc_titleis neverNone:
- Both helper functions (
_process_http_sourceand_process_document_id) computedoc_titlewith fallback expressions that always return a string.- The streaming path (line 201) filters metadata before assignment:
if "docs_url" in v and "title" in vensuresv.get("title")never returnsNone.- No code accesses
doc_titlewithout first ensuring it's set by the creation logic.The type annotation change from
strtostr | Nonecorrectly reflects the model's optional initialization contract and does not introduce breaking changes. All existing usages work correctly.tests/unit/app/endpoints/test_query.py (3)
33-33: LGTM! Import follows guidelines.The
CacheEntryimport is necessary for the new test assertions and follows the coding guidelines for absolute imports.
216-222: LGTM! Mock return value correctly structured.The mocked
retrieve_responsereturn value properly includesmock_referenced_documentsin the expected position.
251-262: LGTM! Comprehensive validation of cached entry.The test properly verifies that
store_conversation_into_cachewas called with a correctly populatedCacheEntryobject, including thereferenced_documentsfield. The assertions thoroughly validate the cache entry structure.tests/unit/cache/test_postgres_cache.py (3)
3-3: LGTM! Imports properly structured.The import additions are necessary for the new tests:
jsonfor parsing serialized referenced_documentsConversationDataandReferencedDocumentfrommodels.responses(the canonical source)These follow the coding guidelines.
Also applies to: 12-13
383-442: LGTM! Comprehensive test for referenced_documents serialization.This test properly validates the end-to-end flow:
- Creates a
CacheEntrywithreferenced_documents- Verifies the JSON serialization in the INSERT statement
- Simulates DB retrieval with the serialized data
- Confirms deserialization produces the original structure
The test correctly uses plain strings for URLs and matches INSERT calls by SQL content, avoiding the brittle indexing issues from earlier iterations.
444-487: LGTM! Validates the None case for referenced_documents.This test correctly verifies that
CacheEntryobjects withoutreferenced_documents(or withNone) are handled properly:
- Confirms
Noneis passed to the INSERT statement- Simulates DB retrieval with
None- Validates deserialization preserves
NoneTogether with the previous test, this provides complete coverage of the
referenced_documentsfield behavior.
|
LGTM, ty! |
3 participants
Description
Assisted-by: Gemini 2.5 Pro
This PR implements the
referenced_documentscaching for postgres and sqlite. Similar to what we had in road-core/service.v1/queryandv1/streaming_querysaves thereferenced_documentsfrom the response in the postgres and sqlite db. The cached data can be fetched usingv2/conversations/{conversation id}The API
v2/conversations/{conversation id}returns the cached referenced_documents in format (see screenshot in the comment below):I tried to keep it similar to how conversation API was working in road-core/services
Type of change
Related Tickets & Documents
https://issues.redhat.com/browse/RHDHPAI-1143
Checklist before requesting a review
Testing
v1/queryorv1/streaming_queryreferenced_documentsusingv2/conversations/{conversation id}Summary by CodeRabbit
Improvements
Documentation
Tests