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

fix: Resolve issue with only_context

Browse source
This commit is contained in:
Igor Ilic 2025-09-23 13:41:12 +02:00
parent c329e3a1b4
commit afa47c28b0
2 changed files with 27 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/modules/search/methods/no_access_control_search.py
View file

@ -35,7 +35,7 @@ async def no_access_control_search(
[get_completion, get_context] = search_tools
if only_context:
return await get_context(query_text)
return None, await get_context(query_text), []
context = await get_context(query_text)
result = await get_completion(query_text, context)

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

@ -143,20 +143,35 @@ async def search(
context = prepared_search_results["context"]
datasets = prepared_search_results["datasets"]
return_value.append(
{
"search_result": [result] if result else None,
"dataset_id": datasets[0].id,
"dataset_name": datasets[0].name,
"graphs": graphs,
}
)
if only_context:
return_value.append(
{
"search_result": [context] if context else None,
"dataset_id": datasets[0].id,
"dataset_name": datasets[0].name,
"graphs": graphs,
}
)
else:
return_value.append(
{
"search_result": [result] if result else None,
"dataset_id": datasets[0].id,
"dataset_name": datasets[0].name,
"graphs": graphs,
}
)
return return_value
else:
return_value = []
for search_result in search_results:
result, context, datasets = search_result
return_value.append(result)
if only_context:
for search_result in search_results:
prepared_search_results = await prepare_search_result(search_result)
return_value.append(prepared_search_results["context"])
else:
for search_result in search_results:
result, context, datasets = search_result
return_value.append(result)
# For maintaining backwards compatibility
if len(return_value) == 1 and isinstance(return_value[0], list):
return return_value[0]