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
```

engine/tests/agents/test_math_presentation.py

@@ -14,6 +14,7 @@ from pydantic import ValidationError
 
 from stirling.agents.math_presentation import extract_math_verdict
 from stirling.contracts import (
+    AiFile,
     ExtractedFileText,
     ExtractedTextArtifact,
     MathAuditorToolReportArtifact,
@@ -21,6 +22,7 @@ from stirling.contracts import (
     WorkflowArtifact,
 )
 from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
+from stirling.models import FileId
 
 
 def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
@@ -42,7 +44,7 @@ def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
 def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest:
     return OrchestratorRequest(
         user_message="review the math",
-        file_names=["report.pdf"],
+        files=[AiFile(id=FileId("report-id"), name="report.pdf")],
         artifacts=artifacts,
     )
 

engine/tests/agents/test_orchestrator_pdf_comment.py

@@ -23,8 +23,9 @@ from unittest.mock import AsyncMock, patch
 import pytest
 
 from stirling.agents import OrchestratorAgent
-from stirling.contracts import OrchestratorRequest
+from stirling.contracts import AiFile, OrchestratorRequest
 from stirling.contracts.pdf_edit import EditPlanResponse
+from stirling.models import FileId
 from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
 from stirling.services.runtime import AppRuntime
 
@@ -39,7 +40,7 @@ async def test_delegate_pdf_review_wires_prompt_to_tool_step(runtime: AppRuntime
     orchestrator = OrchestratorAgent(runtime)
     request = OrchestratorRequest(
         user_message="please add review comments flagging ambiguous dates",
-        file_names=["contract.pdf"],
+        files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
     )
     ctx = SimpleNamespace(deps=_FakeDeps(request=request))
 

engine/tests/agents/test_pdf_questions_orchestrate.py

@@ -12,12 +12,15 @@ import pytest
 
 from stirling.agents.pdf_questions import _MATH_SYNTH_SYSTEM_PROMPT, PdfQuestionAgent
 from stirling.contracts import (
+    AiFile,
+    EditPlanResponse,
     MathAuditorToolReportArtifact,
     OrchestratorRequest,
     PdfQuestionAnswerResponse,
     SupportedCapability,
 )
 from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
+from stirling.models import FileId
 from stirling.models.agent_tool_models import AgentToolId
 from stirling.services.runtime import AppRuntime
 
@@ -49,25 +52,23 @@ def _make_verdict() -> Verdict:
 
 
 @pytest.mark.anyio
-async def test_orchestrate_classifier_true_embeds_plan_in_answer(runtime: AppRuntime) -> None:
-    """First turn — classifier says math; the response is a PdfQuestionAnswerResponse
-    with the math-auditor plan attached as a nullable ``edit_plan`` field. The
-    answer is empty on this turn; the caller runs the embedded plan and resumes."""
+async def test_orchestrate_classifier_true_returns_math_audit_plan(runtime: AppRuntime) -> None:
+    """First turn — classifier says math; the response is an EditPlanResponse
+    (``outcome=PLAN``) with ``resume_with=PDF_QUESTION``. The caller runs the
+    plan and re-invokes the orchestrator with the verdict in artifacts."""
     agent = PdfQuestionAgent(runtime)
     request = OrchestratorRequest(
         user_message="ist die mathematik korrekt?",
-        file_names=["report.pdf"],
+        files=[AiFile(id=FileId("report-id"), name="report.pdf")],
     )
 
     with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
         response = await agent.orchestrate(request)
 
-    assert isinstance(response, PdfQuestionAnswerResponse)
-    assert response.answer == ""
-    assert response.edit_plan is not None
-    assert response.edit_plan.resume_with == SupportedCapability.PDF_QUESTION
-    assert len(response.edit_plan.steps) == 1
-    assert response.edit_plan.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
+    assert isinstance(response, EditPlanResponse)
+    assert response.resume_with == SupportedCapability.PDF_QUESTION
+    assert len(response.steps) == 1
+    assert response.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
 
 
 @pytest.mark.anyio
@@ -81,7 +82,7 @@ async def test_orchestrate_resume_synthesises_answer_without_calling_classifier(
     verdict = _make_verdict()
     request = OrchestratorRequest(
         user_message="ist die mathematik korrekt?",
-        file_names=["report.pdf"],
+        files=[AiFile(id=FileId("report-id"), name="report.pdf")],
         artifacts=[MathAuditorToolReportArtifact(report=verdict)],
     )
     canned_answer = "Die Summe stimmt nicht: angegeben $215,000, erwartet $215,500."

engine/tests/agents/test_pdf_review.py

@@ -20,9 +20,9 @@ from stirling.agents.pdf_review import (
     _LocalisedComment,
     _LocalisedVerdict,
 )
-from stirling.contracts import EditPlanResponse, OrchestratorRequest, SupportedCapability
+from stirling.contracts import AiFile, EditPlanResponse, OrchestratorRequest, SupportedCapability
 from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
-from stirling.models import ToolEndpoint
+from stirling.models import FileId, ToolEndpoint
 from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
 from stirling.services.runtime import AppRuntime
 
@@ -144,7 +144,10 @@ async def test_orchestrate_classifier_true_emits_math_audit_plan(runtime: AppRun
     """First turn — when the math-intent classifier says yes, emit a one-step plan
     calling the math auditor with resume_with=PDF_REVIEW."""
     agent = PdfReviewAgent(runtime)
-    request = OrchestratorRequest(user_message="vérifie les totaux", file_names=["report.pdf"])
+    request = OrchestratorRequest(
+        user_message="vérifie les totaux",
+        files=[AiFile(id=FileId("report-id"), name="report.pdf")],
+    )
 
     with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
         response = await agent.orchestrate(request)
@@ -161,7 +164,7 @@ async def test_orchestrate_classifier_false_routes_to_pdf_comment_agent(runtime:
     agent = PdfReviewAgent(runtime)
     request = OrchestratorRequest(
         user_message="review the invoices for ambiguous wording",
-        file_names=["contract.pdf"],
+        files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
     )
 
     with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=False)):
@@ -187,7 +190,7 @@ async def test_orchestrate_resume_uses_verdict_without_calling_classifier(
     verdict = _make_verdict([_discrepancy(page=0, stated="$100")])
     request = OrchestratorRequest(
         user_message="flag math errors",
-        file_names=["report.pdf"],
+        files=[AiFile(id=FileId("report-id"), name="report.pdf")],
         artifacts=[MathAuditorToolReportArtifact(report=verdict)],
     )
     canned = _LocalisedVerdict(

engine/tests/conftest.py

@@ -30,6 +30,7 @@ def build_app_settings() -> AppSettings:
         rag_chunk_size=512,
         rag_chunk_overlap=64,
         rag_default_top_k=5,
+        rag_max_searches=5,
         max_pages=200,
         max_characters=200_000,
         posthog_enabled=False,

engine/tests/pdf_comment/test_routes.py

@@ -42,6 +42,7 @@ class StubSettingsProvider:
             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,

engine/tests/test_pdf_edit_agent.py

@@ -7,6 +7,7 @@ import pytest
 from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
 from stirling.agents.pdf_edit import PdfEditPlanOutput
 from stirling.contracts import (
+    AiFile,
     EditCannotDoResponse,
     EditClarificationRequest,
     EditPlanResponse,
@@ -19,6 +20,7 @@ from stirling.contracts import (
     SupportedCapability,
     ToolOperationStep,
 )
+from stirling.models import FileId
 from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
 from stirling.services.runtime import AppRuntime
 
@@ -91,7 +93,7 @@ async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> Non
     response = await agent.handle(
         PdfEditRequest(
             user_message="Rotate the PDF clockwise and then compress it.",
-            file_names=["scan.pdf"],
+            files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
         )
     )
 
@@ -117,7 +119,7 @@ async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtim
 
     request = PdfEditRequest(
         user_message="Rotate the PDF clockwise and then compress it.",
-        file_names=["scan.pdf"],
+        files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
     )
     response = await agent.handle(request)
 
@@ -181,13 +183,18 @@ async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime
     response = await agent.handle(
         PdfEditRequest(
             user_message="Split after every page that says 'NEW PAGE'.",
-            file_names=["report.pdf"],
+            files=[AiFile(id=FileId("report-id"), name="report.pdf")],
         )
     )
 
     assert isinstance(response, NeedContentResponse)
     assert response.resume_with == SupportedCapability.PDF_EDIT
-    assert response.files == [NeedContentFileRequest(file_name="report.pdf", content_types=[PdfContentType.PAGE_TEXT])]
+    assert response.files == [
+        NeedContentFileRequest(
+            file=AiFile(id=FileId("report-id"), name="report.pdf"),
+            content_types=[PdfContentType.PAGE_TEXT],
+        )
+    ]
     assert response.max_pages == runtime.settings.max_pages
     assert response.max_characters == runtime.settings.max_characters
     assert parameter_selector.calls == []
@@ -269,7 +276,7 @@ async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: Ap
     await agent.handle(
         PdfEditRequest(
             user_message="Rotate clockwise.",
-            file_names=["report.pdf"],
+            files=[AiFile(id=FileId("report-id"), name="report.pdf")],
             page_text=page_text,
         )
     )

engine/tests/test_pdf_question_agent.py

@@ -1,64 +1,133 @@
 from __future__ import annotations
 
+from dataclasses import replace
+
 import pytest
 
 from stirling.agents import PdfQuestionAgent
 from stirling.contracts import (
+    AiFile,
     ExtractedFileText,
-    NeedContentResponse,
+    NeedIngestResponse,
+    PdfContentType,
     PdfQuestionAnswerResponse,
     PdfQuestionNotFoundResponse,
     PdfQuestionRequest,
+    PdfQuestionTerminalResponse,
     PdfTextSelection,
+    SupportedCapability,
 )
+from stirling.models import FileId
+from stirling.rag import Document, RagService, SqliteVecStore
 from stirling.services.runtime import AppRuntime
 
 
-class StubPdfQuestionAgent(PdfQuestionAgent):
-    def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
-        super().__init__(runtime)
-        self.response = response
+class StubEmbedder:
+    """Deterministic embeddings so RAG lookups work in tests without network."""
 
-    async def _run_answer_agent(
+    def __init__(self, dim: int = 8) -> None:
+        self._dim = dim
+
+    async def embed_query(self, text: str) -> list[float]:
+        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,
-        request: PdfQuestionRequest,
-    ) -> PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse:
-        return self.response
+        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: list[Document] = []
+        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
 
 
-def invoice_page() -> ExtractedFileText:
-    return ExtractedFileText(
-        file_name="invoice.pdf",
-        pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
+class StubPdfQuestionAgent(PdfQuestionAgent):
+    def __init__(self, runtime: AppRuntime, response: PdfQuestionTerminalResponse) -> None:
+        super().__init__(runtime)
+        self._response = response
+
+    async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
+        return self._response
+
+
+@pytest.fixture
+def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
+    """A runtime whose RAG service uses a stub embedder + ephemeral store."""
+    stub = RagService(
+        embedder=StubEmbedder(),  # type: ignore[arg-type]
+        store=SqliteVecStore.ephemeral(),
+        default_top_k=runtime.settings.rag_default_top_k,
     )
+    return replace(runtime, rag_service=stub)
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None:
-    agent = PdfQuestionAgent(runtime)
+async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag: AppRuntime) -> None:
+    agent = PdfQuestionAgent(runtime_with_stub_rag)
 
-    response = await agent.handle(
-        PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
-    )
+    missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf")
+    response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[missing_file]))
 
-    assert isinstance(response, NeedContentResponse)
+    assert isinstance(response, NeedIngestResponse)
+    assert response.resume_with == SupportedCapability.PDF_QUESTION
+    assert response.files_to_ingest == [missing_file]
+    assert PdfContentType.PAGE_TEXT in response.content_types
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None:
+async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
+    await runtime_with_stub_rag.rag_service.index_text(
+        collection=FileId("present-id"),
+        text="Invoice total: 120.00.",
+        source="present.pdf",
+    )
+    agent = PdfQuestionAgent(runtime_with_stub_rag)
+
+    present_file = AiFile(id=FileId("present-id"), name="present.pdf")
+    missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf")
+    response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[present_file, missing_file]))
+
+    assert isinstance(response, NeedIngestResponse)
+    assert response.files_to_ingest == [missing_file]
+
+
+@pytest.mark.anyio
+async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub_rag: AppRuntime) -> None:
+    await runtime_with_stub_rag.rag_service.index_text(
+        collection=FileId("invoice-id"),
+        text="Invoice total: 120.00.",
+        source="invoice.pdf",
+    )
     agent = StubPdfQuestionAgent(
-        runtime,
+        runtime_with_stub_rag,
         PdfQuestionAnswerResponse(
             answer="The invoice total is 120.00.",
-            evidence=[invoice_page()],
+            evidence=[
+                ExtractedFileText(
+                    file_name="invoice.pdf",
+                    pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
+                )
+            ],
         ),
     )
 
     response = await agent.handle(
         PdfQuestionRequest(
             question="What is the total?",
-            page_text=[invoice_page()],
-            file_names=["invoice.pdf"],
+            files=[AiFile(id=FileId("invoice-id"), name="invoice.pdf")],
         )
     )
 
@@ -67,19 +136,21 @@ async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -
 
 
 @pytest.mark.anyio
-async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None:
-    agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
+async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
+    await runtime_with_stub_rag.rag_service.index_text(
+        collection=FileId("shipping-id"),
+        text="This page contains only a shipping address.",
+        source="shipping.pdf",
+    )
+    agent = StubPdfQuestionAgent(
+        runtime_with_stub_rag,
+        PdfQuestionNotFoundResponse(reason="The answer is not present in the text."),
+    )
 
     response = await agent.handle(
         PdfQuestionRequest(
             question="What is the total?",
-            page_text=[
-                ExtractedFileText(
-                    file_name="invoice.pdf",
-                    pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")],
-                )
-            ],
-            file_names=["invoice.pdf"],
+            files=[AiFile(id=FileId("shipping-id"), name="shipping.pdf")],
         )
     )
 

engine/tests/test_rag.py

@@ -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()

engine/tests/test_rag_routes.py

@@ -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"))

engine/tests/test_stirling_api.py

@@ -88,7 +88,10 @@ def test_health_route() -> None:
 
 
 def test_orchestrator_route() -> None:
-    response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]})
+    response = client.post(
+        "/api/v1/orchestrator",
+        json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
+    )
 
     assert response.status_code == 200
     assert response.json()["outcome"] == "need_content"
@@ -106,8 +109,7 @@ def test_pdf_questions_route() -> None:
         "/api/v1/pdf/questions",
         json={
             "question": "what is this?",
-            "fileNames": ["test.pdf"],
-            "pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}],
+            "files": [{"id": "test-id", "name": "test.pdf"}],
         },
     )
 

engine/tests/test_stirling_contracts.py

@@ -3,6 +3,7 @@ from stirling.contracts import (
     AgentExecutionRequest,
     AgentSpec,
     AgentSpecStep,
+    AiFile,
     EditPlanResponse,
     ExecutionContext,
     ExtractedFileText,
@@ -12,13 +13,14 @@ from stirling.contracts import (
     PdfTextSelection,
     ToolOperationStep,
 )
+from stirling.models import FileId
 from stirling.models.tool_models import Angle, RotatePdfParams, ToolEndpoint
 
 
 def test_orchestrator_request_accepts_user_message() -> None:
     request = OrchestratorRequest(
         user_message="Rotate the PDF",
-        file_names=["test.pdf"],
+        files=[AiFile(id=FileId("test-id"), name="test.pdf")],
         artifacts=[
             ExtractedTextArtifact(
                 files=[
@@ -89,6 +91,7 @@ def test_app_settings_accepts_model_configuration() -> None:
         rag_chunk_size=512,
         rag_chunk_overlap=64,
         rag_default_top_k=5,
+        rag_max_searches=5,
         max_pages=200,
         max_characters=200_000,
         posthog_enabled=False,