Flesh out RAG system (#6197)

# Description of Changes
Flesh out the RAG system and connect it to the PDF Question Agent so it
can respond to questions about PDFs of an extremely large size.

I'd expect lots more work will need to be done to finish off the RAG
system to really be what we need, but this should be a reasonable start
which will let us connect it to tools and have the ingestion mostly
handled automatically. I'm leaving file deletion and proper file ID
management to be done in a future PR. We also need to consider whether
all tools should retrieve content exclusively via RAG, or whether it's
beneficial to have tools sometimes fetch the direct content and other
times fetch it from RAG.

A diagram of the expected interaction is as follows:

```mermaid
sequenceDiagram
    autonumber
    actor U as User
    participant FE as Frontend<br/>(ChatPanel)
    participant J as Java<br/>(AiWorkflowService)
    participant O as Engine:<br/>OrchestratorAgent
    participant QA as Engine:<br/>PdfQuestionAgent
    participant RAG as Engine:<br/>RagService + SqliteVecStore
    participant V as VoyageAI<br/>(embeddings)
    participant L as LLM<br/>(Claude / etc.)

    U->>FE: types "Summarise this PDF"<br/>(PDF already uploaded)
    FE->>J: POST /api/v1/ai/orchestrate/stream<br/>multipart: fileInputs[], userMessage
    Note over J: ByteHashFileIdStrategy<br/>id = sha256(bytes)[:16]
    J->>O: POST /api/v1/orchestrator<br/>{ files:[{id,name}], userMessage }

    O->>L: route via fast model
    L-->>O: delegate_pdf_question
    O->>QA: PdfQuestionRequest

    loop for each file
        QA->>RAG: has_collection(file.id)
        RAG-->>QA: false
    end
    QA-->>O: NeedIngestResponse(files_to_ingest)
    O-->>J: { outcome:"need_ingest", filesToIngest:[...] }

    Note over J: onNeedIngest
    loop per file
        J->>J: PDFBox: extract page text
        J->>O: POST /api/v1/rag/documents<br/>(long-running timeout)
        O->>RAG: chunk + stage documents
        O->>V: embed_documents (batches of 256)
        V-->>O: embeddings
        O->>RAG: add_documents
        O-->>J: { chunks_indexed: N }
    end

    Note over J: retry with resumeWith=pdf_question
    J->>O: POST /api/v1/orchestrator
    Note over O: fast-path to PdfQuestionAgent

    O->>QA: PdfQuestionRequest
    Note over QA: build RagCapability<br/>pinned to file IDs
    QA->>L: run(prompt) with search_knowledge tool

    loop up to max_searches
        L->>QA: search_knowledge(query)
        QA->>V: embed_query
        V-->>QA: query vector
        QA->>RAG: search(vector, collections=[file.id])
        RAG-->>QA: top-k chunks
        QA-->>L: formatted chunks
    end

    Note over QA: once budget spent,<br/>prepare() hides the tool
    L-->>QA: PdfQuestionAnswerResponse
    QA-->>O: answer
    O-->>J: { outcome:"answer", answer, evidence }
    J-->>FE: SSE "result"
    FE->>U: assistant bubble
```

engine/src/stirling/agents/orchestrator.py

@@ -18,10 +18,11 @@ from stirling.contracts import (
     OrchestratorRequest,
     OrchestratorResponse,
     PdfEditResponse,
-    PdfQuestionResponse,
+    PdfQuestionOrchestrateResponse,
     SupportedCapability,
     UnsupportedCapabilityResponse,
     format_conversation_history,
+    format_file_names,
 )
 from stirling.contracts.pdf_edit import EditPlanResponse
 from stirling.services import AppRuntime
@@ -78,12 +79,13 @@ class OrchestratorAgent:
                 "You are the top-level orchestrator. "
                 "Choose exactly one output function that best handles the request. "
                 "Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
-                "Use delegate_pdf_question for questions about PDF contents. "
+                "Use delegate_pdf_question for questions about the contents of the attached PDFs. "
                 "Use delegate_user_spec for requests to create or define an agent spec. "
                 "Use delegate_pdf_review when the user wants the PDF returned with review"
                 " comments attached — anything like 'review this', 'annotate with comments',"
                 " 'leave feedback on the PDF'. "
-                "Use unsupported_capability only when none of the other outputs fit."
+                "Use unsupported_capability when the user asks about the assistant itself "
+                "or when none of the other outputs fit; supply a helpful message."
             ),
             model_settings=runtime.fast_model_settings,
         )
