arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-15 11:00:47 +02:00
Code Issues Actions 3 Packages Projects Releases 1 Wiki Activity

Flesh out RAG system (#6197)

Browse Source 
# 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
Download Patch File Download Diff File Expand all files Collapse all files
engine/tests/test_rag.py
+41 -16
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import pytest
from stirling.models import FileId
from stirling.rag.capability import RagCapability
from stirling.rag.chunker import chunk_text
from stirling.rag.service import RagService
@@ -163,38 +164,38 @@ class TestRagService:
@pytest.mark.anyio
async def test_index_and_search(self, rag_service: RagService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await rag_service.index_text("docs", text, source="guide.pdf")
count = await rag_service.index_text(FileId("docs"), text, source="guide.pdf")
assert count > 0
results = await rag_service.search("Python libraries", collection="docs")
results = await rag_service.search("Python libraries", collection=FileId("docs"))
assert len(results) > 0
assert results[0].document.text # non-empty text
@pytest.mark.anyio
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
count = await rag_service.index_text("docs", "", source="empty.pdf")
count = await rag_service.index_text(FileId("docs"), "", source="empty.pdf")
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
results = await rag_service.search("anything", collection="nonexistent")
results = await rag_service.search("anything", collection=FileId("nonexistent"))
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("col-a", "Machine learning overview.", source="ml.pdf")
await rag_service.index_text("col-b", "Deep learning with neural networks.", source="dl.pdf")
await rag_service.index_text(FileId("col-a"), "Machine learning overview.", source="ml.pdf")
await rag_service.index_text(FileId("col-b"), "Deep learning with neural networks.", source="dl.pdf")
results = await rag_service.search("neural networks")
assert len(results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, rag_service: RagService) -> None:
await rag_service.index_text("temp", "Temporary data.", source="tmp.pdf")
await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf")
collections = await rag_service.list_collections()
assert "temp" in collections
await rag_service.delete_collection("temp")
await rag_service.delete_collection(FileId("temp"))
collections = await rag_service.list_collections()
assert "temp" not in collections
@@ -214,7 +215,7 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service, collections=["docs", "manuals"])
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
@@ -227,8 +228,8 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("col-a", "Alpha content.", source="a.pdf")
await rag_service.index_text("col-b", "Beta content.", source="b.pdf")
await rag_service.index_text(FileId("col-a"), "Alpha content.", source="a.pdf")
await rag_service.index_text(FileId("col-b"), "Beta content.", source="b.pdf")
cap = RagCapability(rag_service)
instructions_fn = cap.instructions
assert callable(instructions_fn)
@@ -252,7 +253,7 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
await rag_service.index_text("docs", "Python is a programming language.", source="guide.pdf")
await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf")
cap = RagCapability(rag_service)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
@@ -262,10 +263,10 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
await rag_service.index_text("pinned", "Pinned collection content.", source="pinned.pdf")
await rag_service.index_text("other", "Content in another collection.", source="other.pdf")
await rag_service.index_text(FileId("pinned"), "Pinned collection content.", source="pinned.pdf")
await rag_service.index_text(FileId("other"), "Content in another collection.", source="other.pdf")
cap = RagCapability(rag_service, collections=["pinned"])
cap = RagCapability(rag_service, collections=[FileId("pinned")])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@@ -273,7 +274,7 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await rag_service.index_text("bulk", paragraphs, source="bulk.pdf")
await rag_service.index_text(FileId("bulk"), paragraphs, source="bulk.pdf")
cap = RagCapability(rag_service)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
@@ -281,3 +282,27 @@ class TestRagCapability:
assert "[Result 1" in output
assert "[Result 2" in output
assert "[Result 3" not in output
@pytest.mark.anyio
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, rag_service: RagService) -> None:
"""The prepare callback must return None once max_searches has been reached
so the agent can no longer call the tool on subsequent turns."""
await rag_service.index_text(FileId("docs"), "Some content.", source="x.pdf")
cap = RagCapability(rag_service, max_searches=2)
tool_def = _dummy_tool_def()
# Budget intact: prepare returns the tool definition.
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
# Use the budget.
await _invoke_search_knowledge(cap, "content")
await _invoke_search_knowledge(cap, "content")
# Budget spent: prepare returns None, removing the tool from the agent's next turn.
assert await cap._prepare_search_knowledge(None, tool_def) is None # type: ignore[arg-type]
def _dummy_tool_def() -> object:
"""Sentinel passed to ``_prepare_search_knowledge``. The callback only inspects
``_search_count``; it doesn't read anything off the tool_def or context."""
return object()