cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2.dev0__py3-none-any.whl

cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -1,21 +1,18 @@
-import asyncio
-import time
-from typing import List, Optional, Type
+from typing import List, Optional, Type, Union
 
 from cognee.shared.logging_utils import get_logger, ERROR
 from cognee.modules.graph.exceptions.exceptions import EntityNotFoundError
-from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
 from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
 from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.users.models import User
-from cognee.shared.utils import send_telemetry
+from cognee.modules.retrieval.utils.node_edge_vector_search import NodeEdgeVectorSearch
 
 logger = get_logger(level=ERROR)
 
 
 def format_triplets(edges):
+    """Formats edges into human-readable triplet strings."""
     triplets = []
     for edge in edges:
         node1 = edge.node1
@@ -24,12 +21,10 @@ def format_triplets(edges):
         node1_attributes = node1.attributes
         node2_attributes = node2.attributes
 
-        # Filter only non-None properties
         node1_info = {key: value for key, value in node1_attributes.items() if value is not None}
         node2_info = {key: value for key, value in node2_attributes.items() if value is not None}
         edge_info = {key: value for key, value in edge_attributes.items() if value is not None}
 
-        # Create the formatted triplet
         triplet = f"Node1: {node1_info}\nEdge: {edge_info}\nNode2: {node2_info}\n\n\n"
         triplets.append(triplet)
 