@@ -91,7 +93,7 @@ class OrchestratorAgent:
     async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
         logger.info(
             "[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r",
-            request.file_names,
+            [file.name for file in request.files],
             request.resume_with,
             [type(a).__name__ for a in request.artifacts],
             request.user_message,
@@ -137,10 +139,10 @@ class OrchestratorAgent:
     async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
         return await PdfEditAgent(self.runtime).orchestrate(request)
 
-    async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
+    async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionOrchestrateResponse:
         return await self._run_pdf_question(ctx.deps.request)
 
-    async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse:
+    async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
         return await PdfQuestionAgent(self.runtime).orchestrate(request)
 
     async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
@@ -165,12 +167,11 @@ class OrchestratorAgent:
 
     def _build_prompt(self, request: OrchestratorRequest) -> str:
         artifact_summary = self._describe_artifacts(request)
-        file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
         history = format_conversation_history(request.conversation_history)
         return (
             f"Conversation history:\n{history}\n"
             f"User message: {request.user_message}\n"
-            f"Files: {file_names}\n"
+            f"Files: {format_file_names(request.files)}\n"
             f"Available artifacts:\n{artifact_summary}"
         )
 

engine/src/stirling/agents/pdf_edit.py

@@ -22,6 +22,7 @@ from stirling.contracts import (
     SupportedCapability,
     ToolOperationStep,
     format_conversation_history,
+    format_file_names,
 )
 from stirling.logging import Pretty
 from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
@@ -116,7 +117,6 @@ class PdfEditParameterSelector:
     ) -> str:
         operation_id = operation_plan[operation_index]
         operation_list = ", ".join(operation.name for operation in operation_plan)
