Transition to new retrievers, update searches #585
Conversation
WalkthroughThis pull request transitions the project’s retrieval logic from a procedural to an object‐oriented architecture. The changes replace direct function calls with the instantiation and use of various Retriever classes (e.g., CodeRetriever, ChunksRetriever, CompletionRetriever, and others) across API routers, modules, and tasks. Legacy files and functions related to code graph retrieval, search comparisons, and query implementations have been removed. New methods, such as Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant API_Router
participant CodeRetriever
participant Backend_Service
User->>API_Router: Send retrieval query
API_Router->>CodeRetriever: Instantiate & call get_context(query)
CodeRetriever->>Backend_Service: Process query & retrieve context
Backend_Service-->>CodeRetriever: Return context data
CodeRetriever-->>API_Router: Return context result
API_Router-->>User: Deliver JSON response
Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (9)
💤 Files with no reviewable changes (3)
🚧 Files skipped from review as they are similar to previous changes (6)
⏰ Context from checks skipped due to timeout of 90000ms (29)
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
cognee/tasks/chunks/__init__.py (1)
1-14: Implementation looks good, consider adding error handlingThe transition to an object-oriented retriever approach looks clean. The function correctly awaits the retriever's get_completion method.
For improved robustness, consider adding error handling:
async def query_chunks(query: str): + try: return await chunks_retriever.get_completion(query) + except Exception as e: + # Log the error + logger.error(f"Error querying chunks: {e}") + # Re-raise or handle appropriately + raisecognee/tasks/completion/__init__.py (1)
6-6: Unused exception import.
NoRelevantDataFoundis imported but never used. Consider removing it if not needed, or handle the exception if it is expected to be raised.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (19)
cognee/api/v1/cognify/routers/get_code_pipeline_router.py(2 hunks)cognee/modules/retrieval/__init__.py(1 hunks)cognee/modules/retrieval/base_retriever.py(2 hunks)cognee/modules/retrieval/utils/code_graph_retrieval.py(0 hunks)cognee/modules/retrieval/utils/run_search_comparisons.py(0 hunks)cognee/modules/search/methods/search.py(2 hunks)cognee/tasks/chunks/__init__.py(1 hunks)cognee/tasks/chunks/query_chunks.py(0 hunks)cognee/tasks/completion/__init__.py(1 hunks)cognee/tasks/completion/graph_query_completion.py(0 hunks)cognee/tasks/completion/graph_query_summary_completion.py(0 hunks)cognee/tasks/completion/query_completion.py(0 hunks)cognee/tasks/graph/__init__.py(1 hunks)cognee/tasks/graph/query_graph_connections.py(0 hunks)cognee/tasks/summarization/__init__.py(1 hunks)cognee/tasks/summarization/query_summaries.py(0 hunks)evals/qa_context_provider_utils.py(2 hunks)examples/python/graphiti_example.py(2 hunks)notebooks/cognee_graphiti_demo.ipynb(2 hunks)
💤 Files with no reviewable changes (8)
- cognee/tasks/summarization/query_summaries.py
- cognee/tasks/chunks/query_chunks.py
- cognee/tasks/graph/query_graph_connections.py
- cognee/tasks/completion/query_completion.py
- cognee/tasks/completion/graph_query_completion.py
- cognee/modules/retrieval/utils/run_search_comparisons.py
- cognee/tasks/completion/graph_query_summary_completion.py
- cognee/modules/retrieval/utils/code_graph_retrieval.py
✅ Files skipped from review due to trivial changes (1)
- cognee/modules/retrieval/init.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: lint (ubuntu-latest, 3.10.x)
- GitHub Check: Build Cognee Backend Docker App Image
- GitHub Check: docker-compose-test
🔇 Additional comments (18)
cognee/tasks/graph/__init__.py (1)
3-12:Details
❓ Verification inconclusive
Implementation looks good, import path consistency should be reviewed
The transition to using a retriever instance is well-implemented with proper async/await pattern.
I notice that
InsightsRetrieveris imported fromcognee.modules.retrieval.retrieverswhile other retrievers in this PR are imported fromcognee.modules.retrieval.base_retriever. This inconsistency might be intentional due to different implementations, but it's worth verifying that this is the correct import path to maintain consistency in code organization.
🏁 Script executed:
#!/bin/bash # Verify import paths for retriever classes echo "Checking import paths for various retriever classes:" rg -A 1 "from cognee\.modules\.retrieval" --glob "*.py"Length of output: 7042
Ensure Consistent Import Paths for Retrievers
The asynchronous implementation using
InsightsRetrieveris correctly set up with proper async/await usage. However, there's an inconsistency in the import paths. Incognee/tasks/graph/__init__.py,InsightsRetrieveris imported fromcognee.modules.retrieval.retrievers, whereas in other parts of the code (e.g., incognee/modules/search/methods/search.py), it's imported directly fromcognee.modules.retrieval.insights_retriever. Please verify whether this discrepancy is intentional. If not, consider unifying the import paths to maintain consistent code organization.cognee/tasks/summarization/__init__.py (1)
1-11: Consistent retriever implementation, nice job!This implementation follows the same pattern used in the other task modules, creating a single reusable retriever instance and wrapping it with an async function that maintains the same interface as before.
The code is clean, concise, and properly implements the async/await pattern. This approach provides better modularity and makes the codebase more maintainable.
examples/python/graphiti_example.py (2)
15-15: Import updated to use the new retriever classThis change aligns with the PR objective to transition to the new retriever classes, replacing the legacy search implementations.
52-54: Refactored to use object-oriented GraphCompletionRetrieverGood replacement of the legacy function call with the instantiation of GraphCompletionRetriever and its resolve_edges_to_text method. This makes the code more modular and follows the new retriever architecture.
cognee/api/v1/cognify/routers/get_code_pipeline_router.py (2)
7-7: Import updated to use the new retriever classThis change aligns with the PR objective to transition to the new retriever classes, replacing the legacy search implementations.
46-47: Refactored to use object-oriented CodeRetrieverGood replacement of the legacy code_graph_retrieval function with the instantiation of CodeRetriever and its get_context method. This follows the new retriever architecture pattern being established across the codebase.
notebooks/cognee_graphiti_demo.ipynb (2)
40-40: Import updated to use the new retriever classThis change aligns with the PR objective to transition to the new retriever classes, replacing the legacy search implementations.
189-190: Refactored to use object-oriented GraphCompletionRetrieverGood replacement of the legacy function call with the instantiation of GraphCompletionRetriever and its resolve_edges_to_text method. This change is consistent with similar refactoring in other files, showing a systematic approach to the transition.
evals/qa_context_provider_utils.py (2)
5-5: Import updated to use the new retriever classThis change aligns with the PR objective to transition to the new retriever classes, replacing the legacy search implementations.
125-126: Refactored to use object-oriented GraphCompletionRetrieverGood replacement of the legacy function call with the instantiation of GraphCompletionRetriever and its resolve_edges_to_text method. This change maintains consistent behavior while adopting the new retriever architecture.
cognee/tasks/completion/__init__.py (6)
1-5: All good on the base retriever imports.These imports correctly reference the new retriever classes. No issues spotted with the import statements.
8-11: Singleton usage of retrievers.Instantiating retrievers at the module level might be safe if they have no risky shared state. Otherwise, consider scoping them per-request or function for improved concurrency and debugging.
14-16: Async entry point is straightforward.
query_completiondefers everything tocompletion_retrieverwith no intermediate logic. This is fine for clarity, though ensure error handling is done upstream ifNoRelevantDataFoundis raised.
19-21: Graph completion function is consistent.Similar to
query_completion. Confirm you have tests that cover scenarios with no or partial graph results.
23-25: Graph summary completion function.Implementation is consistent with the overall pattern, returning
graph_summary_completion_retriever.get_completion(query).
27-28: Edge resolution logic.
resolve_graph_contextdefers tograph_completion_retriever.resolve_edges_to_text. Looks good. Verify the call meets expected input/output format.cognee/modules/search/methods/search.py (2)
6-14: Retriever-based imports successfully replace legacy implementations.The new imports for
ChunksRetriever,InsightsRetriever,SummariesRetriever, and others align with the transition to an object-oriented retriever architecture. These changes appear consistent with the project’s broader refactor.
50-56: search_tasks dictionary refactor.Switching to
.as_search()references centralizes retrieval logic. This approach is clean and maintainable, enabling further extension as new retrievers are added.
| @classmethod | ||
| def as_search(cls) -> Callable: | ||
| """Creates a search function from the retriever class.""" | ||
| return lambda query: cls().get_completion(query) |
There was a problem hiding this comment.
Add await to async method call in lambda
The as_search method returns a lambda that calls get_completion without awaiting it, but get_completion is an async method as defined in lines 13-16. This will return a coroutine object instead of the actual completion result.
@classmethod
def as_search(cls) -> Callable:
"""Creates a search function from the retriever class."""
- return lambda query: cls().get_completion(query)
+ return lambda query: await cls().get_completion(query)Additionally, since this returns an awaitable function, you should update the return type annotation to reflect this:
@classmethod
-def as_search(cls) -> Callable:
+def as_search(cls) -> Callable[[str], Awaitable[Any]]:This will require adding Awaitable to your imports.
Committable suggestion skipped: line range outside the PR's diff.
There was a problem hiding this comment.
I think I agree with the rabbit here, but I tested and both seem to work. Conceptually, it does seem like we should await for the .get_completion somewhere.
I would probably go with the implementation that defines an inner function, for neatness, but up to you:
async def search(query):
return await cls().get_completion(query)
return search
@hajdul88 , please chime in as you are the one with the most experience with this.
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
I am not sure we have to await it here as it is a lambda definition no?
lxobr
left a comment
lxobr
left a comment
There was a problem hiding this comment.
Nice work, looks good to me.
soobrosa
deleted the
feature/cog-1403-transition-to-new-retrievers-update-searches
branch
Description
Delete legacy search implementations after migrating to new retriever classes
DCO Affirmation
I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin
Summary by CodeRabbit
New Features
Refactor
Bug Fixes