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
@@ -0,0 +1,140 @@
"""Tool capability that lets an agent read whole documents end-to-end.
Companion to :class:`stirling.documents.RagCapability`. Where ``RagCapability``
gives an agent targeted vector retrieval, this gives it map-style whole-document
reading: every page is read in parallel by fast-model workers, and the
question-relevant notes are returned for the agent to synthesise.
Use both capabilities together when the agent should pick its strategy:
``search_knowledge`` for specific lookups, ``read_full_document`` for
aggregations, comparisons, and summaries.
"""
from __future__ import annotations
import logging
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner
from stirling.contracts import AiFile
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
# Cap on per-run calls. One pass already reads every page of every attached
# document, so a second call is almost always the smart model second-guessing
# itself on a near-identical query (and doubles wall-clock time for a sizeable
# document). If a follow-up genuinely needs more, ``search_knowledge`` is the
# right escape hatch. Configurable per-construction in case a future caller
# can prove a real two-read use case; the default stays at 1.
DEFAULT_MAX_READS = 1
class WholeDocReaderCapability:
"""Bundles instructions and the ``read_full_document`` toolset for agent injection.
Lifecycle: a ``WholeDocReaderCapability`` instance is intended to live for
the duration of a single agent run.
The agent picks between this and :class:`RagCapability` per the tool
descriptions: targeted retrieval vs whole-document reading.
"""
def __init__(
self,
runtime: AppRuntime,
files: list[AiFile],
*,
reasoner: ChunkedReasoner | None = None,
max_reads: int = DEFAULT_MAX_READS,
) -> None:
self._runtime = runtime
self._files = files
self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime)
self._max_reads = max_reads
self._read_count = 0
toolset: FunctionToolset[None] = FunctionToolset()
toolset.add_function(
self._read_full_document,
name="read_full_document",
prepare=self._prepare_read_full_document,
)
self._toolset = toolset
@property
def instructions(self) -> str:
names = ", ".join(f.name for f in self._files) if self._files else "the attached documents"
return (
"You have a 'read_full_document' tool that reads every page of "
f"{names} in parallel and returns notes relevant to a query. "
"Use it when answering requires seeing the whole document end-to-end "
"(summaries, aggregations, comparisons across sections). One call "
"already reads everything; phrase the query to cover all the angles "
"you need in a single pass. For follow-ups or specific lookups use "
"'search_knowledge', which is cheaper and targeted."
)
@property
def toolset(self) -> AbstractToolset[None]:
return self._toolset
async def _prepare_read_full_document(
self,
ctx: RunContext[None],
tool_def: ToolDefinition,
) -> ToolDefinition | None:
"""Hide the tool from the agent's toolset once the per-run budget is spent.
Mirrors the search_knowledge prepare callback."""
if self._read_count >= self._max_reads:
return None
return tool_def
async def _read_full_document(self, query: str) -> str:
"""Read every page of the attached documents and return notes relevant to the query.
Use this when answering needs the whole document end-to-end - summaries,
aggregations like 'largest number' or 'shortest chapter', or comparisons
across sections. Slow and expensive (one fast-model call per slice per
document); prefer search_knowledge for targeted lookups.
Args:
query: A focused description of what to extract from the documents,
phrased so a worker reading just one slice can decide what's
relevant to the user's question.
Returns:
Per-document sections of structured notes (page numbers, summary,
relevant excerpts, extracted facts), already ordered by page.
"""
self._read_count += 1
if not self._files:
return "No documents attached to read."
sections: list[str] = []
for file in self._files:
pages = await self._runtime.documents.read_pages(file.id)
if not pages:
logger.info(
"[whole-doc-reader] no stored pages for %s (id=%s); skipping",
file.name,
file.id,
)
continue
notes = await self._reasoner.gather_notes(pages, query)
if not notes:
continue
sections.append(f"=== {file.name} ===\n{ChunkedReasoner.format_notes(notes)}")
if not sections:
return "Could not read any document content."
logger.info(
"[whole-doc-reader] read query=%r files=%d -> %d chars",
query,
len(self._files),
sum(len(s) for s in sections),
)
return "\n\n".join(sections)