mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 03:20:46 +02:00
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:
@@ -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)
|
||||
Reference in New Issue
Block a user