-        file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
         generated_steps_text = (
             "\n".join(
                 f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
@@ -127,7 +127,7 @@ class PdfEditParameterSelector:
         return (
             f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
             f"User request: {request.user_message}\n"
-            f"Files: {file_names}\n"
+            f"Files: {format_file_names(request.files)}\n"
             f"Operation plan: {operation_list}\n"
             f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
             f"Selected operation: {operation_id.name}\n"
@@ -153,7 +153,7 @@ class PdfEditAgent:
         return await self.handle(
             PdfEditRequest(
                 user_message=request.user_message,
-                file_names=request.file_names,
+                files=request.files,
                 conversation_history=request.conversation_history,
                 page_text=extracted_text.files if extracted_text is not None else [],
             )
@@ -166,7 +166,7 @@ class PdfEditAgent:
     async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse:
         logger.info(
             "[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r",
-            request.file_names,
+            [file.name for file in request.files],
             has_page_text(request.page_text),
             allow_need_content,
             request.user_message,
@@ -225,11 +225,10 @@ class PdfEditAgent:
         )
 
     def _build_selection_prompt(self, request: PdfEditRequest) -> str:
-        file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
         return (
             f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
             f"User request: {request.user_message}\n"
-            f"Files: {file_names}\n"
+            f"Files: {format_file_names(request.files)}\n"
             f"Supported operations: {self._supported_operations_prompt()}\n"
             f"Extracted page text:\n{format_page_text(request.page_text)}"
         )
@@ -243,8 +242,7 @@ class PdfEditAgent:
         request: PdfEditRequest,
     ) -> NeedContentResponse:
         files = selection.files or [
-            NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT])
-            for file_name in request.file_names
+            NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in request.files
         ]
         return NeedContentResponse(
             resume_with=SupportedCapability.PDF_EDIT,

engine/src/stirling/agents/pdf_questions.py

@@ -1,32 +1,67 @@
 from __future__ import annotations
 
+import logging
+
 from pydantic_ai import Agent
 from pydantic_ai.output import NativeOutput
 
-from stirling.agents._page_text import (
-    format_page_text,
-    get_extracted_text_artifact,
-    has_page_text,
-)
 from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
 from stirling.contracts import (
+    AiFile,
     EditPlanResponse,
-    NeedContentFileRequest,
-    NeedContentResponse,
+    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 "
@@ -41,26 +76,6 @@ _MATH_SYNTH_SYSTEM_PROMPT = (
 class PdfQuestionAgent:
     def __init__(self, runtime: AppRuntime) -> None:
         self.runtime = runtime
-        rag = runtime.rag_capability
-        self.agent = Agent(
-            model=runtime.smart_model,
-            output_type=NativeOutput(
-                [
-                    PdfQuestionAnswerResponse,
-                    PdfQuestionNotFoundResponse,
-                ]
-            ),
-            system_prompt=(
-                "Answer questions about PDFs using only the extracted page text provided in the prompt. "
-                "Do not guess or use outside knowledge. "
-                "If the answer is not supported by the provided text, return not_found. "
-                "When answering, include a short list of evidence snippets with their page numbers. "
-                "Reply in the SAME LANGUAGE as the question."
-            ),
-            instructions=rag.instructions,
-            toolsets=[rag.toolset],
-            model_settings=runtime.smart_model_settings,
-        )
         self._math_synth_agent: Agent[None, str] = Agent(
             model=runtime.fast_model,
             output_type=str,
@@ -70,30 +85,31 @@ class PdfQuestionAgent:
         self._math_intent_classifier = MathIntentClassifier(runtime)
 
     async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
-        if not has_page_text(request.page_text):
-            return NeedContentResponse(
+        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="No extracted PDF page text was provided, so the question cannot be answered yet.",
-                files=[
-                    NeedContentFileRequest(
-                        file_name=file_name,
-                        content_types=[PdfContentType.PAGE_TEXT],
-                    )
-                    for file_name in request.file_names
-                ],
-                max_pages=self.runtime.settings.max_pages,
-                max_characters=self.runtime.settings.max_characters,
+                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) -> PdfQuestionResponse:
+    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, embeds an :class:`EditPlanResponse` in the answer
-        response; on the resume turn, digests the captured :class:`Verdict` into
-        a localised prose answer. Non-math first turns fall through to the
-        text-grounded :meth:`handle` pipeline.
+        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:
@@ -104,36 +120,53 @@ class PdfQuestionAgent:
             return PdfQuestionAnswerResponse(answer=answer, evidence=[])
 
         if await self._math_intent_classifier.classify(request.user_message):
-            # First turn — ask the caller to run the math specialist and come back.
-            # The plan rides on the answer response as a nullable member; ``answer``
-            # is empty on this turn and the caller resumes once the plan is run.
-            return PdfQuestionAnswerResponse(
-                answer="",
-                evidence=[],
-                edit_plan=EditPlanResponse(
-                    summary="",
-                    steps=[
-                        ToolOperationStep(
-                            tool=AgentToolId.MATH_AUDITOR_AGENT,
-                            parameters=MathAuditorAgentParams(),
-                        )
-                    ],
-                    resume_with=SupportedCapability.PDF_QUESTION,
-                ),
+            # 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,
             )
 
-        extracted_text = get_extracted_text_artifact(request)
         return await self.handle(
             PdfQuestionRequest(
                 question=request.user_message,
-                file_names=request.file_names,
-                page_text=extracted_text.files if extracted_text is not None else [],
+                files=request.files,
                 conversation_history=request.conversation_history,
             )
         )
 
-    async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
-        result = await self.agent.run(self._build_prompt(request))
+    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:
@@ -146,12 +179,10 @@ class PdfQuestionAgent:
         return result.output
 
     def _build_prompt(self, request: PdfQuestionRequest) -> str:
-        file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
-        pages = format_page_text(request.page_text, empty="")
         history = format_conversation_history(request.conversation_history)
         return (
             f"Conversation history:\n{history}\n"
-            f"Files: {file_names}\n"
+            f"Files: {format_file_names(request.files)}\n"
             f"Question: {request.question}\n"
-            f"Extracted page text:\n{pages}"
+            "Use search_knowledge to retrieve the relevant content, then answer."
         )