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

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

from pydantic import BaseModel

from cognee.base_config import get_base_config
from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
from cognee.modules.retrieval.utils.query_state import QueryState
from cognee.modules.retrieval.utils.validate_queries import validate_retriever_input
from cognee.shared.logging_utils import get_logger
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
from cognee.modules.retrieval.utils.completion import (
    batch_llm_completion,
    generate_completion_batch,
)
from cognee.infrastructure.session.get_session_manager import get_session_manager
from cognee.context_global_variables import session_user
from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt

logger = get_logger()


def _as_answer_text(completion: Any) -> str:
    """Convert completion to human-readable text for validation and follow-up prompts."""
    if isinstance(completion, str):
        return completion
    if isinstance(completion, BaseModel):
        json_str = completion.model_dump_json(indent=2)
        return f"[Structured Response]\n{json_str}"
    try:
        return json.dumps(completion, indent=2)
    except TypeError:
        return str(completion)


class GraphCompletionCotRetriever(GraphCompletionRetriever):
    """
    Handles graph completion by generating responses based on a series of interactions with
    a language model. This class extends from GraphCompletionRetriever and is designed to
    manage the retrieval and validation process for user queries, integrating follow-up
    questions based on reasoning. The public methods are:

    - get_completion

    Instance variables include:
    - validation_system_prompt_path
    - validation_user_prompt_path
    - followup_system_prompt_path
    - followup_user_prompt_path
    """

    def __init__(
        self,
        user_prompt_path: str = "graph_context_for_question.txt",
        system_prompt_path: str = "answer_simple_question.txt",
        validation_user_prompt_path: str = "cot_validation_user_prompt.txt",
        validation_system_prompt_path: str = "cot_validation_system_prompt.txt",
        followup_system_prompt_path: str = "cot_followup_system_prompt.txt",
        followup_user_prompt_path: str = "cot_followup_user_prompt.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,
        max_iter: int = 4,
        session_id: Optional[str] = None,
        response_model: Type = str,
        neighborhood_depth: Optional[int] = None,
        neighborhood_seed_top_k: Optional[int] = 10,
        include_references: bool = False,
    ):
        super().__init__(
            user_prompt_path=user_prompt_path,
            system_prompt_path=system_prompt_path,
            system_prompt=system_prompt,
            top_k=top_k,
            node_type=node_type,
            node_name=node_name,
            node_name_filter_operator=node_name_filter_operator,
            wide_search_top_k=wide_search_top_k,
            triplet_distance_penalty=triplet_distance_penalty,
            feedback_influence=feedback_influence,
            session_id=session_id,
            response_model=response_model,
            neighborhood_depth=neighborhood_depth,
            neighborhood_seed_top_k=neighborhood_seed_top_k,
            include_references=include_references,
        )
        self.validation_system_prompt_path = validation_system_prompt_path
        self.validation_user_prompt_path = validation_user_prompt_path
        self.followup_system_prompt_path = followup_system_prompt_path
        self.followup_user_prompt_path = followup_user_prompt_path
        self._cot_final_context = None
        self.max_iter = max_iter

    async def get_retrieved_objects(
        self, query: Optional[str] = None, query_batch: Optional[List[str]] = None
    ) -> Union[List[Edge], List[List[Edge]]]:
        """
        Run chain-of-thought completion with optional structured output.

        Parameters:
        -----------
            - query: User query

        Returns:
        --------
            - List of retrieved edges
        """
        session_save = self._use_session_cache()
        validate_retriever_input(query, query_batch, session_save)

        conversation_history = ""
        if session_save:
            user = session_user.get()
            user_id = getattr(user, "id", None)
            if user_id:
                sm = get_session_manager()
                history = await sm.get_session(
                    user_id=str(user_id),
                    session_id=self.session_id,
                    formatted=True,
                    last_n=sm.session_history_last_n,
                    include_context=False,
                )
                conversation_history = history if isinstance(history, str) else ""
                # Prepend the active session-context guidance block above the (history-only) CoT
                # intermediate-round context. Gated + fail-open: never blocks CoT reasoning.
                block = await self._maybe_active_context_block(sm, str(user_id), query)
                if block:
                    conversation_history = (
                        block + "\n\n" + conversation_history if conversation_history else block
                    )

        # Normalize single query to batch for uniform processing
        effective_batch = [query] if query else query_batch

        _completion, context_text, triplets = await self._run_cot_completion(
            effective_batch, conversation_history, skip_final_completion=True
        )

        self._cot_final_context = context_text

        if query:
            return triplets[0]
        return triplets

    async def _maybe_active_context_block(self, sm, user_id: str, query: Optional[str]) -> str:
        """Render the active session-context block for the CoT intermediate rounds.

        Gated on caching + auto_feedback. Fully fail-open: returns "" on any error or
        when the layer is disabled, so it can never block chain-of-thought reasoning.
        """
        try:
            from cognee.infrastructure.databases.cache.config import CacheConfig
            from cognee.infrastructure.session.session_context_builder import (
                build_active_context_block,
            )

            cache_config = CacheConfig()
            if not (cache_config.caching and cache_config.auto_feedback):
                return ""

            block, _served_ids = await build_active_context_block(
                session_manager=sm,
                user_id=user_id,
                session_id=sm._resolve_session_id(self.session_id),
                query=query or "",
            )
            return block or ""
        except Exception as exc:
            logger.warning("CoT active session-context block failed: %s", exc)
            return ""

    # -- CoT orchestrator --

    async def _run_cot_completion(
        self,
        query_batch: List[str],
        conversation_history: str = "",
        skip_final_completion: bool = False,
    ) -> tuple[List[Any], List[str], List[List[Edge]]]:
        """
        Run chain-of-thought completion with optional structured output.

        Parameters:
        -----------
            - query_batch: Batch of user queries
            - conversation_history: Optional conversation history string
            - skip_final_completion: If True, do not run _generate_completions on the last
              iteration; return [] for completions (final completion done in get_completion_from_context).

        Returns:
        --------
            - completion_result: The generated completion (string or structured model), or [] if skip_final_completion.
            - context_text: The resolved context text
            - triplets: The list of triplets used
        """
        states = {q: QueryState() for q in query_batch}
        await self._fetch_initial_triplets_and_context(states)
        await self._generate_completions(states, conversation_history)

        for reasoning_iteration in range(self.max_iter):
            followup_queries = await self._run_cot_round(states)
            await self._merge_followup_triplets(states, followup_queries)
            if not (skip_final_completion and reasoning_iteration == self.max_iter - 1):
                await self._generate_completions(states, conversation_history)

        return self._collect_results(states, query_batch, skip_final_completion)

    # -- Helper methods called by the orchestrator --

    async def _fetch_initial_triplets_and_context(self, states: dict):
        """Fetch triplets and resolve context text for all queries."""
        queries = list(states.keys())
        triplets_batch = await self.get_triplets_batch(queries)
        context_batch = await asyncio.gather(
            *[self.resolve_edges_to_text(t) for t in triplets_batch]
        )
        for q, triplets, context in zip(queries, triplets_batch, context_batch):
            states[q].triplets = triplets
            states[q].context_text = context

    async def _generate_completions(self, states: dict, conversation_history: str):
        """Generate completions for all queries in parallel."""
        queries = list(states.keys())
        contexts = [states[q].context_text for q in queries]
        completions = await generate_completion_batch(
            query_batch=queries,
            context=contexts,
            user_prompt_path=self.user_prompt_path,
            system_prompt_path=self.system_prompt_path,
            system_prompt=self.system_prompt,
            response_model=self.response_model,
            conversation_history=conversation_history if conversation_history else None,
        )
        for q, comp in zip(queries, completions):
            states[q].completion = comp
        logger.info(f"Chain-of-thought: generated completions for {len(queries)} queries")

    async def _run_cot_round(self, states: dict) -> List[str]:
        """Run one CoT round: validate answers, generate follow-up questions."""
        validation_prompts, validation_system = self._build_validation_prompts(states)
        reasoning_batch = await batch_llm_completion(validation_prompts, validation_system)

        followup_prompts, followup_system = self._build_followup_prompts(states, reasoning_batch)
        followup_questions = await batch_llm_completion(followup_prompts, followup_system)

        logger.info(f"Chain-of-thought: follow-up questions: {followup_questions}")
        return followup_questions

    # -- Prompt builders --

    def _build_cot_prompts(self, template_path, states, extras):
        """Build prompts with common query+answer fields and per-query extras."""
        return [
            render_prompt(
                filename=template_path,
                context={"query": q, "answer": _as_answer_text(states[q].completion), **extra},
            )
            for q, extra in zip(states.keys(), extras)
        ]

    def _build_validation_prompts(self, states):
        """Build validation user prompts and load system prompt."""
        system_prompt = read_query_prompt(prompt_file_name=self.validation_system_prompt_path)
        user_prompts = self._build_cot_prompts(
            self.validation_user_prompt_path,
            states,
            [{"context": s.context_text} for s in states.values()],
        )
        return user_prompts, system_prompt

    def _build_followup_prompts(self, states, reasoning_batch):
        """Build followup user prompts and load system prompt."""
        system_prompt = read_query_prompt(prompt_file_name=self.followup_system_prompt_path)
        user_prompts = self._build_cot_prompts(
            self.followup_user_prompt_path,
            states,
            [{"reasoning": r} for r in reasoning_batch],
        )
        return user_prompts, system_prompt

    async def _merge_followup_triplets(self, states: dict, followup_questions: List[str]):
        """Fetch triplets for follow-up questions and merge with existing state."""
        queries = list(states.keys())
        new_triplets_batch = await self.get_triplets_batch(followup_questions)

        for q, new_triplets in zip(queries, new_triplets_batch):
            states[q].merge_triplets(new_triplets)

        context_batch = await asyncio.gather(
            *[self.resolve_edges_to_text(states[q].triplets) for q in queries]
        )
        for q, context in zip(queries, context_batch):
            states[q].context_text = context

    def _collect_results(
        self,
        states: dict,
        query_batch: List[str],
        skip_final_completion: bool = False,
    ) -> tuple[List[Any], List[str], List[List[Edge]]]:
        """Extract final completions, context texts, and triplets from states."""
        completions = [] if skip_final_completion else [states[q].completion for q in query_batch]
        contexts = [states[q].context_text for q in query_batch]
        triplets = [states[q].triplets for q in query_batch]
        return completions, contexts, triplets

    async def get_context_from_objects(
        self,
        query: Optional[str] = None,
        query_batch: Optional[List[str]] = None,
        retrieved_objects=None,
    ) -> Union[str, List[str]]:
        """Return stored CoT final context when set; otherwise delegate to parent."""
        cot_context = getattr(self, "_cot_final_context", None)
        if cot_context is not None:
            if query_batch:
                return cot_context
            return cot_context[0] if cot_context else ""

        return await super().get_context_from_objects(
            query=query,
            query_batch=query_batch,
            retrieved_objects=retrieved_objects,
        )