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

chore: tweak type hints

Browse source
This commit is contained in:
lxobr 2026-01-12 13:44:04 +01:00
parent fce018d43d
commit 5ac288afa3
1 changed files with 7 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
cognee/modules/retrieval/utils/node_edge_vector_search.py
View file

@ -16,8 +16,8 @@ class NodeEdgeVectorSearch:
self.edge_collection = edge_collection self.edge_collection = edge_collection
self.vector_engine = vector_engine or self._init_vector_engine() self.vector_engine = vector_engine or self._init_vector_engine()
self.query_vector: Optional[Any] = None self.query_vector: Optional[Any] = None
self.node_distances: dict[str, list[list[Any]]] = {} self.node_distances: dict[str, list[Any]] = {}
self.edge_distances: list[list[Any]] = [] self.edge_distances: list[Any] = []
self.query_list_length: Optional[int] = None self.query_list_length: Optional[int] = None
def _init_vector_engine(self): def _init_vector_engine(self):
@ -109,7 +109,11 @@ class NodeEdgeVectorSearch:
async def _run_single_search( async def _run_single_search(
self, collections: List[str], query: str, wide_search_limit: Optional[int] self, collections: List[str], query: str, wide_search_limit: Optional[int]
) -> List[List[Any]]: ) -> 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) await self._embed_query(query)
search_tasks = [ search_tasks = [
self._search_single_collection(self.vector_engine, wide_search_limit, collection) self._search_single_collection(self.vector_engine, wide_search_limit, collection)