Files
Stirling-PDF/engine/tests/pdf_comment/test_routes.py
T
James BruntonandGitHub 5541dd666c 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
```
2026-05-01 14:11:54 +01:00

205 lines
7.2 KiB
Python

"""
PDF Comment Agent — FastAPI route tests.
Uses the FastAPI :class:`TestClient` with dependency overrides so the tests
exercise HTTP parsing, validation, and serialisation only — never the real
pydantic-ai agent.
"""
from __future__ import annotations
from collections.abc import Iterator
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from stirling.api import app
from stirling.api.dependencies import get_pdf_comment_agent
from stirling.config import AppSettings, RagBackend, load_settings
from stirling.contracts.pdf_comments import (
PdfCommentInstruction,
PdfCommentRequest,
PdfCommentResponse,
)
# ---------------------------------------------------------------------------
# Stubs
# ---------------------------------------------------------------------------
class StubSettingsProvider:
def __call__(self) -> AppSettings:
return AppSettings(
smart_model_name="test",
fast_model_name="test",
smart_model_max_tokens=8192,
fast_model_max_tokens=2048,
rag_backend=RagBackend.SQLITE,
rag_embedding_model="test-embed",
rag_store_path=Path(":memory:"),
rag_pgvector_dsn="",
rag_chunk_size=512,
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
max_pages=100,
max_characters=100_000,
posthog_enabled=False,
posthog_api_key="",
posthog_host="https://eu.i.posthog.com",
)
class StubPdfCommentAgent:
"""Stub that echoes the session id and returns a canned comment."""
def __init__(self, response: PdfCommentResponse | None = None) -> None:
self._response = response
self.generate_calls: list[PdfCommentRequest] = []
async def generate(self, request: PdfCommentRequest) -> PdfCommentResponse:
self.generate_calls.append(request)
if self._response is not None:
return self._response
return PdfCommentResponse(
session_id=request.session_id,
comments=[
PdfCommentInstruction(
chunk_id=request.chunks[0].id if request.chunks else "p0-c0",
comment_text="Stub comment.",
)
],
rationale="stubbed response",
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def stub_agent() -> StubPdfCommentAgent:
return StubPdfCommentAgent()
@pytest.fixture
def client(stub_agent: StubPdfCommentAgent) -> Iterator[TestClient]:
app.dependency_overrides[load_settings] = StubSettingsProvider()
app.dependency_overrides[get_pdf_comment_agent] = lambda: stub_agent
yield TestClient(app, raise_server_exceptions=False)
app.dependency_overrides.pop(load_settings, None)
app.dependency_overrides.pop(get_pdf_comment_agent, None)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _camel_request_body() -> dict[str, object]:
return {
"sessionId": "sess-1",
"userMessage": "flag dates",
"chunks": [
{
"id": "p0-c0",
"page": 0,
"x": 72.0,
"y": 700.0,
"width": 200.0,
"height": 12.0,
"text": "Signed on 5/6/2026",
}
],
}
def _snake_request_body() -> dict[str, object]:
return {
"session_id": "sess-snake",
"user_message": "flag dates",
"chunks": [
{
"id": "p0-c0",
"page": 0,
"x": 72.0,
"y": 700.0,
"width": 200.0,
"height": 12.0,
"text": "Snake case text",
}
],
}
# ---------------------------------------------------------------------------
# POST /api/v1/ai/pdf-comment-agent/generate
# ---------------------------------------------------------------------------
class TestGenerateEndpoint:
def test_camel_case_body_returns_200(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/pdf-comment-agent/generate", json=_camel_request_body())
assert resp.status_code == 200
def test_camel_case_body_response_has_expected_shape(self, client: TestClient) -> None:
resp = client.post("/api/v1/ai/pdf-comment-agent/generate", json=_camel_request_body())
body = resp.json()
assert body["sessionId"] == "sess-1"
assert isinstance(body["comments"], list)
assert len(body["comments"]) == 1
comment = body["comments"][0]
assert comment["chunkId"] == "p0-c0"
assert comment["commentText"] == "Stub comment."
assert "rationale" in body
def test_snake_case_body_is_still_accepted(self, client: TestClient) -> None:
"""ApiModel has validate_by_name=True & validate_by_alias=True, so snake_case
payloads must still be accepted."""
resp = client.post("/api/v1/ai/pdf-comment-agent/generate", json=_snake_request_body())
assert resp.status_code == 200
body = resp.json()
# Response is always serialised in camelCase regardless of request form.
assert body["sessionId"] == "sess-snake"
def test_missing_required_field_returns_422(self, client: TestClient) -> None:
body = _camel_request_body()
del body["sessionId"]
resp = client.post("/api/v1/ai/pdf-comment-agent/generate", json=body)
assert resp.status_code == 422
def test_agent_is_called_with_parsed_request(
self,
client: TestClient,
stub_agent: StubPdfCommentAgent,
) -> None:
client.post("/api/v1/ai/pdf-comment-agent/generate", json=_camel_request_body())
assert len(stub_agent.generate_calls) == 1
call = stub_agent.generate_calls[0]
assert call.session_id == "sess-1"
assert call.user_message == "flag dates"
assert len(call.chunks) == 1
assert call.chunks[0].id == "p0-c0"
def test_agent_exception_surfaces_as_500(self) -> None:
"""If the agent raises (LLM outage, auth failure, OOM), the route must
surface it as HTTP 500 so Java's AiEngineClient maps it to 502 — rather
than silently returning an empty/successful response that the Java caller
would mis-apply as 'zero comments to place'."""
class FailingAgent:
async def generate(self, _request: PdfCommentRequest) -> PdfCommentResponse:
raise RuntimeError("model provider unreachable")
app.dependency_overrides[load_settings] = StubSettingsProvider()
app.dependency_overrides[get_pdf_comment_agent] = lambda: FailingAgent()
try:
with TestClient(app, raise_server_exceptions=False) as failing_client:
resp = failing_client.post("/api/v1/ai/pdf-comment-agent/generate", json=_camel_request_body())
assert resp.status_code == 500
finally:
app.dependency_overrides.pop(load_settings, None)
app.dependency_overrides.pop(get_pdf_comment_agent, None)