Stirling-PDF/engine/src/stirling/agents/pdf_questions.py

from __future__ import annotations

import logging

from pydantic_ai import Agent
from pydantic_ai.output import NativeOutput

from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
from stirling.contracts import (
    AiFile,
    EditPlanResponse,
    NeedIngestResponse,
    OrchestratorRequest,
    PdfContentType,
    PdfQuestionAnswerResponse,
    PdfQuestionNotFoundResponse,
    PdfQuestionOrchestrateResponse,
    PdfQuestionRequest,
    PdfQuestionResponse,
    PdfQuestionTerminalResponse,
    SupportedCapability,
    ToolOperationStep,
    Verdict,
    format_conversation_history,
    format_file_names,
)
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
from stirling.rag import RagCapability
from stirling.services import AppRuntime

logger = logging.getLogger(__name__)


PDF_QUESTION_SYSTEM_PROMPT = (
    "You answer questions about PDF documents by retrieving relevant content with the "
    "search_knowledge tool. Use it before answering. Do not guess or use outside knowledge.\n"
    "\n"
    "The search_knowledge tool has a finite call budget per run. When it is no longer "
    "available, answer from what you have already retrieved.\n"
    "\n"
    "Guidelines:\n"
    "- Make targeted search_knowledge calls. Typically one or two is enough.\n"
    "- Answer from the retrieved text. If the retrieved content doesn't support a confident "
    "answer, return not_found.\n"
    "- For questions that would require reading the entire document end-to-end (e.g. "
    "'what's the shortest chapter', 'how many X are there'), return not_found.\n"
    "- Include a short list of evidence snippets (with page numbers where available) drawn "
    "from what search_knowledge returned.\n"
    "\n"
    "Writing the not_found reason:\n"
    "- The reason is shown directly to the end user, so write it in plain, friendly "
    "language. One or two short sentences.\n"
    "- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
    "'search_knowledge', or other implementation details.\n"
    "- Be honest about the actual limitation. For questions that require full-document "
    "analysis (shortest chapter, word counts, etc.), explain that the document is too "
    "long to analyse end-to-end: you can only look up specific passages, and that's "
    "not enough to compare every part of the document against every other.\n"
    "- For questions where the answer just isn't in the document, say so directly: "
    "'I couldn't find that information in the document.'\n"
    "- Do not make it sound like you're choosing not to answer. Be clear that it's "
    "a genuine constraint."
)

_MATH_SYNTH_SYSTEM_PROMPT = (
    "You are given a math-audit Verdict (structured JSON) and the user's "
    "original question. Answer the question in plain prose using only "
    "facts from the Verdict; do not invent figures or pages. "
    "Reply in the SAME LANGUAGE as the user's question. Keep the answer "
    "concise — a sentence or short paragraph. "
    "Quote any stated/expected numeric values from the Verdict verbatim — "
    "do not paraphrase, abbreviate, or convert units."
)


class PdfQuestionAgent:
    def __init__(self, runtime: AppRuntime) -> None:
        self.runtime = runtime
        self._math_synth_agent: Agent[None, str] = Agent(
            model=runtime.fast_model,
            output_type=str,
            system_prompt=_MATH_SYNTH_SYSTEM_PROMPT,
            model_settings=runtime.fast_model_settings,
        )
        self._math_intent_classifier = MathIntentClassifier(runtime)

    async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
        logger.info(
            "[pdf-question] handle: files=%s question=%r",
            [file.name for file in request.files],
            request.question,
        )
        missing = await self._find_missing_files(request.files)
        if missing:
            logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
            return NeedIngestResponse(
                resume_with=SupportedCapability.PDF_QUESTION,
                reason="Some files have not been ingested into RAG yet.",
                files_to_ingest=missing,
                content_types=[PdfContentType.PAGE_TEXT],
            )
        return await self._run_answer_agent(request)

    async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
        """Entry point for the orchestrator delegate.

        Decides math intent locally via a small classifier LLM (language-agnostic).
        On a math first turn, returns an :class:`EditPlanResponse` (``outcome=PLAN``)
        with ``resume_with=PDF_QUESTION`` so the caller runs the math specialist
        and re-invokes the orchestrator. On the resume turn, the captured
        :class:`Verdict` is digested into a localised prose answer. Non-math
        first turns fall through to the text-grounded :meth:`handle` pipeline.
        """
        verdict = extract_math_verdict(request)
        if verdict is not None:
            # Resume turn — Verdict in hand. Synthesise a localised answer from
            # the structured verdict via a small LLM that mirrors the user's
            # language; no English glue in the response.
            answer = await self._synthesise_math_answer(request.user_message, verdict)
            return PdfQuestionAnswerResponse(answer=answer, evidence=[])

        if await self._math_intent_classifier.classify(request.user_message):
            # First turn — emit a one-step plan calling the math specialist,
            # with resume_with set so the caller comes back with the verdict
            # in artifacts (handled by the resume branch above).
            return EditPlanResponse(
                summary="",
                steps=[
                    ToolOperationStep(
                        tool=AgentToolId.MATH_AUDITOR_AGENT,
                        parameters=MathAuditorAgentParams(),
                    )
                ],
                resume_with=SupportedCapability.PDF_QUESTION,
            )

        return await self.handle(
            PdfQuestionRequest(
                question=request.user_message,
                files=request.files,
                conversation_history=request.conversation_history,
            )
        )

    async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
        missing: list[AiFile] = []
        for file in files:
            if not await self.runtime.rag_service.has_collection(file.id):
                missing.append(file)
        return missing

    async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
        rag = RagCapability(
            rag_service=self.runtime.rag_service,
            collections=[file.id for file in request.files],
            top_k=self.runtime.settings.rag_default_top_k,
            max_searches=self.runtime.settings.rag_max_searches,
        )
        agent = Agent(
            model=self.runtime.smart_model,
            output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
            system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
            instructions=rag.instructions,
            toolsets=[rag.toolset],
            model_settings=self.runtime.smart_model_settings,
        )
        prompt = self._build_prompt(request)
        logger.debug("[pdf-question] prompt:\n%s", prompt)
        result = await agent.run(prompt)
        return result.output

    async def _synthesise_math_answer(self, user_message: str, verdict: Verdict) -> str:
        """Use a small LLM to render the structured Verdict as a natural-language
        answer in the same language as the user's question. The system prompt
        forbids invented figures; the LLM only restates Verdict facts.
        """
        prompt = f"User question:\n{user_message}\n\nMath audit Verdict (JSON):\n{verdict.model_dump_json()}"
        result = await self._math_synth_agent.run(prompt)
        return result.output

    def _build_prompt(self, request: PdfQuestionRequest) -> str:
        history = format_conversation_history(request.conversation_history)
        return (
            f"Conversation history:\n{history}\n"
            f"Files: {format_file_names(request.files)}\n"
            f"Question: {request.question}\n"
            "Use search_knowledge to retrieve the relevant content, then answer."
        )