From 5ac288afa3e82c40579901a0463ea290e56e8197 Mon Sep 17 00:00:00 2001 From: lxobr <122801072+lxobr@users.noreply.github.com> Date: Mon, 12 Jan 2026 13:44:04 +0100 Subject: [PATCH] chore: tweak type hints --- .../modules/retrieval/utils/node_edge_vector_search.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/cognee/modules/retrieval/utils/node_edge_vector_search.py b/cognee/modules/retrieval/utils/node_edge_vector_search.py index db9acc121..ff2d98eb8 100644 --- a/cognee/modules/retrieval/utils/node_edge_vector_search.py +++ b/cognee/modules/retrieval/utils/node_edge_vector_search.py @@ -16,8 +16,8 @@ class NodeEdgeVectorSearch: self.edge_collection = edge_collection self.vector_engine = vector_engine or self._init_vector_engine() self.query_vector: Optional[Any] = None - self.node_distances: dict[str, list[list[Any]]] = {} - self.edge_distances: list[list[Any]] = [] + self.node_distances: dict[str, list[Any]] = {} + self.edge_distances: list[Any] = [] self.query_list_length: Optional[int] = None def _init_vector_engine(self): @@ -109,7 +109,11 @@ class NodeEdgeVectorSearch: async def _run_single_search( self, collections: List[str], query: str, wide_search_limit: Optional[int] ) -> List[List[Any]]: - """Runs single query search and returns list-of-lists per collection.""" + """Runs single query search and returns flat lists per collection. + + Returns a list where each element is a collection's results (flat list). + These are stored as flat lists in node_distances/edge_distances for single-query mode. + """ await self._embed_query(query) search_tasks = [ self._search_single_collection(self.vector_engine, wide_search_limit, collection)