gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: small changes based on PR comments

Browse source
This commit is contained in:
Andrej Milicevic 2025-12-01 18:32:55 +01:00
parent 76d054b6a5
commit d473ef12ae
2 changed files with 7 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
cognee/modules/retrieval/register_retriever.py
View file

@ -1,7 +1,10 @@
from typing import Type
from .base_retriever import BaseRetriever
from .registered_community_retrievers import registered_community_retrievers
from ..search.types import SearchType
def register_retriever(search_type: SearchType, retriever: BaseRetriever):
def use_retriever(search_type: SearchType, retriever: Type[BaseRetriever]):
"""Register a retriever class for a given search type."""
registered_community_retrievers[search_type] = retriever

5
cognee/modules/search/methods/get_search_type_tools.py
View file

@ -196,9 +196,10 @@ async def get_search_type_tools(
if query_type in registered_community_retrievers:
retriever = registered_community_retrievers[query_type]
retriever_instance = retriever(top_k=top_k)
search_type_tools = [
retriever(top_k=top_k).get_completion,
retriever(top_k=top_k).get_context,
retriever_instance.get_completion,
retriever_instance.get_context,
]
else:
search_type_tools = search_tasks.get(query_type)