mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 18:44:05 +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
```
309 lines
14 KiB
Python
309 lines
14 KiB
Python
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
|
|
from stirling.rag.sqlite_vec_store import SqliteVecStore
|
|
from stirling.rag.store import Document, SearchResult
|
|
|
|
# ── chunk_text ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestChunkText:
|
|
def test_empty_input_returns_empty(self) -> None:
|
|
assert chunk_text("") == []
|
|
assert chunk_text(" ") == []
|
|
|
|
def test_short_text_returns_single_chunk(self) -> None:
|
|
text = "Hello world."
|
|
chunks = chunk_text(text, chunk_size=100)
|
|
assert len(chunks) == 1
|
|
assert chunks[0] == "Hello world."
|
|
|
|
def test_splits_on_paragraph_boundaries(self) -> None:
|
|
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
|
|
chunks = chunk_text(text, chunk_size=30, overlap=0)
|
|
# Each paragraph fits in 30 chars, so they should be split
|
|
assert len(chunks) >= 2
|
|
assert "First paragraph." in chunks[0]
|
|
|
|
def test_long_text_produces_multiple_chunks(self) -> None:
|
|
text = " ".join(["word"] * 200)
|
|
chunks = chunk_text(text, chunk_size=100, overlap=10)
|
|
assert len(chunks) > 1
|
|
for chunk in chunks:
|
|
# Chunks may slightly exceed due to sentence boundary snapping
|
|
assert len(chunk) <= 200 # generous upper bound
|
|
|
|
def test_overlap_produces_shared_content(self) -> None:
|
|
sentences = [f"Sentence number {i}." for i in range(20)]
|
|
text = " ".join(sentences)
|
|
chunks = chunk_text(text, chunk_size=100, overlap=30)
|
|
if len(chunks) >= 2:
|
|
# After word-boundary snapping, the second chunk should share
|
|
# some content with the tail of the first chunk
|
|
words_in_first_tail = chunks[0].split()[-3:] # last 3 words
|
|
overlap_text = " ".join(words_in_first_tail)
|
|
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
|
|
|
|
|
|
# ── SqliteVecStore ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSqliteVecStore:
|
|
"""Each test gets its own ephemeral store to avoid cross-test dimension conflicts."""
|
|
|
|
@pytest.mark.anyio
|
|
async def test_add_and_search(self) -> None:
|
|
store = SqliteVecStore.ephemeral()
|
|
docs = [
|
|
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
|
|
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
|
|
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
|
|
]
|
|
# Simple 3-dimensional embeddings for testing
|
|
embeddings = [
|
|
[1.0, 0.0, 0.0],
|
|
[0.9, 0.1, 0.0],
|
|
[0.0, 0.0, 1.0],
|
|
]
|
|
await store.add_documents("test-col", docs, embeddings)
|
|
|
|
# Search with a query close to the programming-related docs
|
|
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
|
|
assert len(results) == 2
|
|
assert isinstance(results[0], SearchResult)
|
|
# The closest should be doc "1" (exact match on first dimension)
|
|
assert results[0].document.id == "1"
|
|
assert results[0].score > 0.5
|
|
|
|
@pytest.mark.anyio
|
|
async def test_list_and_has_collection(self) -> None:
|
|
store = SqliteVecStore.ephemeral()
|
|
docs = [Document(id="1", text="test", metadata={})]
|
|
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
|
|
|
|
collections = await store.list_collections()
|
|
assert "my-collection" in collections
|
|
assert await store.has_collection("my-collection") is True
|
|
assert await store.has_collection("nonexistent") is False
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_collection(self) -> None:
|
|
store = SqliteVecStore.ephemeral()
|
|
docs = [Document(id="1", text="test", metadata={})]
|
|
await store.add_documents("to-delete", docs, [[1.0]])
|
|
|
|
assert await store.has_collection("to-delete") is True
|
|
await store.delete_collection("to-delete")
|
|
assert await store.has_collection("to-delete") is False
|
|
|
|
@pytest.mark.anyio
|
|
async def test_search_empty_collection(self) -> None:
|
|
store = SqliteVecStore.ephemeral()
|
|
docs = [Document(id="1", text="test", metadata={})]
|
|
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
|
|
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
|
|
assert len(results) == 1
|
|
|
|
@pytest.mark.anyio
|
|
async def test_mismatched_docs_embeddings_raises(self) -> None:
|
|
store = SqliteVecStore.ephemeral()
|
|
docs = [Document(id="1", text="test", metadata={})]
|
|
with pytest.raises(ValueError, match="documents.*embeddings"):
|
|
await store.add_documents("bad", docs, [[1.0], [2.0]])
|
|
|
|
|
|
# ── RagService (with stub embedder) ────────────────────────────────────
|
|
|
|
|
|
class StubEmbeddingService:
|
|
"""A minimal stub that returns fixed-dimension embeddings for testing."""
|
|
|
|
def __init__(self, dim: int = 8) -> None:
|
|
self._dim = dim
|
|
|
|
async def embed_query(self, text: str) -> list[float]:
|
|
# Deterministic embedding based on hash of text
|
|
h = hash(text) % 1000
|
|
return [(h + i) / 1000.0 for i in range(self._dim)]
|
|
|
|
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
return [await self.embed_query(t) for t in texts]
|
|
|
|
def chunk_and_prepare(
|
|
self,
|
|
text: str,
|
|
source: str = "",
|
|
base_metadata: dict[str, str] | None = None,
|
|
) -> list[Document]:
|
|
from stirling.rag.chunker import chunk_text
|
|
|
|
chunks = chunk_text(text, 100, 10)
|
|
docs = []
|
|
for i, chunk in enumerate(chunks):
|
|
meta = dict(base_metadata) if base_metadata else {}
|
|
meta["source"] = source
|
|
meta["chunk_index"] = str(i)
|
|
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
|
|
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
|
|
return docs
|
|
|
|
|
|
@pytest.fixture
|
|
def rag_service() -> RagService:
|
|
"""Each RagService test gets its own fresh ephemeral store to avoid dimension conflicts."""
|
|
store = SqliteVecStore.ephemeral()
|
|
return RagService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
|
|
|
|
|
|
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(FileId("docs"), text, source="guide.pdf")
|
|
assert count > 0
|
|
|
|
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(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=FileId("nonexistent"))
|
|
assert results == []
|
|
|
|
@pytest.mark.anyio
|
|
async def test_search_all_collections(self, rag_service: RagService) -> None:
|
|
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(FileId("temp"), "Temporary data.", source="tmp.pdf")
|
|
collections = await rag_service.list_collections()
|
|
assert "temp" in collections
|
|
|
|
await rag_service.delete_collection(FileId("temp"))
|
|
collections = await rag_service.list_collections()
|
|
assert "temp" not in collections
|
|
|
|
|
|
# ── RagCapability ──────────────────────────────────────────────────────
|
|
|
|
|
|
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
|
|
"""Extract and call the search_knowledge tool function from a RagCapability's toolset."""
|
|
from pydantic_ai import FunctionToolset
|
|
|
|
toolset = capability.toolset
|
|
assert isinstance(toolset, FunctionToolset)
|
|
tool = toolset.tools["search_knowledge"]
|
|
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg] — pyright can't infer the generic tool function's kwargs
|
|
|
|
|
|
class TestRagCapability:
|
|
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
|
|
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
|
|
instructions = cap.instructions
|
|
assert isinstance(instructions, str)
|
|
assert "docs, manuals" in instructions
|
|
assert "search_knowledge" in instructions
|
|
|
|
def test_instructions_dynamic_when_no_collections(self, rag_service: RagService) -> None:
|
|
cap = RagCapability(rag_service)
|
|
instructions = cap.instructions
|
|
assert callable(instructions)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
|
|
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)
|
|
text = await instructions_fn()
|
|
assert "col-a" in text
|
|
assert "col-b" in text
|
|
|
|
@pytest.mark.anyio
|
|
async def test_dynamic_instructions_when_store_empty(self, rag_service: RagService) -> None:
|
|
cap = RagCapability(rag_service)
|
|
instructions_fn = cap.instructions
|
|
assert callable(instructions_fn)
|
|
text = await instructions_fn()
|
|
assert "empty" in text.lower()
|
|
|
|
@pytest.mark.anyio
|
|
async def test_search_knowledge_returns_no_results_message_when_empty(self, rag_service: RagService) -> None:
|
|
cap = RagCapability(rag_service)
|
|
output = await _invoke_search_knowledge(cap, "anything")
|
|
assert output == "No relevant results found in the knowledge base."
|
|
|
|
@pytest.mark.anyio
|
|
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
|
|
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
|
|
assert "source: guide.pdf" in output
|
|
assert "chunk:" in output
|
|
assert "relevance:" in output
|
|
|
|
@pytest.mark.anyio
|
|
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
|
|
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=[FileId("pinned")])
|
|
output = await _invoke_search_knowledge(cap, "content")
|
|
assert "pinned.pdf" in output
|
|
assert "other.pdf" not in output
|
|
|
|
@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(FileId("bulk"), paragraphs, source="bulk.pdf")
|
|
|
|
cap = RagCapability(rag_service)
|
|
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
|
|
# Only two results requested, so only Result 1 and Result 2 should appear
|
|
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()
|