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 .base_retriever import BaseRetriever
from .registered_community_retrievers import registered_community_retrievers from .registered_community_retrievers import registered_community_retrievers
from ..search.types import SearchType 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 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: if query_type in registered_community_retrievers:
retriever = registered_community_retrievers[query_type] retriever = registered_community_retrievers[query_type]
retriever_instance = retriever(top_k=top_k)
search_type_tools = [ search_type_tools = [
retriever(top_k=top_k).get_completion, retriever_instance.get_completion,
retriever(top_k=top_k).get_context, retriever_instance.get_context,
] ]
else: else:
search_type_tools = search_tasks.get(query_type) search_type_tools = search_tasks.get(query_type)