Skip to content

Transition to new retrievers, update searches #585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
soobrosa merged 7 commits into from
Feb 27, 2025
Merged

Transition to new retrievers, update searches #585

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
30927d7
First render.
soobrosa Feb 27, 2025
b02231d
Small fixes.
soobrosa Feb 27, 2025
2d90221
coderabbit don't be smart
soobrosa Feb 27, 2025
f482593
Fixing BaseRetriever refs.
soobrosa Feb 27, 2025
158a68d
Catch an orphan.
soobrosa Feb 27, 2025
931717e
Lean inits.
soobrosa Feb 27, 2025
9537221
Merge branch 'dev' into feature/cog-1403-transition-to-new-retrievers…
soobrosa Feb 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions cognee/api/v1/cognify/routers/get_code_pipeline_router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from fastapi.responses import JSONResponse
from cognee.api.DTO import InDTO
from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline
from cognee.modules.retrieval import code_graph_retrieval
from cognee.modules.retrieval.code_retriever import CodeRetriever
from cognee.modules.storage.utils import JSONEncoder


Expand Down Expand Up @@ -43,7 +43,8 @@ async def code_pipeline_retrieve(payload: CodePipelineRetrievePayloadDTO):
else payload.full_input
)

retrieved_files = await code_graph_retrieval(query)
retriever = CodeRetriever()
retrieved_files = await retriever.get_context(query)

return json.dumps(retrieved_files, cls=JSONEncoder)
except Exception as error:
Expand Down
2 changes: 1 addition & 1 deletion cognee/modules/retrieval/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
from cognee.modules.retrieval.utils.code_graph_retrieval import code_graph_retrieval
from cognee.modules.retrieval.code_retriever import CodeRetriever
7 changes: 6 additions & 1 deletion cognee/modules/retrieval/base_retriever.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from abc import ABC, abstractmethod
from typing import Any, Optional
from typing import Any, Optional, Callable


class BaseRetriever(ABC):
Expand All @@ -14,3 +14,8 @@ async def get_context(self, query: str) -> Any:
async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
"""Generates a response using the query and optional context."""
pass

@classmethod
def as_search(cls) -> Callable:
"""Creates a search function from the retriever class."""
return lambda query: cls().get_completion(query)
Comment on lines +18 to +21
Copy link
Contributor

@coderabbitai coderabbitai bot Feb 27, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

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.

Copy link
Collaborator

@lxobr lxobr Feb 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

@coderabbitai coderabbitai bot Feb 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

Copy link
Collaborator

@hajdul88 hajdul88 Feb 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure we have to await it here as it is a lambda definition no?

128 changes: 0 additions & 128 deletions cognee/modules/retrieval/utils/code_graph_retrieval.py
View file
Open in desktop

This file was deleted.

221 changes: 0 additions & 221 deletions cognee/modules/retrieval/utils/run_search_comparisons.py
View file
Open in desktop

This file was deleted.

Loading
Loading