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

import asyncio
from typing import Any, List, Optional, Tuple, Type

from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.infrastructure.llm.pipeline_stage import pipeline_stage
from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
from cognee.modules.observability import new_span, COGNEE_RESULT_SUMMARY


async def generate_completion(
    query: str,
    context: str,
    user_prompt_path: str,
    system_prompt_path: str,
    system_prompt: Optional[str] = None,
    conversation_history: Optional[str] = None,
    response_model: Type = str,
) -> Any:
    """Generates a completion using LLM with given context and prompts."""
    args = {"question": query, "context": context}
    user_prompt = render_prompt(user_prompt_path, args)
    system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path)

    if conversation_history:
        system_prompt = conversation_history + "\nTASK:" + system_prompt

    with pipeline_stage("query"):
        with new_span("cognee.llm.completion") as span:
            span.set_attribute("cognee.llm.prompt_path", system_prompt_path)
            span.set_attribute("cognee.llm.context_length", len(context))
            span.set_attribute("cognee.llm.query_length", len(query))
            result = await LLMGateway.acreate_structured_output(
                text_input=user_prompt,
                system_prompt=system_prompt,
                response_model=response_model,
            )
            if isinstance(result, str):
                span.set_attribute("cognee.llm.response_length", len(result))
            span.set_attribute(COGNEE_RESULT_SUMMARY, "LLM completion generated")
            return result


async def generate_completion_batch(
    query_batch: List[str],
    context: List[str],
    user_prompt_path: str,
    system_prompt_path: str,
    system_prompt: Optional[str] = None,
    conversation_history: Optional[str] = "",
    response_model: Type = str,
) -> List[Any]:
    """Generates completions for a batch of queries in parallel."""
    return await asyncio.gather(
        *[
            generate_completion(
                query=q,
                context=c,
                user_prompt_path=user_prompt_path,
                system_prompt_path=system_prompt_path,
                system_prompt=system_prompt,
                conversation_history=conversation_history,
                response_model=response_model,
            )
            for q, c in zip(query_batch, context)
        ]
    )


async def generate_session_completion_with_optional_summary(
    *,
    query: str,
    context: str,
    conversation_history: str,
    user_prompt_path: str,
    system_prompt_path: str,
    system_prompt: Optional[str] = None,
    response_model: Type = str,
    summarize_context: bool = False,
) -> Tuple[Any, str, Any]:
    """
    Run LLM completion (and optionally summarization) for the session-manager flow.
    Returns (completion, context_to_store, feedback_result).
    When summarize_context is True, context_to_store is the summarized context; otherwise "".
    Feedback analysis runs before retrieval/generation in SessionManager.prepare_session_turn.
    """
    if summarize_context:
        context_summary, completion = await asyncio.gather(
            summarize_text(context),
            generate_completion(
                query=query,
                context=context,
                user_prompt_path=user_prompt_path,
                system_prompt_path=system_prompt_path,
                system_prompt=system_prompt,
                conversation_history=conversation_history,
                response_model=response_model,
            ),
        )
        return (completion, context_summary, None)

    completion = await generate_completion(
        query=query,
        context=context,
        user_prompt_path=user_prompt_path,
        system_prompt_path=system_prompt_path,
        system_prompt=system_prompt,
        conversation_history=conversation_history,
        response_model=response_model,
    )
    return (completion, "", None)


async def batch_llm_completion(
    user_prompts: List[str],
    system_prompt: str,
    response_model: Type = str,
) -> List[Any]:
    """Run a batch of pre-built prompts through the LLM in parallel."""
    return list(
        await asyncio.gather(
            *[
                LLMGateway.acreate_structured_output(
                    text_input=prompt, system_prompt=system_prompt, response_model=response_model
                )
                for prompt in user_prompts
            ]
        )
    )


async def summarize_text(
    text: str,
    system_prompt_path: str = "summarize_search_results.txt",
    system_prompt: str = None,
) -> str:
    """Summarizes text using LLM with the specified prompt."""
    system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path)

    with new_span("cognee.llm.summarize") as span:
        span.set_attribute("cognee.llm.input_length", len(text))
        result = await LLMGateway.acreate_structured_output(
            text_input=text,
            system_prompt=system_prompt,
            response_model=str,
        )
        if isinstance(result, str):
            span.set_attribute("cognee.llm.response_length", len(result))
        span.set_attribute(COGNEE_RESULT_SUMMARY, "Text summarized")
        return result