refactor: Add back verbose parameter to search

This commit is contained in:
Igor Ilic 2026-01-13 21:11:57 +01:00
parent bd03a43efa
commit a27b4b5cd0
3 changed files with 28 additions and 18 deletions

View file

@ -31,6 +31,7 @@ class SearchPayloadDTO(InDTO):
node_name: Optional[list[str]] = Field(default=None, example=[]) node_name: Optional[list[str]] = Field(default=None, example=[])
top_k: Optional[int] = Field(default=10) top_k: Optional[int] = Field(default=10)
only_context: bool = Field(default=False) only_context: bool = Field(default=False)
verbose: bool = Field(default=False)
def get_search_router() -> APIRouter: def get_search_router() -> APIRouter:
@ -117,6 +118,7 @@ def get_search_router() -> APIRouter:
"node_name": payload.node_name, "node_name": payload.node_name,
"top_k": payload.top_k, "top_k": payload.top_k,
"only_context": payload.only_context, "only_context": payload.only_context,
"verbose": payload.verbose,
"cognee_version": cognee_version, "cognee_version": cognee_version,
}, },
) )
@ -133,6 +135,7 @@ def get_search_router() -> APIRouter:
system_prompt=payload.system_prompt, system_prompt=payload.system_prompt,
node_name=payload.node_name, node_name=payload.node_name,
top_k=payload.top_k, top_k=payload.top_k,
verbose=payload.verbose,
only_context=payload.only_context, only_context=payload.only_context,
) )

View file

@ -35,6 +35,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,
) -> List[SearchResult]: ) -> List[SearchResult]:
""" """
Search and query the knowledge graph for insights, information, and connections. Search and query the knowledge graph for insights, information, and connections.
@ -218,6 +219,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

View file

@ -46,6 +46,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=False,
) -> List[SearchResult]: ) -> List[SearchResult]:
""" """
@ -141,25 +142,29 @@ 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
) 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
) search_result_dict["graphs"] = graphs
return_value.append(search_result_dict)
return return_value return return_value
else: else:
return_value = [] return_value = []