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

Fix: fixes unit test for codepart search

Browse source
This commit is contained in:
hajdul88 2025-01-10 13:52:26 +01:00
parent 6177d04b44
commit 06e8d2268b
1 changed files with 8 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
cognee/modules/retrieval/description_to_codepart_search.py
View file

@ -27,7 +27,7 @@ async def code_description_to_code_part_search(
async def code_description_to_code_part(
query: str, user: User, top_k: int, include_docs: bool
query: str, user: User, top_k: int, include_docs: bool = False
) -> List[str]:
"""
Maps a code description query to relevant code parts using a CodeGraph pipeline.
@ -45,7 +45,6 @@ async def code_description_to_code_part(
ValueError: If arguments are invalid.
RuntimeError: If an unexpected error occurs during execution.
"""
print(include_docs)
if not query or not isinstance(query, str):
raise ValueError("The query must be a non-empty string.")
if top_k <= 0 or not isinstance(top_k, int):
@ -94,7 +93,13 @@ async def code_description_to_code_part(
memory_fragment = CogneeGraph()
await memory_fragment.project_graph_from_db(
graph_engine,
node_properties_to_project=["id", "type", "text", "source_code", "pydantic_type"],
node_properties_to_project=[
"id",
"type",
"text",
"source_code",
"pydantic_type",
],
edge_properties_to_project=["relationship_name"],
)