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
```
This commit is contained in:
James Brunton
2026-05-01 14:11:54 +01:00
committed by GitHub
parent 5605062153
commit 5541dd666c
48 changed files with 1067 additions and 534 deletions
@@ -14,6 +14,7 @@ from pydantic import ValidationError
from stirling.agents.math_presentation import extract_math_verdict
from stirling.contracts import (
AiFile,
ExtractedFileText,
ExtractedTextArtifact,
MathAuditorToolReportArtifact,
@@ -21,6 +22,7 @@ from stirling.contracts import (
WorkflowArtifact,
)
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
from stirling.models import FileId
def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
@@ -42,7 +44,7 @@ def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest:
return OrchestratorRequest(
user_message="review the math",
file_names=["report.pdf"],
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
artifacts=artifacts,
)
@@ -23,8 +23,9 @@ from unittest.mock import AsyncMock, patch
import pytest
from stirling.agents import OrchestratorAgent
from stirling.contracts import OrchestratorRequest
from stirling.contracts import AiFile, OrchestratorRequest
from stirling.contracts.pdf_edit import EditPlanResponse
from stirling.models import FileId
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
from stirling.services.runtime import AppRuntime
@@ -39,7 +40,7 @@ async def test_delegate_pdf_review_wires_prompt_to_tool_step(runtime: AppRuntime
orchestrator = OrchestratorAgent(runtime)
request = OrchestratorRequest(
user_message="please add review comments flagging ambiguous dates",
file_names=["contract.pdf"],
files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
)
ctx = SimpleNamespace(deps=_FakeDeps(request=request))
@@ -12,12 +12,15 @@ import pytest
from stirling.agents.pdf_questions import _MATH_SYNTH_SYSTEM_PROMPT, PdfQuestionAgent
from stirling.contracts import (
AiFile,
EditPlanResponse,
MathAuditorToolReportArtifact,
OrchestratorRequest,
PdfQuestionAnswerResponse,
SupportedCapability,
)
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
from stirling.models import FileId
from stirling.models.agent_tool_models import AgentToolId
from stirling.services.runtime import AppRuntime
@@ -49,25 +52,23 @@ def _make_verdict() -> Verdict:
@pytest.mark.anyio
async def test_orchestrate_classifier_true_embeds_plan_in_answer(runtime: AppRuntime) -> None:
"""First turn — classifier says math; the response is a PdfQuestionAnswerResponse
with the math-auditor plan attached as a nullable ``edit_plan`` field. The
answer is empty on this turn; the caller runs the embedded plan and resumes."""
async def test_orchestrate_classifier_true_returns_math_audit_plan(runtime: AppRuntime) -> None:
"""First turn — classifier says math; the response is an EditPlanResponse
(``outcome=PLAN``) with ``resume_with=PDF_QUESTION``. The caller runs the
plan and re-invokes the orchestrator with the verdict in artifacts."""
agent = PdfQuestionAgent(runtime)
request = OrchestratorRequest(
user_message="ist die mathematik korrekt?",
file_names=["report.pdf"],
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
)
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
response = await agent.orchestrate(request)
assert isinstance(response, PdfQuestionAnswerResponse)
assert response.answer == ""
assert response.edit_plan is not None
assert response.edit_plan.resume_with == SupportedCapability.PDF_QUESTION
assert len(response.edit_plan.steps) == 1
assert response.edit_plan.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
assert isinstance(response, EditPlanResponse)
assert response.resume_with == SupportedCapability.PDF_QUESTION
assert len(response.steps) == 1
assert response.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
@pytest.mark.anyio
@@ -81,7 +82,7 @@ async def test_orchestrate_resume_synthesises_answer_without_calling_classifier(
verdict = _make_verdict()
request = OrchestratorRequest(
user_message="ist die mathematik korrekt?",
file_names=["report.pdf"],
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
)
canned_answer = "Die Summe stimmt nicht: angegeben $215,000, erwartet $215,500."
+8 -5
View File
@@ -20,9 +20,9 @@ from stirling.agents.pdf_review import (
_LocalisedComment,
_LocalisedVerdict,
)
from stirling.contracts import EditPlanResponse, OrchestratorRequest, SupportedCapability
from stirling.contracts import AiFile, EditPlanResponse, OrchestratorRequest, SupportedCapability
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
from stirling.models import ToolEndpoint
from stirling.models import FileId, ToolEndpoint
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
from stirling.services.runtime import AppRuntime
@@ -144,7 +144,10 @@ async def test_orchestrate_classifier_true_emits_math_audit_plan(runtime: AppRun
"""First turn — when the math-intent classifier says yes, emit a one-step plan
calling the math auditor with resume_with=PDF_REVIEW."""
agent = PdfReviewAgent(runtime)
request = OrchestratorRequest(user_message="vérifie les totaux", file_names=["report.pdf"])
request = OrchestratorRequest(
user_message="vérifie les totaux",
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
)
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
response = await agent.orchestrate(request)
@@ -161,7 +164,7 @@ async def test_orchestrate_classifier_false_routes_to_pdf_comment_agent(runtime:
agent = PdfReviewAgent(runtime)
request = OrchestratorRequest(
user_message="review the invoices for ambiguous wording",
file_names=["contract.pdf"],
files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
)
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=False)):
@@ -187,7 +190,7 @@ async def test_orchestrate_resume_uses_verdict_without_calling_classifier(
verdict = _make_verdict([_discrepancy(page=0, stated="$100")])
request = OrchestratorRequest(
user_message="flag math errors",
file_names=["report.pdf"],
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
)
canned = _LocalisedVerdict(