Add ability for Stirling engine to reason across large documents (#6314)

# Description of Changes
Adds storage in the database for full document content alongside the RAG
content (and changes the service to `DocumentService` instead of
`RagService`). Then adds a generic capability that should be usable by
any agent (currently just used by the Question Agent) which allows the
agent to pull out the full contents of the doc, chunks it into various
sections that will fit in the context window, and then processes them in
parallel to create an intermediate result, and then processes the
intermediate result into a final answer. It will re-chunk as many times
as necessary to get the content small enough for the actual answer to be
analysed (I've tested on PDFs ~3500 pages long, which is well above the
context limit and requires maybe 3 rounds of compression to get an
answer).

The new full doc analysis stuff is heavier than the RAG lookup so both
remain. The agents should use RAG for targeted info and the chunked
reasoner for info that requires reading the full doc.
This commit is contained in:
James Brunton
2026-05-14 13:19:38 +00:00
committed by GitHub
parent 8abe734f0b
commit 672e81d286
55 changed files with 3327 additions and 578 deletions
+15 -14
View File
@@ -9,6 +9,7 @@ from stirling.contracts import (
AiFile,
ExtractedFileText,
NeedIngestResponse,
PageText,
PdfContentType,
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
@@ -17,8 +18,8 @@ from stirling.contracts import (
PdfTextSelection,
SupportedCapability,
)
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.rag import Document, RagService, SqliteVecStore
from stirling.services.runtime import AppRuntime
@@ -41,7 +42,7 @@ class StubEmbedder:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs: list[Document] = []
@@ -65,13 +66,13 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
@pytest.fixture
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
"""A runtime whose RAG service uses a stub embedder + ephemeral store."""
stub = RagService(
"""A runtime whose document service uses a stub embedder + ephemeral store."""
stub = DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=runtime.settings.rag_default_top_k,
)
return replace(runtime, rag_service=stub)
return replace(runtime, documents=stub)
@pytest.mark.anyio
@@ -89,9 +90,9 @@ async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag:
@pytest.mark.anyio
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.",
await runtime_with_stub_rag.documents.ingest(
FileId("present-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="present.pdf",
)
agent = PdfQuestionAgent(runtime_with_stub_rag)
@@ -106,9 +107,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) ->
@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.",
await runtime_with_stub_rag.documents.ingest(
FileId("invoice-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="invoice.pdf",
)
agent = StubPdfQuestionAgent(
@@ -137,9 +138,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub
@pytest.mark.anyio
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.",
await runtime_with_stub_rag.documents.ingest(
FileId("shipping-id"),
[PageText(page_number=1, text="This page contains only a shipping address.")],
source="shipping.pdf",
)
agent = StubPdfQuestionAgent(