topoteretes--cognee/cognee/modules/retrieval/utils/node_edge_vector_search.py

import asyncio
import time
from typing import Any, List, Optional

from cognee.shared.logging_utils import get_logger, ERROR
from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
from cognee.infrastructure.databases.vector import get_vector_engine_async
from cognee.modules.observability import new_span, COGNEE_VECTOR_COLLECTION

logger = get_logger(level=ERROR)


class NodeEdgeVectorSearch:
    """Manages vector search and distance retrieval for graph nodes and edges."""

    def __init__(self, edge_collection: str = "EdgeType_relationship_name", vector_engine=None):
        self.edge_collection = edge_collection
        # ``get_vector_engine_async()`` is async, so this sync ``__init__`` can't
        # eagerly resolve it. Keep the (possibly-None) injected engine and
        # resolve lazily in the first async method via ``_get_vector_engine()``.
        self.vector_engine = vector_engine
        self.query_vector: Optional[Any] = None
        self.node_distances: dict[str, list[Any]] = {}
        self.edge_distances: list[Any] = []
        self.query_list_length: Optional[int] = None

    async def _get_vector_engine(self):
        if self.vector_engine is None:
            try:
                self.vector_engine = await get_vector_engine_async()
            except Exception as e:
                logger.error("Failed to initialize vector engine: %s", e)
                raise RuntimeError("Initialization error") from e
        return self.vector_engine

    async def embed_and_retrieve_distances(
        self,
        query: Optional[str] = None,
        query_batch: Optional[List[str]] = None,
        collections: List[str] = None,
        wide_search_limit: Optional[int] = None,
        node_name: Optional[List[str]] = None,
        node_name_filter_operator: str = "OR",
    ):
        """Embeds query/queries and retrieves vector distances from all collections."""
        if query is not None and query_batch is not None:
            raise ValueError("Cannot provide both 'query' and 'query_batch'; use exactly one.")
        if query is None and query_batch is None:
            raise ValueError("Must provide either 'query' or 'query_batch'.")
        if not collections:
            raise ValueError("'collections' must be a non-empty list.")

        with new_span("cognee.retrieval.vector_search") as span:
            span.set_attribute("cognee.vector.collection_count", len(collections))
            span.set_attribute(COGNEE_VECTOR_COLLECTION, ", ".join(collections))
            span.set_attribute(
                "cognee.vector.mode", "batch" if query_batch is not None else "single"
            )
            if wide_search_limit is not None:
                span.set_attribute("cognee.vector.wide_search_limit", wide_search_limit)

            start_time = time.time()

            if query_batch is not None:
                self.query_list_length = len(query_batch)
                span.set_attribute("cognee.vector.batch_size", len(query_batch))
                search_results = await self._run_batch_search(collections, query_batch)
            else:
                self.query_list_length = None
                search_results = await self._run_single_search(
                    collections, query, wide_search_limit, node_name, node_name_filter_operator
                )

            elapsed_time = time.time() - start_time
            collections_with_results = sum(1 for result in search_results if any(result))
            logger.info(
                f"Vector collection retrieval completed: Retrieved distances from "
                f"{collections_with_results} collections in {elapsed_time:.2f}s"
            )

            span.set_attribute("cognee.vector.collections_with_results", collections_with_results)
            span.set_attribute("cognee.vector.duration_ms", round(elapsed_time * 1000, 1))

            self.set_distances_from_results(collections, search_results, self.query_list_length)

    def has_results(self) -> bool:
        """Checks if any collections returned results."""
        if self.query_list_length is None:
            if self.edge_distances and any(self.edge_distances):
                return True
            return any(
                bool(collection_results) for collection_results in self.node_distances.values()
            )

        if self.edge_distances and any(inner_list for inner_list in self.edge_distances):
            return True
        return any(
            any(results_per_query for results_per_query in collection_results)
            for collection_results in self.node_distances.values()
        )

    def extract_relevant_node_ids(self) -> List[str]:
        """Extracts unique node IDs from search results."""
        if self.query_list_length is not None:
            return []
        relevant_node_ids = set()
        for scored_results in self.node_distances.values():
            for scored_node in scored_results:
                node_id = getattr(scored_node, "id", None)
                if node_id:
                    relevant_node_ids.add(str(node_id))
        return list(relevant_node_ids)

    def set_distances_from_results(
        self,
        collections: List[str],
        search_results: List[List[Any]],
        query_list_length: Optional[int] = None,
    ):
        """Separates search results into node and edge distances with stable shapes.

        Ensures all collections are present in the output, even if empty:
        - Batch mode: missing/empty collections become [[]] * query_list_length
        - Single mode: missing/empty collections become []
        """
        self.node_distances = {}
        self.edge_distances = (
            [] if query_list_length is None else [[] for _ in range(query_list_length)]
        )
        for collection, result in zip(collections, search_results):
            if not result:
                empty_result = (
                    [] if query_list_length is None else [[] for _ in range(query_list_length)]
                )
                if collection == self.edge_collection:
                    self.edge_distances = empty_result
                else:
                    self.node_distances[collection] = empty_result
            else:
                if collection == self.edge_collection:
                    self.edge_distances = result
                else:
                    self.node_distances[collection] = result

    async def _run_batch_search(
        self, collections: List[str], query_batch: List[str]
    ) -> List[List[Any]]:
        """Runs batch search across all collections and returns list-of-lists per collection."""
        search_tasks = [
            self._search_batch_collection(collection, query_batch) for collection in collections
        ]
        return await asyncio.gather(*search_tasks)

    async def _search_batch_collection(
        self, collection_name: str, query_batch: List[str]
    ) -> List[List[Any]]:
        """Searches one collection with batch queries and returns list-of-lists."""
        try:
            vector_engine = await self._get_vector_engine()
            return await vector_engine.batch_search(
                collection_name=collection_name, query_texts=query_batch, limit=None
            )
        except CollectionNotFoundError:
            return [[]] * len(query_batch)

    async def _run_single_search(
        self,
        collections: List[str],
        query: str,
        wide_search_limit: Optional[int],
        node_name: Optional[List[str]],
        node_name_filter_operator: str,
    ) -> List[List[Any]]:
        """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)
        vector_engine = await self._get_vector_engine()
        search_tasks = [
            self._search_single_collection(
                vector_engine,
                wide_search_limit,
                collection,
                node_name,
                node_name_filter_operator,
            )
            for collection in collections
        ]
        search_results = await asyncio.gather(*search_tasks)
        return search_results

    async def _embed_query(self, query: str):
        """Embeds the query and stores the resulting vector."""
        with new_span("cognee.retrieval.embed_query") as span:
            span.set_attribute("cognee.vector.query_length", len(query))
            vector_engine = await self._get_vector_engine()
            query_embeddings = await vector_engine.embedding_engine.embed_text([query])
            self.query_vector = query_embeddings[0]
            span.set_attribute("cognee.vector.embedding_dimensions", len(self.query_vector))

    async def _search_single_collection(
        self,
        vector_engine: Any,
        wide_search_limit: Optional[int],
        collection_name: str,
        node_name: Optional[List[str]],
        node_name_filter_operator: str,
    ):
        """Searches one collection and returns results or empty list if not found."""
        try:
            return await vector_engine.search(
                collection_name=collection_name,
                query_vector=self.query_vector,
                limit=wide_search_limit,
                node_name=node_name,
                node_name_filter_operator=node_name_filter_operator,
            )
        except CollectionNotFoundError:
            return []