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

refactor: Make graphs return optional

Browse source
This commit is contained in:
Igor Ilic 2025-12-18 11:54:05 +01:00
parent c7e94e296e
commit dd9aad90cb
2 changed files with 26 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/api/v1/search/search.py
View file

@ -33,6 +33,7 @@ async def search(
session_id: Optional[str] = None, session_id: Optional[str] = None,
wide_search_top_k: Optional[int] = 100, wide_search_top_k: Optional[int] = 100,
triplet_distance_penalty: Optional[float] = 3.5, triplet_distance_penalty: Optional[float] = 3.5,
verbose: bool = False,
) -> Union[List[SearchResult], CombinedSearchResult]: ) -> Union[List[SearchResult], CombinedSearchResult]:
""" """
Search and query the knowledge graph for insights, information, and connections. Search and query the knowledge graph for insights, information, and connections.
@ -204,6 +205,7 @@ async def search(
session_id=session_id, session_id=session_id,
wide_search_top_k=wide_search_top_k, wide_search_top_k=wide_search_top_k,
triplet_distance_penalty=triplet_distance_penalty, triplet_distance_penalty=triplet_distance_penalty,
verbose=verbose,
) )
return filtered_search_results return filtered_search_results

42
cognee/modules/search/methods/search.py
View file

@ -49,6 +49,7 @@ async def search(
session_id: Optional[str] = None, session_id: Optional[str] = None,
wide_search_top_k: Optional[int] = 100, wide_search_top_k: Optional[int] = 100,
triplet_distance_penalty: Optional[float] = 3.5, triplet_distance_penalty: Optional[float] = 3.5,
verbose: bool = False,
) -> Union[CombinedSearchResult, List[SearchResult]]: ) -> Union[CombinedSearchResult, List[SearchResult]]:
""" """
@ -173,25 +174,30 @@ async def search(
datasets = prepared_search_results["datasets"] datasets = prepared_search_results["datasets"]
if only_context: if only_context:
return_value.append( search_result_dict = {
{ "search_result": [context] if context else None,
"search_result": [context] if context else None, "dataset_id": datasets[0].id,
"dataset_id": datasets[0].id, "dataset_name": datasets[0].name,
"dataset_name": datasets[0].name, "dataset_tenant_id": datasets[0].tenant_id,
"dataset_tenant_id": datasets[0].tenant_id, }
"graphs": graphs, if verbose:
} # Include graphs only in verbose mode to reduce payload size
) search_result_dict["graphs"] = graphs
return_value.append(search_result_dict)
else: else:
return_value.append( search_result_dict = {
{ "search_result": [result] if result else None,
"search_result": [result] if result else None, "dataset_id": datasets[0].id,
"dataset_id": datasets[0].id, "dataset_name": datasets[0].name,
"dataset_name": datasets[0].name, "dataset_tenant_id": datasets[0].tenant_id,
"dataset_tenant_id": datasets[0].tenant_id, }
"graphs": graphs, if verbose:
} # Include graphs only in verbose mode to reduce payload size
) search_result_dict["graphs"] = graphs
return_value.append(search_result_dict)
return return_value return return_value
else: else:
return_value = [] return_value = []