topoteretes--cognee/cognee/modules/retrieval/graph_completion_retriever.py

import asyncio
from typing import Any, Dict, List, Optional, Type, Union

from cognee.base_config import get_base_config
from cognee.infrastructure.engine import DataPoint
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
from cognee.modules.retrieval.utils.validate_queries import validate_retriever_input
from cognee.modules.graph.utils import resolve_edges_to_text
from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
from cognee.modules.retrieval.base_retriever import BaseRetriever
from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
from cognee.modules.retrieval.utils.global_context import (
    format_global_context_prelude,
    load_root_text,
    search_top_global_context_summaries,
)
from cognee.modules.retrieval.utils.used_graph_elements import (
    is_edge_list,
    extract_from_edges,
)
from cognee.modules.retrieval.utils.references import append_answer_grounded_evidence
from cognee.modules.retrieval.utils.completion import (
    generate_completion,
    generate_completion_batch,
)
from cognee.infrastructure.session.get_session_manager import get_session_manager
from cognee.shared.logging_utils import get_logger
from cognee.infrastructure.databases.unified import get_unified_engine
from cognee.context_global_variables import session_user
from cognee.infrastructure.databases.cache.config import CacheConfig

logger = get_logger("GraphCompletionRetriever")


class GraphCompletionRetriever(BaseRetriever):
    """
    Retriever for handling graph-based completion searches.

    This class implements the retrieval pipeline by searching for graph triplets (get_retrieved_objects function),
    resolving those triplets into human-readable text context (get_context_from_objects function), and generating
    LLM completions using the retrieved graph data (get_completion_from_context function).
    """

    def __init__(
        self,
        user_prompt_path: str = "graph_context_for_question.txt",
        system_prompt_path: str = "answer_simple_question.txt",
        system_prompt: Optional[str] = None,
        top_k: Optional[int] = 5,
        node_type: Optional[Type] = None,
        node_name: Optional[List[str]] = None,
        node_name_filter_operator: str = "OR",
        wide_search_top_k: Optional[int] = 100,
        triplet_distance_penalty: Optional[float] = 6.5,
        feedback_influence: float = get_base_config().default_feedback_influence,
        session_id: Optional[str] = None,
        response_model: Type = str,
        neighborhood_depth: Optional[int] = None,
        neighborhood_seed_top_k: Optional[int] = 10,
        include_global_context_index: bool = False,
        global_context_index_top_k: int = 3,
        include_references: bool = False,
    ):
        """Initialize retriever with prompt paths and search parameters."""
        self.user_prompt_path = user_prompt_path
        self.system_prompt_path = system_prompt_path
        self.system_prompt = system_prompt
        self.top_k = top_k if top_k is not None else 5
        self.wide_search_top_k = wide_search_top_k
        self.node_type = node_type
        self.node_name = node_name
        self.node_name_filter_operator = node_name_filter_operator
        self.triplet_distance_penalty = triplet_distance_penalty
        self.feedback_influence = feedback_influence
        # session_id (Optional[str]): Identifier for managing conversation history.
        self.session_id = session_id
        # response_model (Type): The Pydantic model or type for the expected response.
        self.response_model = response_model
        self.neighborhood_depth = neighborhood_depth
        self.neighborhood_seed_top_k = neighborhood_seed_top_k
        self.include_global_context_index = include_global_context_index
        self.global_context_index_top_k = global_context_index_top_k
        self.include_references = include_references

    def _use_session_cache(self) -> bool:
        """Check if session caching is enabled for the current user."""
        user = session_user.get()
        user_id = getattr(user, "id", None)
        return bool(user_id and CacheConfig().caching)

    @staticmethod
    def _get_vector_index_collections() -> List[str]:
        """Collect vector index collection names from all DataPoint subclasses."""
        collections = []
        for subclass in get_all_subclasses(DataPoint):
            metadata = subclass.model_fields.get("metadata")
            if metadata is None:
                continue
            default = getattr(metadata, "default", None)
            if isinstance(default, dict):
                for field_name in default.get("index_fields", []):
                    collections.append(f"{subclass.__name__}_{field_name}")
        return collections

    async def get_retrieved_objects(
        self, query: Optional[str] = None, query_batch: Optional[List[str]] = None
    ) -> Union[List[Edge], List[List[Edge]]]:
        """
        Performs a brute-force triplet search on the graph and updates access timestamps.

        Args:
            query (str): The search query to find relevant graph triplets.
            query_batch (str): The batch of search queries to find relevant graph triplets.

        Returns:
            List[Edge]: A list of retrieved Edge objects (triplets).
                       Returns an empty list if the graph is empty or no results are found.
        """

        validate_retriever_input(query, query_batch, self._use_session_cache())

        self._unified_engine = await get_unified_engine()
        is_empty = await self._unified_engine.graph.is_empty()

        if is_empty:
            logger.warning("Search attempt on an empty knowledge graph")
            return []

        triplets = await self.get_triplets(query, query_batch)

        # Check if all triplets are empty, in case of batch queries
        if query_batch and all(len(batched_triplets) == 0 for batched_triplets in triplets):
            logger.warning("Empty context was provided to the completion")
            return []

        if len(triplets) == 0:
            logger.warning("Empty context was provided to the completion")
            return []

        return triplets

    async def resolve_edges_to_text(self, retrieved_edges: list) -> str:
        """
        Converts retrieved graph edges into a human-readable string format.

        Parameters:
        -----------

            - retrieved_edges (list): A list of edges retrieved from the graph.

        Returns:
        --------

            - str: A formatted string representation of the nodes and their connections.
        """
        return await resolve_edges_to_text(retrieved_edges)

    async def get_triplets(
        self,
        query: Optional[str] = None,
        query_batch: Optional[List[str]] = None,
    ) -> Union[List[Edge], List[List[Edge]]]:
        """
        Retrieves relevant graph triplets based on a query string.

        Parameters:
        -----------

            - query (str): The query string used to search for relevant triplets in the graph.

        Returns:
        --------

            - list: A list of found triplets that match the query.
        """
        collections = self._get_vector_index_collections()
        unified_engine = getattr(self, "_unified_engine", None)
        return await brute_force_triplet_search(
            query,
            query_batch,
            top_k=self.top_k,
            collections=collections or None,
            node_type=self.node_type,
            node_name=self.node_name,
            node_name_filter_operator=self.node_name_filter_operator,
            wide_search_top_k=self.wide_search_top_k,
            triplet_distance_penalty=self.triplet_distance_penalty,
            feedback_influence=self.feedback_influence,
            unified_engine=unified_engine,
            neighborhood_depth=self.neighborhood_depth,
            neighborhood_seed_top_k=self.neighborhood_seed_top_k,
        )

    async def get_triplets_batch(
        self,
        queries: List[str],
    ) -> List[List[Edge]]:
        """
        Retrieves triplets for a list of queries, using single-query mode when
        possible to enable ID-filtered graph projection.

        When there is only one query, delegates to single-query mode (query=)
        which computes relevant node IDs and filters the graph projection.
        For multiple queries, uses batch mode (query_batch=).

        Returns:
            List[List[Edge]]: One list of edges per query.
        """
        if len(queries) == 1:
            triplets = await self.get_triplets(query=queries[0])
            return [triplets]
        return await self.get_triplets(query_batch=queries)

    async def get_context_from_objects(
        self,
        query: Optional[str] = None,
        query_batch: Optional[List[str]] = None,
        retrieved_objects=None,
    ) -> Union[str, List[str]]:
        """
        Transforms raw retrieved graph triplets into a textual context string.

        Args:
            query (str): The original search query.
            query_batch (List[str]): The batch of original search queries.
            retrieved_objects (List[Edge]): The raw triplets returned from the search.
                                            Output of the get_retrieved_objects method.

        Returns:
            str: A string representing the resolved graph context.
                 Returns an empty list (as string) if no triplets are provided.

        Note: To avoid duplicate retrievals, ensure that retrieved_objects
              are provided from get_retrieved_objects method call.
        """

        triplets = retrieved_objects

        if query_batch:
            # Check if all triplets are empty, in case of batch queries
            if not triplets or all(len(batched_triplets) == 0 for batched_triplets in triplets):
                logger.warning("Empty context was provided to the completion")
                return ["" for _ in query_batch]

            return await asyncio.gather(
                *[self.resolve_edges_to_text(batched_triplets) for batched_triplets in triplets]
            )

        graph_context = await self.resolve_edges_to_text(triplets) if triplets else ""

        if not self.include_global_context_index:
            if not triplets:
                logger.warning("Empty context was provided to the completion")
                return ""
            return graph_context

        prelude = await self._build_global_context_prelude(query)
        if not prelude and not graph_context:
            logger.warning("Empty context was provided to the completion")
            return ""
        if not prelude:
            return graph_context
        if not graph_context:
            return prelude
        return f"{prelude}\n\n{graph_context}"

    async def _build_global_context_prelude(self, query: Optional[str]) -> str:
        if not query:
            return ""
        if getattr(self, "_unified_engine", None) is None:
            self._unified_engine = await get_unified_engine()
        root_text = await load_root_text()
        top_summaries = await search_top_global_context_summaries(
            query, self.global_context_index_top_k, self._unified_engine.vector
        )
        return format_global_context_prelude(root_text, top_summaries)

    def _extract_context_object_ids(self, retrieved_objects: Any) -> Optional[Dict[str, List[str]]]:
        """Extract node_ids and edge_ids from list of Edge. Only used for single-query session path."""
        if not isinstance(retrieved_objects, list) or not retrieved_objects:
            return None
        if not is_edge_list(retrieved_objects):
            return None
        return extract_from_edges(retrieved_objects)

    def _completion_kwargs(self, context: str) -> dict:
        """Common kwargs for completion calls (no session)."""
        return {
            "context": context,
            "user_prompt_path": self.user_prompt_path,
            "system_prompt_path": self.system_prompt_path,
            "system_prompt": self.system_prompt,
            "response_model": self.response_model,
        }

    async def _generate_completion_without_session(
        self,
        query: Optional[str],
        query_batch: Optional[List[str]],
        context: str,
    ) -> List[Any]:
        """Generate completion(s) without session; returns list of completions."""
        kwargs = self._completion_kwargs(context)
        if query_batch:
            return await generate_completion_batch(query_batch=query_batch, **kwargs)
        completion = await generate_completion(query=query, **kwargs)
        return [completion]

    async def _append_graph_evidence(self, completions: List[Any]) -> List[Any]:
        """Append an answer-grounded chunk Evidence block to string completions.

        Each answer is run as a vector query against the chunk index, so the
        Evidence bullets reflect where the answer text is grounded in the corpus
        rather than which graph elements happened to be retrieved. Evidence is
        appended only when references are enabled and the completion is a plain
        string (never corrupt a structured response_model); search failures
        degrade to no Evidence.
        """
        return await append_answer_grounded_evidence(
            completions,
            enabled=self.include_references and self.response_model is str,
        )

    async def get_completion_from_context(
        self,
        query: Optional[str] = None,
        query_batch: Optional[List[str]] = None,
        retrieved_objects: Optional[List[Edge]] = None,
        context: str = None,
        effective_query: Optional[str] = None,
        turn_preparation=None,
    ) -> List[Any]:
        """
        Generates an LLM response based on the query, context, and conversation history.
        Optionally saves the interaction and updates the session cache.

        Args:
            query (str): The user's question or prompt.
            query_batch (List[str]): The batch of user queries.
            retrieved_objects (Optional[List[Edge]]): Raw triplets used for interaction mapping.
                                                     Output of get_retrieved_objects method.
            context (str): The text-resolved graph context.
                           Output of the get_context_from_objects method.

        Returns:
            List[Any]: A list containing the generated response (completion).

        Note: To avoid duplicate retrievals, ensure that retrieved_objects and context
              are provided from previous method calls.
        """
        use_session = self._use_session_cache() and not query_batch
        if use_session:
            sm = get_session_manager()
            used_graph_element_ids = self._extract_context_object_ids(retrieved_objects)
            completion = await sm.generate_completion_with_session(
                session_id=self.session_id,
                query=query,
                context=context,
                user_prompt_path=self.user_prompt_path,
                system_prompt_path=self.system_prompt_path,
                system_prompt=self.system_prompt,
                response_model=self.response_model,
                summarize_context=False,
                used_graph_element_ids=used_graph_element_ids,
                max_context_chars=getattr(self, "max_context_chars", None),
                effective_query=effective_query,
                turn_preparation=turn_preparation,
            )
            completions = [completion]
        else:
            completions = await self._generate_completion_without_session(
                query, query_batch, context
            )

        # Session and non-session branches rejoin here so every variant that calls
        # this method (including via super()) appends references once. Evidence is
        # grounded in each completion's own text, so a cache-hit answer never
        # cites chunks that share nothing with it.
        return await self._append_graph_evidence(completions)

    async def get_completion(
        self, query: Optional[str] = None, query_batch: Optional[List[str]] = None
    ) -> List[Any]:
        """
        Generates a final output or answer based on the query and retrieved context.

        Args:
            query (str): The original user query.
            query_batch (List[str]): The batch of user queries.

        Returns:
            List[Any]: A list containing the generated completions or response objects.
        """
        validate_retriever_input(query, query_batch)

        effective_query = query
        turn_preparation = None
        if query is not None and not query_batch:
            turn_preparation = await self.prepare_session_turn_for_retrieval(query)
            if not turn_preparation.should_answer:
                return [turn_preparation.response_to_user or "Got it."]
            effective_query = turn_preparation.effective_query or query

        retrieved_objects = await self.get_retrieved_objects(
            query=effective_query,
            query_batch=query_batch,
        )
        context = await self.get_context_from_objects(
            query=effective_query,
            query_batch=query_batch,
            retrieved_objects=retrieved_objects,
        )
        completion = await self.get_completion_from_context(
            query=query,
            query_batch=query_batch,
            retrieved_objects=retrieved_objects,
            context=context,
            effective_query=effective_query,
            turn_preparation=turn_preparation,
        )

        return completion