From 06e8d2268b231abb680a89ff86c2525d16599389 Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Fri, 10 Jan 2025 13:52:26 +0100 Subject: [PATCH] Fix: fixes unit test for codepart search --- .../retrieval/description_to_codepart_search.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/cognee/modules/retrieval/description_to_codepart_search.py b/cognee/modules/retrieval/description_to_codepart_search.py index 538f76a6e..bd506f76a 100644 --- a/cognee/modules/retrieval/description_to_codepart_search.py +++ b/cognee/modules/retrieval/description_to_codepart_search.py @@ -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"], )