refactor: Add back verbose parameter to search
This commit is contained in:
parent
bd03a43efa
commit
a27b4b5cd0
3 changed files with 28 additions and 18 deletions
|
|
@ -31,6 +31,7 @@ class SearchPayloadDTO(InDTO):
|
|||
node_name: Optional[list[str]] = Field(default=None, example=[])
|
||||
top_k: Optional[int] = Field(default=10)
|
||||
only_context: bool = Field(default=False)
|
||||
verbose: bool = Field(default=False)
|
||||
|
||||
|
||||
def get_search_router() -> APIRouter:
|
||||
|
|
@ -117,6 +118,7 @@ def get_search_router() -> APIRouter:
|
|||
"node_name": payload.node_name,
|
||||
"top_k": payload.top_k,
|
||||
"only_context": payload.only_context,
|
||||
"verbose": payload.verbose,
|
||||
"cognee_version": cognee_version,
|
||||
},
|
||||
)
|
||||
|
|
@ -133,6 +135,7 @@ def get_search_router() -> APIRouter:
|
|||
system_prompt=payload.system_prompt,
|
||||
node_name=payload.node_name,
|
||||
top_k=payload.top_k,
|
||||
verbose=payload.verbose,
|
||||
only_context=payload.only_context,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ async def search(
|
|||
session_id: Optional[str] = None,
|
||||
wide_search_top_k: Optional[int] = 100,
|
||||
triplet_distance_penalty: Optional[float] = 3.5,
|
||||
verbose: bool = False,
|
||||
) -> List[SearchResult]:
|
||||
"""
|
||||
Search and query the knowledge graph for insights, information, and connections.
|
||||
|
|
@ -218,6 +219,7 @@ async def search(
|
|||
session_id=session_id,
|
||||
wide_search_top_k=wide_search_top_k,
|
||||
triplet_distance_penalty=triplet_distance_penalty,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
return filtered_search_results
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ async def search(
|
|||
session_id: Optional[str] = None,
|
||||
wide_search_top_k: Optional[int] = 100,
|
||||
triplet_distance_penalty: Optional[float] = 3.5,
|
||||
verbose=False,
|
||||
) -> List[SearchResult]:
|
||||
"""
|
||||
|
||||
|
|
@ -141,25 +142,29 @@ async def search(
|
|||
datasets = prepared_search_results["datasets"]
|
||||
|
||||
if only_context:
|
||||
return_value.append(
|
||||
{
|
||||
"search_result": [context] if context else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
"graphs": graphs,
|
||||
}
|
||||
)
|
||||
search_result_dict = {
|
||||
"search_result": [context] if context else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
}
|
||||
if verbose:
|
||||
# Include graphs only in verbose mode
|
||||
search_result_dict["graphs"] = graphs
|
||||
|
||||
return_value.append(search_result_dict)
|
||||
else:
|
||||
return_value.append(
|
||||
{
|
||||
"search_result": [result] if result else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
"graphs": graphs,
|
||||
}
|
||||
)
|
||||
search_result_dict = {
|
||||
"search_result": [result] if result else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
}
|
||||
if verbose:
|
||||
# Include graphs only in verbose mode
|
||||
search_result_dict["graphs"] = graphs
|
||||
return_value.append(search_result_dict)
|
||||
|
||||
return return_value
|
||||
else:
|
||||
return_value = []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue