mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
James BruntonandGitHub
5541dd666c
# 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
```
109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
"""Tests for ``stirling.agents.math_presentation``.
|
|
|
|
Only one helper lives in this module now: Verdict-artifact extraction
|
|
on the resume turn. Math intent itself is decided by the orchestrator's
|
|
top-level LLM and passed in as a flag, so there's no English regex to
|
|
test here. Verdict → prose / sticky-note text are the consumer agents'
|
|
responsibility — those projections are tested with each consumer.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from stirling.agents.math_presentation import extract_math_verdict
|
|
from stirling.contracts import (
|
|
AiFile,
|
|
ExtractedFileText,
|
|
ExtractedTextArtifact,
|
|
MathAuditorToolReportArtifact,
|
|
OrchestratorRequest,
|
|
WorkflowArtifact,
|
|
)
|
|
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
|
from stirling.models import FileId
|
|
|
|
|
|
def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
|
|
return Verdict(
|
|
session_id="s1",
|
|
discrepancies=discrepancies,
|
|
pages_examined=[d.page for d in discrepancies] or [0],
|
|
rounds_taken=1,
|
|
summary="Test verdict.",
|
|
clean=not discrepancies,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Resume-turn round-trip — ToolReportArtifact → Verdict
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest:
|
|
return OrchestratorRequest(
|
|
user_message="review the math",
|
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
|
artifacts=artifacts,
|
|
)
|
|
|
|
|
|
def test_extract_math_verdict_roundtrips_a_math_auditor_report() -> None:
|
|
"""When the math auditor has already run, Java re-enters the orchestrator with
|
|
a ToolReportArtifact carrying the serialised Verdict; the meta-agent's first
|
|
job on the resume turn is to hydrate that back into a Verdict."""
|
|
original = _make_verdict(
|
|
[
|
|
Discrepancy(
|
|
page=0,
|
|
kind=DiscrepancyKind.TALLY,
|
|
severity=Severity.ERROR,
|
|
description="Total mismatch.",
|
|
stated="$215,000",
|
|
expected="$215,500",
|
|
context="Revenue row",
|
|
)
|
|
]
|
|
)
|
|
artifact = MathAuditorToolReportArtifact(report=original)
|
|
request = _orchestrator_request_with_artifacts([artifact])
|
|
|
|
verdict = extract_math_verdict(request)
|
|
|
|
assert verdict is not None
|
|
assert len(verdict.discrepancies) == 1
|
|
assert verdict.discrepancies[0].stated == "$215,000"
|
|
assert verdict.discrepancies[0].expected == "$215,500"
|
|
|
|
|
|
def test_extract_math_verdict_returns_none_when_no_artifacts_present() -> None:
|
|
"""First turn — the plan has not yet run, so artifacts is empty."""
|
|
request = _orchestrator_request_with_artifacts([])
|
|
assert extract_math_verdict(request) is None
|
|
|
|
|
|
def test_extract_math_verdict_ignores_other_artifact_kinds() -> None:
|
|
"""Only MathAuditorToolReportArtifact counts. Other artifact kinds (e.g.
|
|
extracted page text from a NeedContent round-trip) must be ignored here so
|
|
meta-agents don't misinterpret them as math reports."""
|
|
unrelated = ExtractedTextArtifact(
|
|
files=[ExtractedFileText(file_name="report.pdf", pages=[])],
|
|
)
|
|
request = _orchestrator_request_with_artifacts([unrelated])
|
|
assert extract_math_verdict(request) is None
|
|
|
|
|
|
def test_malformed_math_auditor_report_is_rejected_at_validation_time() -> None:
|
|
"""The discriminated-union contract validates the report payload as a
|
|
:class:`Verdict` on receipt — a corrupt body raises at construction time
|
|
rather than silently surviving until the meta-agent tries to read it."""
|
|
with pytest.raises(ValidationError):
|
|
MathAuditorToolReportArtifact.model_validate(
|
|
{
|
|
"kind": "tool_report",
|
|
"source_tool": "math_auditor_agent",
|
|
"report": {"not_a_verdict_field": "garbage"},
|
|
}
|
|
)
|