mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 11:23:10 +02:00
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
+108
-108
@@ -6,14 +6,13 @@ import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from stirling.api import app
|
||||
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
|
||||
from stirling.api.dependencies import get_rag_service
|
||||
from stirling.models import FileId
|
||||
from stirling.rag import Document, RagService, SqliteVecStore
|
||||
|
||||
TEST_EMBEDDING_MODEL = "test-embedder"
|
||||
|
||||
|
||||
class StubEmbedder:
|
||||
"""Deterministic embeddings for route tests — no network, no provider needed."""
|
||||
"""Deterministic embeddings for route tests: no network, no provider needed."""
|
||||
|
||||
def __init__(self, dim: int = 8) -> None:
|
||||
self._dim = dim
|
||||
@@ -53,153 +52,154 @@ def _build_service() -> RagService:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> Iterator[TestClient]:
|
||||
service = _build_service()
|
||||
def service() -> RagService:
|
||||
return _build_service()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(service: RagService) -> Iterator[TestClient]:
|
||||
app.dependency_overrides[get_rag_service] = lambda: service
|
||||
app.dependency_overrides[get_rag_embedding_model] = lambda: TEST_EMBEDDING_MODEL
|
||||
try:
|
||||
yield TestClient(app)
|
||||
finally:
|
||||
app.dependency_overrides.pop(get_rag_service, None)
|
||||
app.dependency_overrides.pop(get_rag_embedding_model, None)
|
||||
|
||||
|
||||
# ── /status ─────────────────────────────────────────────────────────────
|
||||
# ── POST /documents ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_status_reports_embedding_model_and_collections(client: TestClient) -> None:
|
||||
def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "doc-123",
|
||||
"source": "report.pdf",
|
||||
"pageText": [
|
||||
{"pageNumber": 1, "text": "The introduction covers the main topic."},
|
||||
{"pageNumber": 2, "text": "The conclusion summarises the findings."},
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["documentId"] == "doc-123"
|
||||
assert body["chunksIndexed"] >= 2
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None:
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "my-docs", "text": "Hello world.", "source": "a.pdf"},
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "replace-me",
|
||||
"source": "replace-me.pdf",
|
||||
"pageText": [{"pageNumber": 1, "text": "Original content that existed before."}],
|
||||
},
|
||||
)
|
||||
response = client.get("/api/v1/rag/status")
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["embeddingModel"] == TEST_EMBEDDING_MODEL
|
||||
assert "my-docs" in body["collections"]
|
||||
|
||||
|
||||
def test_status_when_empty(client: TestClient) -> None:
|
||||
response = client.get("/api/v1/rag/status")
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body == {"embeddingModel": TEST_EMBEDDING_MODEL, "collections": []}
|
||||
|
||||
|
||||
# ── /index ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_index_returns_chunk_count(client: TestClient) -> None:
|
||||
# Second ingest with different content should replace the first entirely
|
||||
response = client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "indexed", "text": "Short text.", "source": "doc.pdf"},
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "replace-me",
|
||||
"source": "replace-me.pdf",
|
||||
"pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["collection"] == "indexed"
|
||||
assert body["chunksIndexed"] >= 1
|
||||
|
||||
results = await service.search("New content", collection=FileId("replace-me"), top_k=5)
|
||||
texts = [r.document.text for r in results]
|
||||
assert any("New content" in t for t in texts)
|
||||
assert not any("Original content" in t for t in texts)
|
||||
|
||||
|
||||
def test_index_rejects_empty_collection_name(client: TestClient) -> None:
|
||||
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "", "text": "Text.", "source": "x.pdf"},
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "mixed",
|
||||
"source": "mixed.pdf",
|
||||
"pageText": [
|
||||
{"pageNumber": 1, "text": " "},
|
||||
{"pageNumber": 2, "text": "Real content on page 2."},
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["chunksIndexed"] >= 1
|
||||
|
||||
|
||||
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
|
||||
response = client.post("/api/v1/rag/documents", json={"documentId": "empty", "source": "empty.pdf"})
|
||||
assert response.status_code == 200
|
||||
assert response.json()["chunksIndexed"] == 0
|
||||
|
||||
|
||||
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/documents",
|
||||
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_index_rejects_oversized_text(client: TestClient) -> None:
|
||||
huge = "x" * 1_000_001 # Just over the 1MB cap
|
||||
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "toobig", "text": huge},
|
||||
"/api/v1/rag/documents",
|
||||
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# ── /search ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_search_returns_results(client: TestClient) -> None:
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "search-test", "text": "Python is fun.", "source": "guide.pdf"},
|
||||
)
|
||||
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/search",
|
||||
json={"query": "Python", "collection": "search-test", "topK": 3},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["query"] == "Python"
|
||||
assert len(body["results"]) >= 1
|
||||
first = body["results"][0]
|
||||
assert first["source"] == "guide.pdf"
|
||||
assert "score" in first
|
||||
|
||||
|
||||
def test_search_rejects_empty_collection_name(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/search",
|
||||
json={"query": "anything", "collection": ""},
|
||||
"/api/v1/rag/documents",
|
||||
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_search_without_collection_searches_all(client: TestClient) -> None:
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "col-one", "text": "Alpha content.", "source": "one.pdf"},
|
||||
)
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "col-two", "text": "Beta content.", "source": "two.pdf"},
|
||||
)
|
||||
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
|
||||
response = client.post(
|
||||
"/api/v1/rag/search",
|
||||
json={"query": "content"},
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "bad-page",
|
||||
"source": "bad-page.pdf",
|
||||
"pageText": [{"pageNumber": 0, "text": "something"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert len(body["results"]) >= 1
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
# ── /collections ────────────────────────────────────────────────────────
|
||||
# ── DELETE /documents/{id} ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_collections_empty_when_no_data(client: TestClient) -> None:
|
||||
response = client.get("/api/v1/rag/collections")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"collections": []}
|
||||
|
||||
|
||||
def test_collections_lists_indexed(client: TestClient) -> None:
|
||||
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "list-me", "text": "Text.", "source": "x.pdf"},
|
||||
"/api/v1/rag/documents",
|
||||
json={
|
||||
"documentId": "to-delete",
|
||||
"source": "to-delete.pdf",
|
||||
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
||||
},
|
||||
)
|
||||
response = client.get("/api/v1/rag/collections")
|
||||
response = client.delete("/api/v1/rag/documents/to-delete")
|
||||
assert response.status_code == 200
|
||||
assert "list-me" in response.json()["collections"]
|
||||
assert response.json() == {"documentId": "to-delete", "deleted": True}
|
||||
|
||||
|
||||
# ── DELETE /collections/{name} ──────────────────────────────────────────
|
||||
def test_delete_document_is_idempotent(client: TestClient) -> None:
|
||||
response = client.delete("/api/v1/rag/documents/never-existed")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"documentId": "never-existed", "deleted": False}
|
||||
|
||||
|
||||
def test_delete_collection_removes_it(client: TestClient) -> None:
|
||||
@pytest.mark.anyio
|
||||
async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None:
|
||||
client.post(
|
||||
"/api/v1/rag/index",
|
||||
json={"collection": "to-delete", "text": "Text.", "source": "x.pdf"},
|
||||
"/api/v1/rag/documents",
|
||||
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
|
||||
)
|
||||
response = client.delete("/api/v1/rag/collections/to-delete")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "deleted", "collection": "to-delete"}
|
||||
|
||||
listing = client.get("/api/v1/rag/collections").json()
|
||||
assert "to-delete" not in listing["collections"]
|
||||
|
||||
|
||||
def test_delete_nonexistent_collection_is_idempotent(client: TestClient) -> None:
|
||||
response = client.delete("/api/v1/rag/collections/never-existed")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "deleted", "collection": "never-existed"}
|
||||
assert await service.has_collection(FileId("gone"))
|
||||
client.delete("/api/v1/rag/documents/gone")
|
||||
assert not await service.has_collection(FileId("gone"))
|
||||
|
||||
Reference in New Issue
Block a user