mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 18:44:05 +02:00
# 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
```
205 lines
7.2 KiB
Python
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)
|