@@ -51,7 +46,6 @@ async def get_memory_fragment(
 
     try:
         graph_engine = await get_graph_engine()
-
         await memory_fragment.project_graph_from_db(
             graph_engine,
             node_properties_to_project=properties_to_project,
@@ -61,20 +55,64 @@ async def get_memory_fragment(
             relevant_ids_to_filter=relevant_ids_to_filter,
             triplet_distance_penalty=triplet_distance_penalty,
         )
-
     except EntityNotFoundError:
-        # This is expected behavior - continue with empty fragment
         pass
     except Exception as e:
         logger.error(f"Error during memory fragment creation: {str(e)}")
-        # Still return the fragment even if projection failed
-        pass
 
     return memory_fragment
 
 
+async def _get_top_triplet_importances(
+    memory_fragment: Optional[CogneeGraph],
+    vector_search: NodeEdgeVectorSearch,
+    properties_to_project: Optional[List[str]],
+    node_type: Optional[Type],
+    node_name: Optional[List[str]],
+    triplet_distance_penalty: float,
+    wide_search_limit: Optional[int],
+    top_k: int,
+    query_list_length: Optional[int] = None,
+) -> Union[List[Edge], List[List[Edge]]]:
+    """Creates memory fragment (if needed), maps distances, and calculates top triplet importances.
+
+    Args:
+        query_list_length: Number of queries in batch mode (None for single-query mode).
+            When None, node_distances/edge_distances are flat lists; when set, they are list-of-lists.
+
+    Returns:
+        List[Edge]: For single-query mode (query_list_length is None).
+        List[List[Edge]]: For batch mode (query_list_length is set), one list per query.
+    """
+    if memory_fragment is None:
+        if wide_search_limit is None:
+            relevant_node_ids = None
+        else:
+            relevant_node_ids = vector_search.extract_relevant_node_ids()
+
+        memory_fragment = await get_memory_fragment(
+            properties_to_project=properties_to_project,
+            node_type=node_type,
+            node_name=node_name,
+            relevant_ids_to_filter=relevant_node_ids,
+            triplet_distance_penalty=triplet_distance_penalty,
+        )
+
+    await memory_fragment.map_vector_distances_to_graph_nodes(
+        node_distances=vector_search.node_distances, query_list_length=query_list_length
+    )
+    await memory_fragment.map_vector_distances_to_graph_edges(
+        edge_distances=vector_search.edge_distances, query_list_length=query_list_length
+    )
+
+    return await memory_fragment.calculate_top_triplet_importances(
+        k=top_k, query_list_length=query_list_length
+    )
+
+
 async def brute_force_triplet_search(
-    query: str,
+    query: Optional[str] = None,
+    query_batch: Optional[List[str]] = None,
     top_k: int = 5,
     collections: Optional[List[str]] = None,
     properties_to_project: Optional[List[str]] = None,
@@ -83,33 +121,49 @@ async def brute_force_triplet_search(
     node_name: Optional[List[str]] = None,
     wide_search_top_k: Optional[int] = 100,
     triplet_distance_penalty: Optional[float] = 3.5,
-) -> List[Edge]:
+) -> Union[List[Edge], List[List[Edge]]]:
     """
     Performs a brute force search to retrieve the top triplets from the graph.
 
     Args:
-        query (str): The search query.
+        query (Optional[str]): The search query (single query mode). Exactly one of query or query_batch must be provided.
+        query_batch (Optional[List[str]]): List of search queries (batch mode). Exactly one of query or query_batch must be provided.
         top_k (int): The number of top results to retrieve.
         collections (Optional[List[str]]): List of collections to query.
         properties_to_project (Optional[List[str]]): List of properties to project.
         memory_fragment (Optional[CogneeGraph]): Existing memory fragment to reuse.
         node_type: node type to filter
         node_name: node name to filter
-        wide_search_top_k (Optional[int]): Number of initial elements to retrieve from collections
+        wide_search_top_k (Optional[int]): Number of initial elements to retrieve from collections.
+            Ignored in batch mode (always None to project full graph).
         triplet_distance_penalty (Optional[float]): Default distance penalty in graph projection
 
     Returns:
-        list: The top triplet results.
+        List[Edge]: The top triplet results for single query mode (flat list).
+        List[List[Edge]]: List of top triplet results (one per query) for batch mode (list-of-lists).
+
+    Note:
+        In single-query mode, node_distances and edge_distances are stored as flat lists.
+        In batch mode, they are stored as list-of-lists (one list per query).
     """
-    if not query or not isinstance(query, str):
+    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 query is not None and (not query or not isinstance(query, str)):
         raise ValueError("The query must be a non-empty string.")
+    if query_batch is not None:
+        if not isinstance(query_batch, list) or not query_batch:
+            raise ValueError("query_batch must be a non-empty list of strings.")
+        if not all(isinstance(q, str) and q for q in query_batch):
+            raise ValueError("All items in query_batch must be non-empty strings.")
     if top_k <= 0:
         raise ValueError("top_k must be a positive integer.")
 
-    # Setting wide search limit based on the parameters
-    non_global_search = node_name is None
-
-    wide_search_limit = wide_search_top_k if non_global_search else None
+    query_list_length = len(query_batch) if query_batch is not None else None
+    wide_search_limit = (
+        None if query_list_length else (wide_search_top_k if node_name is None else None)
+    )
 
     if collections is None:
         collections = [
@@ -123,77 +177,37 @@ async def brute_force_triplet_search(
         collections.append("EdgeType_relationship_name")
 
     try:
-        vector_engine = get_vector_engine()
-    except Exception as e:
-        logger.error("Failed to initialize vector engine: %s", e)
-        raise RuntimeError("Initialization error") from e
-
-    query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
+        vector_search = NodeEdgeVectorSearch()
 
-    async def search_in_collection(collection_name: str):
-        try:
-            return await vector_engine.search(
-                collection_name=collection_name, query_vector=query_vector, limit=wide_search_limit
-            )
-        except CollectionNotFoundError:
-            return []
-
-    try:
-        start_time = time.time()
-
-        results = await asyncio.gather(
-            *[search_in_collection(collection_name) for collection_name in collections]
+        await vector_search.embed_and_retrieve_distances(
+            query=None if query_list_length else query,
+            query_batch=query_batch if query_list_length else None,
+            collections=collections,
+            wide_search_limit=wide_search_limit,
         )
 
-        if all(not item for item in results):
-            return []
-
-        # Final statistics
-        vector_collection_search_time = time.time() - start_time
-        logger.info(
-            f"Vector collection retrieval completed: Retrieved distances from {sum(1 for res in results if res)} collections in {vector_collection_search_time:.2f}s"
+        if not vector_search.has_results():
+            return [[] for _ in range(query_list_length)] if query_list_length else []
+
+        results = await _get_top_triplet_importances(
+            memory_fragment,
+            vector_search,
+            properties_to_project,
+            node_type,
+            node_name,
+            triplet_distance_penalty,
+            wide_search_limit,
+            top_k,
+            query_list_length=query_list_length,
         )
 
-        node_distances = {collection: result for collection, result in zip(collections, results)}
-
-        edge_distances = node_distances.get("EdgeType_relationship_name", None)
-
-        if wide_search_limit is not None:
-            relevant_ids_to_filter = list(
-                {
-                    str(getattr(scored_node, "id"))
-                    for collection_name, score_collection in node_distances.items()
-                    if collection_name != "EdgeType_relationship_name"
-                    and isinstance(score_collection, (list, tuple))
-                    for scored_node in score_collection
-                    if getattr(scored_node, "id", None)
-                }
-            )
-        else:
-            relevant_ids_to_filter = None
-
-        if memory_fragment is None:
-            memory_fragment = await get_memory_fragment(
-                properties_to_project=properties_to_project,
-                node_type=node_type,
-                node_name=node_name,
-                relevant_ids_to_filter=relevant_ids_to_filter,
-                triplet_distance_penalty=triplet_distance_penalty,
-            )
-
-        await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
-        await memory_fragment.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
-
-        results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
-
         return results
-
     except CollectionNotFoundError:
-        return []
+        return [[] for _ in range(query_list_length)] if query_list_length else []
     except Exception as error:
         logger.error(
             "Error during brute force search for query: %s. Error: %s",
-            query,
+            query_batch if query_list_length else [query],
             error,
         )
         raise error

cognee/modules/retrieval/utils/node_edge_vector_search.py

@@ -0,0 +1,174 @@
+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
+
+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
+        self.vector_engine = vector_engine or self._init_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
+
+    def _init_vector_engine(self):
+        try:
+            return get_vector_engine()
+        except Exception as e:
+            logger.error("Failed to initialize vector engine: %s", e)
+            raise RuntimeError("Initialization error") from e
+
+    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,
+    ):
+        """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.")
+
+        start_time = time.time()
+
+        if query_batch is not None:
+            self.query_list_length = 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)
+
+        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"
+        )
+
+        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:
+            return await self.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]
+    ) -> 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)
+        search_tasks = [
+            self._search_single_collection(self.vector_engine, wide_search_limit, collection)
+            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."""
+        query_embeddings = await self.vector_engine.embedding_engine.embed_text([query])
+        self.query_vector = query_embeddings[0]
+
+    async def _search_single_collection(
+        self, vector_engine: Any, wide_search_limit: Optional[int], collection_name: 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,
+            )
+        except CollectionNotFoundError:
+            return []

cognee/modules/search/methods/__init__.py

@@ -1 +1,2 @@
 from .search import search
+from .get_retriever_output import get_retriever_output

cognee/modules/search/methods/get_retriever_output.py

@@ -0,0 +1,53 @@
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.search.models.SearchResultPayload import SearchResultPayload
+from cognee.modules.search.methods.get_search_type_retriever_instance import (
+    get_search_type_retriever_instance,
+)
+from cognee.modules.search.types import SearchType
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger()
+
+
+async def get_retriever_output(query_type: SearchType, query_text: str, **kwargs):
+    graph_engine = await get_graph_engine()
+    is_empty = await graph_engine.is_empty()
+
+    if is_empty:
+        logger.warning("Search attempt on an empty knowledge graph")
+
+    retriever_instance = await get_search_type_retriever_instance(
+        query_type=query_type, query_text=query_text, **kwargs
+    )
+
+    # Get raw result objects from retriever and forward to context and completion methods to avoid duplicate retrievals.
+    retrieved_objects = await retriever_instance.get_retrieved_objects(query=query_text)
+
+    # Handle raw result object to extract context information
+    context = await retriever_instance.get_context_from_objects(
+        query=query_text, retrieved_objects=retrieved_objects
+    )
+
+    completion = None
+    if not kwargs.get(
+        "only_context", False
+    ):  # If only_context is True, skip getting completion. Performance optimization.
+        # Handle raw result and context object to handle completion operation
+        completion = await retriever_instance.get_completion_from_context(
+            query=query_text,
+            retrieved_objects=retrieved_objects,
+            context=context,
+        )
+
+    search_result = SearchResultPayload(
+        result_object=retrieved_objects,
+        context=context,
+        completion=completion,
+        search_type=query_type,
+        only_context=kwargs.get("only_context", False),
+        dataset_name=kwargs.get("dataset").name if kwargs.get("dataset") else None,
+        dataset_id=kwargs.get("dataset").id if kwargs.get("dataset") else None,
+        dataset_tenant_id=kwargs.get("dataset").tenant_id if kwargs.get("dataset") else None,
+    )
+
+    return search_result