mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +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:
James Brunton
@@ -6,6 +6,7 @@ from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
|
||||
from stirling.agents.shared import ChunkedReasoner, WholeDocReaderCapability
|
||||
from stirling.contracts import (
|
||||
AiFile,
|
||||
EditPlanResponse,
|
||||
@@ -24,44 +25,46 @@ from stirling.contracts import (
|
||||
format_conversation_history,
|
||||
format_file_names,
|
||||
)
|
||||
from stirling.documents import RagCapability
|
||||
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
|
||||
from stirling.rag import RagCapability
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
PDF_QUESTION_SYSTEM_PROMPT = (
|
||||
"You answer questions about PDF documents by retrieving relevant content with the "
|
||||
"search_knowledge tool. Use it before answering. Do not guess or use outside knowledge.\n"
|
||||
"You answer questions about PDF documents using two retrieval tools:\n"
|
||||
"\n"
|
||||
"The search_knowledge tool has a finite call budget per run. When it is no longer "
|
||||
"available, answer from what you have already retrieved.\n"
|
||||
"1. search_knowledge(query) - returns the passages most semantically similar "
|
||||
"to the query. Use it for targeted lookups: a specific fact, a named section, "
|
||||
"a particular passage. Typically one or two calls is enough.\n"
|
||||
"\n"
|
||||
"2. read_full_document(query) - reads every page of the attached documents in "
|
||||
"parallel and returns notes relevant to the query. Use it when answering "
|
||||
"requires seeing the whole document end-to-end: summaries, aggregations "
|
||||
"(largest, shortest, count), comparisons across sections. It is more "
|
||||
"expensive than search_knowledge, so prefer search_knowledge when one or two "
|
||||
"passages would suffice.\n"
|
||||
"\n"
|
||||
"Pick the right tool, call it, then answer from what you got back. Do not "
|
||||
"guess or use outside knowledge.\n"
|
||||
"\n"
|
||||
"Guidelines:\n"
|
||||
"- Make targeted search_knowledge calls. Typically one or two is enough.\n"
|
||||
"- Answer from the retrieved text. If the retrieved content doesn't support a confident "
|
||||
"answer, return not_found.\n"
|
||||
"- For questions that would require reading the entire document end-to-end (e.g. "
|
||||
"'what's the shortest chapter', 'how many X are there'), return not_found.\n"
|
||||
"- Include a short list of evidence snippets (with page numbers where available) drawn "
|
||||
"from what search_knowledge returned.\n"
|
||||
"- If the retrieved content does not support a confident answer, return not_found.\n"
|
||||
"- Include a short list of evidence snippets (with page numbers where available) "
|
||||
"drawn from what the tools returned.\n"
|
||||
"\n"
|
||||
"Writing the not_found reason:\n"
|
||||
"- The reason is shown directly to the end user, so write it in plain, friendly "
|
||||
"language. One or two short sentences.\n"
|
||||
"- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
|
||||
"'search_knowledge', or other implementation details.\n"
|
||||
"- Be honest about the actual limitation. For questions that require full-document "
|
||||
"analysis (shortest chapter, word counts, etc.), explain that the document is too "
|
||||
"long to analyse end-to-end: you can only look up specific passages, and that's "
|
||||
"not enough to compare every part of the document against every other.\n"
|
||||
"'search_knowledge', 'read_full_document', or other implementation details.\n"
|
||||
"- For questions where the answer just isn't in the document, say so directly: "
|
||||
"'I couldn't find that information in the document.'\n"
|
||||
"- Do not make it sound like you're choosing not to answer. Be clear that it's "
|
||||
"a genuine constraint."
|
||||
"- Do not make it sound like you're choosing not to answer."
|
||||
)
|
||||
|
||||
|
||||
_MATH_SYNTH_SYSTEM_PROMPT = (
|
||||
"You are given a math-audit Verdict (structured JSON) and the user's "
|
||||
"original question. Answer the question in plain prose using only "
|
||||
@@ -83,6 +86,9 @@ class PdfQuestionAgent:
|
||||
model_settings=runtime.fast_model_settings,
|
||||
)
|
||||
self._math_intent_classifier = MathIntentClassifier(runtime)
|
||||
# Shared across whole-doc-reader instances so the worker agent and
|
||||
# semaphore are constructed once and reused per request.
|
||||
self._chunked_reasoner = ChunkedReasoner(runtime)
|
||||
|
||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||
logger.info(
|
||||
@@ -95,7 +101,7 @@ class PdfQuestionAgent:
|
||||
logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
|
||||
return NeedIngestResponse(
|
||||
resume_with=SupportedCapability.PDF_QUESTION,
|
||||
reason="Some files have not been ingested into RAG yet.",
|
||||
reason="Some files have not been ingested yet.",
|
||||
files_to_ingest=missing,
|
||||
content_types=[PdfContentType.PAGE_TEXT],
|
||||
)
|
||||
@@ -145,23 +151,37 @@ class PdfQuestionAgent:
|
||||
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
|
||||
missing: list[AiFile] = []
|
||||
for file in files:
|
||||
if not await self.runtime.rag_service.has_collection(file.id):
|
||||
if not await self.runtime.documents.has_collection(file.id):
|
||||
missing.append(file)
|
||||
return missing
|
||||
|
||||
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
|
||||
"""Drive a single smart-model agent with both retrieval tools.
|
||||
|
||||
The agent picks ``search_knowledge`` for targeted lookups and
|
||||
``read_full_document`` for whole-document questions. Removing the
|
||||
upstream classifier keeps that judgement in the same call that writes
|
||||
the answer, and lets the agent mix tools when the question warrants it.
|
||||
"""
|
||||
rag = RagCapability(
|
||||
rag_service=self.runtime.rag_service,
|
||||
documents=self.runtime.documents,
|
||||
collections=[file.id for file in request.files],
|
||||
top_k=self.runtime.settings.rag_default_top_k,
|
||||
max_searches=self.runtime.settings.rag_max_searches,
|
||||
)
|
||||
whole_doc = WholeDocReaderCapability(
|
||||
runtime=self.runtime,
|
||||
files=request.files,
|
||||
reasoner=self._chunked_reasoner,
|
||||
)
|
||||
agent = Agent(
|
||||
model=self.runtime.smart_model,
|
||||
output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
|
||||
system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
|
||||
instructions=rag.instructions,
|
||||
toolsets=[rag.toolset],
|
||||
# pydantic-ai accepts a list of (string-or-callable) instruction sources;
|
||||
# it resolves each at run time and concatenates them for the model.
|
||||
instructions=[rag.instructions, whole_doc.instructions],
|
||||
toolsets=[rag.toolset, whole_doc.toolset],
|
||||
model_settings=self.runtime.smart_model_settings,
|
||||
)
|
||||
prompt = self._build_prompt(request)
|
||||
@@ -184,5 +204,5 @@ class PdfQuestionAgent:
|
||||
f"Conversation history:\n{history}\n"
|
||||
f"Files: {format_file_names(request.files)}\n"
|
||||
f"Question: {request.question}\n"
|
||||
"Use search_knowledge to retrieve the relevant content, then answer."
|
||||
"Pick the right retrieval tool for this question, then answer from what it returns."
|
||||
)
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Reasoning utilities shared across agents."""
|
||||
|
||||
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
|
||||
from stirling.agents.shared.whole_doc_reader import WholeDocReaderCapability
|
||||
|
||||
__all__ = [
|
||||
"ChunkNotes",
|
||||
"ChunkedReasoner",
|
||||
"WholeDocReaderCapability",
|
||||
]
|
||||
@@ -0,0 +1,616 @@
|
||||
"""Chunked reasoning over long documents.
|
||||
|
||||
A reusable primitive for any agent that needs to answer a question that
|
||||
requires reading a whole document end-to-end. The document is split into
|
||||
character-budgeted chunks; each chunk is read by a parallel worker that
|
||||
extracts question-relevant notes; if the gathered notes overflow the
|
||||
synthesis context budget, the resulting notes are regrouped into fresh
|
||||
chunks and run through the same extractor again, until they fit.
|
||||
|
||||
Pages are tracked by the wrapper, never asked of the model: keeps the model
|
||||
output schema small and the page list authoritative.
|
||||
|
||||
Used wherever pure RAG retrieval is the wrong tool: aggregations ("largest
|
||||
number"), comparisons ("shortest chapter"), and full summaries.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.output import NativeOutput
|
||||
|
||||
from stirling.contracts import (
|
||||
WholeDocCompressionRound,
|
||||
WholeDocReadDone,
|
||||
WholeDocReadStarted,
|
||||
WholeDocSliceDone,
|
||||
)
|
||||
from stirling.contracts.documents import Page
|
||||
from stirling.models import ApiModel
|
||||
from stirling.services import AppRuntime, emit_progress
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChunkNotes(ApiModel):
|
||||
"""Public-facing notes for a span of pages.
|
||||
|
||||
Returned to callers of :meth:`ChunkedReasoner.gather_notes` and to the
|
||||
inside of :meth:`ChunkedReasoner.reason`. The wrapper builds these from
|
||||
the model's :class:`_ExtractedNotes` output and a deterministic page list.
|
||||
"""
|
||||
|
||||
pages: list[int] = Field(description="Page numbers covered by these notes (1-indexed).")
|
||||
summary: str = Field(description="One- to three-sentence summary of the covered range.")
|
||||
relevant_excerpts: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Short verbatim quotes from the source content that bear on the user's question.",
|
||||
)
|
||||
facts: list[str] = Field(
|
||||
default_factory=list,
|
||||
description=(
|
||||
"Concrete facts (numbers, names, dates, claims) the synthesiser may need. "
|
||||
"Includes candidate values for aggregation questions."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _ExtractedNotes(BaseModel):
|
||||
"""Model output for one extractor call.
|
||||
|
||||
No ``pages`` field: page numbers are mechanical aggregation the wrapper
|
||||
computes deterministically. Keeping them out of the schema saves output
|
||||
tokens for the bulkier excerpts/facts payload and prevents the model
|
||||
from misreporting page coverage.
|
||||
"""
|
||||
|
||||
summary: str = Field(description="One- to three-sentence summary of the supplied content.")
|
||||
relevant_excerpts: list[str] = Field(
|
||||
default_factory=list,
|
||||
description=(
|
||||
"Short verbatim quotes drawn from the supplied content that bear on the question. "
|
||||
"Deduplicate; drop ones that don't bear on the question."
|
||||
),
|
||||
)
|
||||
facts: list[str] = Field(
|
||||
default_factory=list,
|
||||
description=(
|
||||
"Distinct, deduplicated facts (numbers, names, dates, claims) needed to answer "
|
||||
"the question. For aggregation questions retain ALL candidate values across the "
|
||||
"supplied content so a later round can still pick the global winner."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Chunk:
|
||||
"""A unit of work for the extractor: content + the pages it covers + a fallback.
|
||||
|
||||
``content`` is the formatted text fed to the model: raw page text with
|
||||
``[Page N]`` markers in the first round, formatted prior-pass notes with
|
||||
``[Notes from pages A-B]`` markers in subsequent rounds. ``pages`` is
|
||||
attached to the resulting :class:`ChunkNotes` deterministically.
|
||||
|
||||
``fallback`` is the list of notes to keep if the extractor call fails. For
|
||||
raw page chunks it's empty (a failed slice has no pre-extracted notes to
|
||||
preserve). For chunks built from existing notes it's the input notes
|
||||
themselves, so a failure doesn't lose page coverage.
|
||||
"""
|
||||
|
||||
content: str
|
||||
pages: list[int]
|
||||
fallback: list[ChunkNotes]
|
||||
label: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _RoundResult:
|
||||
"""Outcome of one extraction round.
|
||||
|
||||
``successes`` lets the loop detect rounds that made no forward progress
|
||||
(every chunk failed) and bail rather than spinning. ``slowest`` is the
|
||||
chunk with the longest successful extractor call this round, used for
|
||||
diagnostic log lines on the first round.
|
||||
"""
|
||||
|
||||
notes: list[ChunkNotes]
|
||||
successes: int
|
||||
slowest: tuple[str, float] | None
|
||||
|
||||
|
||||
def _page_range_label(pages: list[Page]) -> str:
|
||||
if not pages:
|
||||
return "pages=?"
|
||||
elif len(pages) == 1:
|
||||
return f"pages={pages[0].page_number}"
|
||||
else:
|
||||
return f"pages={pages[0].page_number}-{pages[-1].page_number}"
|
||||
|
||||
|
||||
def _note_range_label(notes: list[ChunkNotes]) -> str:
|
||||
"""Render a "pages=A-B" label for a group of already-extracted notes."""
|
||||
page_numbers = sorted({p for note in notes for p in note.pages})
|
||||
if not page_numbers:
|
||||
return "pages=?"
|
||||
if len(page_numbers) == 1:
|
||||
return f"pages={page_numbers[0]}"
|
||||
return f"pages={page_numbers[0]}-{page_numbers[-1]}"
|
||||
|
||||
|
||||
_EXTRACTOR_SYSTEM_PROMPT = (
|
||||
"You are reading content from a document - either raw page text or "
|
||||
"condensed notes from an earlier extraction pass - and your job is to "
|
||||
"produce a tight set of notes that captures everything relevant to the "
|
||||
"user's question. The same job runs many times in parallel across the "
|
||||
"document and may run again to consolidate notes into smaller batches, "
|
||||
"so be thorough: anything you skip cannot be recovered later.\n"
|
||||
"\n"
|
||||
"Output:\n"
|
||||
"- summary: 1-3 sentences covering the supplied content.\n"
|
||||
"- relevant_excerpts: short verbatim quotes from the supplied content "
|
||||
"that bear on the question. Deduplicate; drop quotes that don't help.\n"
|
||||
"- facts: concrete facts (numbers, names, dates, claims). Deduplicate; "
|
||||
"drop irrelevant ones. For aggregation questions (largest, smallest, "
|
||||
"count, total) retain ALL candidate values across the content so a "
|
||||
"later step can still pick the global winner.\n"
|
||||
"\n"
|
||||
"Stay grounded in the supplied content. Do not infer or fabricate "
|
||||
"anything that isn't already present. If nothing in the content is "
|
||||
"relevant to the question, return empty excerpts and facts and a short "
|
||||
"neutral summary."
|
||||
)
|
||||
|
||||
|
||||
class ChunkedReasoner:
|
||||
"""Run a question against a long document by chunking, mapping, and looping.
|
||||
|
||||
Two consumption styles:
|
||||
|
||||
* Tools that already have a synthesising LLM call upstream call
|
||||
:meth:`gather_notes` to get the structured notes and format them
|
||||
themselves with :meth:`format_notes`.
|
||||
* Callers that just want an answer call :meth:`reason`, which runs
|
||||
:meth:`gather_notes` and then a single synthesis call governed by the
|
||||
caller's ``answer_prompt`` and ``answer_type``.
|
||||
|
||||
Lifetime:
|
||||
Construct once per agent that uses it. The extractor agent is built
|
||||
at construction time and reused; the synthesis agent in :meth:`reason`
|
||||
is built per call because its output type is generic.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
runtime: AppRuntime,
|
||||
*,
|
||||
chars_per_slice: int | None = None,
|
||||
concurrency: int | None = None,
|
||||
worker_timeout_seconds: float | None = None,
|
||||
notes_char_budget: int | None = None,
|
||||
) -> None:
|
||||
chars = chars_per_slice if chars_per_slice is not None else runtime.settings.chunked_reasoner_chars_per_slice
|
||||
conc = concurrency if concurrency is not None else runtime.settings.chunked_reasoner_concurrency
|
||||
timeout = (
|
||||
worker_timeout_seconds
|
||||
if worker_timeout_seconds is not None
|
||||
else runtime.settings.chunked_reasoner_worker_timeout_seconds
|
||||
)
|
||||
budget = (
|
||||
notes_char_budget if notes_char_budget is not None else runtime.settings.chunked_reasoner_notes_char_budget
|
||||
)
|
||||
if chars <= 0:
|
||||
raise ValueError("chars_per_slice must be positive")
|
||||
if conc <= 0:
|
||||
raise ValueError("concurrency must be positive")
|
||||
if timeout <= 0:
|
||||
raise ValueError("worker_timeout_seconds must be positive")
|
||||
if budget <= 0:
|
||||
raise ValueError("notes_char_budget must be positive")
|
||||
self._runtime = runtime
|
||||
self._chars_per_slice = chars
|
||||
self._worker_timeout_seconds = timeout
|
||||
self._notes_char_budget = budget
|
||||
self._semaphore = asyncio.Semaphore(conc)
|
||||
self._extractor: Agent[None, _ExtractedNotes] = Agent(
|
||||
model=runtime.fast_model,
|
||||
output_type=NativeOutput(_ExtractedNotes),
|
||||
system_prompt=_EXTRACTOR_SYSTEM_PROMPT,
|
||||
model_settings=runtime.fast_model_settings,
|
||||
)
|
||||
|
||||
async def gather_notes(self, pages: list[Page], question: str) -> list[ChunkNotes]:
|
||||
"""Return notes covering every page that fit the synthesis budget.
|
||||
|
||||
Worker failures are tolerated: surviving notes are returned. Returns
|
||||
an empty list only when every first-round chunk raises, which the
|
||||
caller can treat as a hard failure.
|
||||
|
||||
Progress events fire as each first-round chunk finishes (in completion
|
||||
order, not chunk order) carrying a monotonic ``completed`` counter so
|
||||
consumers can render "Read X of Y" with X advancing by exactly one
|
||||
per event. Subsequent compression rounds emit a single round-start
|
||||
event each.
|
||||
"""
|
||||
if not pages:
|
||||
raise ValueError("ChunkedReasoner.gather_notes requires at least one page")
|
||||
|
||||
chunks = [self._chunk_from_pages(slice_pages) for slice_pages in self._slice_pages(pages)]
|
||||
slice_total = len(chunks)
|
||||
logger.info(
|
||||
"[chunked-reasoner] question=%r pages=%d slices=%d",
|
||||
question,
|
||||
len(pages),
|
||||
slice_total,
|
||||
)
|
||||
await emit_progress(WholeDocReadStarted(question=question, pages=len(pages), slices=slice_total))
|
||||
|
||||
gather_start = time.perf_counter()
|
||||
notes = await self._run_chunks(chunks, question)
|
||||
|
||||
await emit_progress(
|
||||
WholeDocReadDone(
|
||||
completed=len(notes),
|
||||
slices=slice_total,
|
||||
duration_seconds=round(time.perf_counter() - gather_start, 2),
|
||||
)
|
||||
)
|
||||
return notes
|
||||
|
||||
async def _run_chunks(self, chunks: list[_Chunk], question: str) -> list[ChunkNotes]:
|
||||
"""Run chunks through the extractor, regrouping and looping until under budget.
|
||||
|
||||
The first round emits per-chunk progress events for streaming UIs;
|
||||
later rounds emit a single round-start event. Each round may produce
|
||||
fewer notes than chunks (every chunk maps to at most one consolidated
|
||||
note); when the rendered notes still exceed the budget, the survivors
|
||||
are regrouped into fresh chunks and the loop runs again.
|
||||
"""
|
||||
round_number = 0
|
||||
while True:
|
||||
chunks_in = len(chunks)
|
||||
result = await self._extract_chunks(chunks, question, round_number)
|
||||
|
||||
if result.slowest is not None:
|
||||
slow_label, slow_duration = result.slowest
|
||||
logger.info(
|
||||
"[chunked-reasoner] round %d: %d/%d chunks succeeded; slowest %s (%.1fs)",
|
||||
round_number,
|
||||
result.successes,
|
||||
chunks_in,
|
||||
slow_label,
|
||||
slow_duration,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"[chunked-reasoner] round %d: 0/%d chunks succeeded",
|
||||
round_number,
|
||||
chunks_in,
|
||||
)
|
||||
|
||||
rendered_size = self._rendered_notes_size(result.notes)
|
||||
if rendered_size <= self._notes_char_budget or len(result.notes) <= 1:
|
||||
if round_number > 0:
|
||||
logger.info(
|
||||
"[chunked-reasoner] compression done after %d round(s): %d notes, %d chars",
|
||||
round_number,
|
||||
len(result.notes),
|
||||
rendered_size,
|
||||
)
|
||||
return result.notes
|
||||
|
||||
if result.successes == 0:
|
||||
# No forward progress this round; further rounds would
|
||||
# reproduce the same shape. Return what we have.
|
||||
logger.warning(
|
||||
"[chunked-reasoner] round %d produced no successful extractions; bailing with %d notes",
|
||||
round_number,
|
||||
len(result.notes),
|
||||
)
|
||||
return result.notes
|
||||
|
||||
round_number += 1
|
||||
groups = self._group_notes_for_compression(result.notes)
|
||||
chunks = [self._chunk_from_notes(group) for group in groups]
|
||||
logger.info(
|
||||
"[chunked-reasoner] compression round %d: %d notes (%d chars) -> %d groups",
|
||||
round_number,
|
||||
len(result.notes),
|
||||
rendered_size,
|
||||
len(groups),
|
||||
)
|
||||
await emit_progress(
|
||||
WholeDocCompressionRound(
|
||||
round_number=round_number,
|
||||
notes_in=len(result.notes),
|
||||
groups=len(groups),
|
||||
)
|
||||
)
|
||||
|
||||
async def _extract_chunks(
|
||||
self,
|
||||
chunks: list[_Chunk],
|
||||
question: str,
|
||||
round_number: int,
|
||||
) -> _RoundResult:
|
||||
"""Run all chunks through the extractor in parallel; collect surviving notes.
|
||||
|
||||
Failures fall back to ``chunk.fallback`` (empty in the first round, so
|
||||
failures drop; populated in compression rounds, so failures preserve
|
||||
their input notes). The first round emits a
|
||||
:class:`WholeDocSliceDone` per successful completion in completion
|
||||
order, with a monotonic ``completed`` counter.
|
||||
|
||||
Returned notes are sorted by first page so downstream grouping packs
|
||||
document-adjacent content together regardless of which task happened
|
||||
to finish first.
|
||||
"""
|
||||
total = len(chunks)
|
||||
pending: dict[asyncio.Task[tuple[ChunkNotes, float]], _Chunk] = {
|
||||
asyncio.create_task(self._extract_chunk(chunk, question)): chunk for chunk in chunks
|
||||
}
|
||||
|
||||
notes: list[ChunkNotes] = []
|
||||
successes = 0
|
||||
slowest: tuple[str, float] | None = None
|
||||
completed = 0
|
||||
|
||||
try:
|
||||
while pending:
|
||||
done, _ = await asyncio.wait(pending.keys(), return_when=asyncio.FIRST_COMPLETED)
|
||||
for task in done:
|
||||
chunk = pending.pop(task)
|
||||
exc = task.exception()
|
||||
if exc is not None:
|
||||
if chunk.fallback:
|
||||
logger.warning(
|
||||
"[chunked-reasoner] chunk %s failed: %s; preserving %d input note(s)",
|
||||
chunk.label,
|
||||
exc,
|
||||
len(chunk.fallback),
|
||||
)
|
||||
notes.extend(chunk.fallback)
|
||||
else:
|
||||
logger.warning("[chunked-reasoner] chunk %s failed: %s", chunk.label, exc)
|
||||
continue
|
||||
extracted, duration = task.result()
|
||||
notes.append(extracted)
|
||||
successes += 1
|
||||
completed += 1
|
||||
if slowest is None or duration > slowest[1]:
|
||||
slowest = (chunk.label, duration)
|
||||
if round_number == 0:
|
||||
await emit_progress(
|
||||
WholeDocSliceDone(
|
||||
completed=completed,
|
||||
total=total,
|
||||
pages=chunk.label,
|
||||
duration_ms=int(duration * 1000),
|
||||
excerpts=len(extracted.relevant_excerpts),
|
||||
facts=len(extracted.facts),
|
||||
)
|
||||
)
|
||||
finally:
|
||||
# On cancellation (typically a frontend disconnect propagating up
|
||||
# through the streaming orchestrator) the per-chunk model calls
|
||||
# would otherwise keep running to completion, billing tokens whose
|
||||
# results nobody is reading. Cancel and drain so the upstream
|
||||
# cancellation is the cancellation that matters.
|
||||
if pending:
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
await asyncio.gather(*pending.keys(), return_exceptions=True)
|
||||
|
||||
notes.sort(key=lambda n: n.pages[0] if n.pages else 0)
|
||||
return _RoundResult(notes=notes, successes=successes, slowest=slowest)
|
||||
|
||||
async def _extract_chunk(self, chunk: _Chunk, question: str) -> tuple[ChunkNotes, float]:
|
||||
"""Run the extractor on one chunk and attach the chunk's pages to the output."""
|
||||
try:
|
||||
extracted, duration = await self._run_extractor(chunk.content, question, chunk.label)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"[chunked-reasoner] chunk %s timed out (limit %.1fs)",
|
||||
chunk.label,
|
||||
self._worker_timeout_seconds,
|
||||
)
|
||||
raise
|
||||
logger.debug(
|
||||
"[chunked-reasoner] chunk %s: %d excerpt(s), %d fact(s) in %dms",
|
||||
chunk.label,
|
||||
len(extracted.relevant_excerpts),
|
||||
len(extracted.facts),
|
||||
int(duration * 1000),
|
||||
)
|
||||
return self._build_chunk_notes(extracted, chunk.pages), duration
|
||||
|
||||
async def _run_extractor(
|
||||
self,
|
||||
content: str,
|
||||
question: str,
|
||||
page_label: str,
|
||||
) -> tuple[_ExtractedNotes, float]:
|
||||
"""Inner primitive: run the extractor agent under semaphore + timeout."""
|
||||
prompt = self._build_extraction_prompt(content, question)
|
||||
async with self._semaphore:
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
result = await asyncio.wait_for(self._extractor.run(prompt), timeout=self._worker_timeout_seconds)
|
||||
except TimeoutError:
|
||||
duration = time.perf_counter() - start
|
||||
logger.debug(
|
||||
"[chunked-reasoner] extractor %s timed out after %dms",
|
||||
page_label,
|
||||
int(duration * 1000),
|
||||
)
|
||||
raise
|
||||
duration = time.perf_counter() - start
|
||||
return result.output, duration
|
||||
|
||||
def _chunk_from_pages(self, pages: list[Page]) -> _Chunk:
|
||||
"""Build a first-round chunk from a slice of raw pages."""
|
||||
return _Chunk(
|
||||
content="\n\n".join(f"[Page {p.page_number}]\n{p.text}" for p in pages),
|
||||
pages=[p.page_number for p in pages],
|
||||
fallback=[],
|
||||
label=_page_range_label(pages),
|
||||
)
|
||||
|
||||
def _chunk_from_notes(self, group: list[ChunkNotes]) -> _Chunk:
|
||||
"""Build a compression-round chunk from a group of prior-pass notes.
|
||||
|
||||
``fallback`` is the input group itself: if the extractor call fails,
|
||||
the originals stay in the working set so page coverage isn't lost.
|
||||
"""
|
||||
return _Chunk(
|
||||
content=self.format_notes(group),
|
||||
pages=sorted({p for note in group for p in note.pages}),
|
||||
fallback=group,
|
||||
label=_note_range_label(group),
|
||||
)
|
||||
|
||||
def _group_notes_for_compression(self, notes: list[ChunkNotes]) -> list[list[ChunkNotes]]:
|
||||
"""Pack consecutive notes into groups whose rendered size fits ``chars_per_slice``.
|
||||
|
||||
Each group becomes one compression-round chunk. Sized to match the
|
||||
first-round slice budget so the extractor sees roughly the same input
|
||||
footprint regardless of which round is running. Single notes that
|
||||
exceed the budget on their own become their own group.
|
||||
"""
|
||||
groups: list[list[ChunkNotes]] = []
|
||||
current: list[ChunkNotes] = []
|
||||
current_chars = 0
|
||||
for note in notes:
|
||||
note_chars = self._rendered_notes_size([note])
|
||||
if current and current_chars + note_chars > self._chars_per_slice:
|
||||
groups.append(current)
|
||||
current = []
|
||||
current_chars = 0
|
||||
current.append(note)
|
||||
current_chars += note_chars
|
||||
if current:
|
||||
groups.append(current)
|
||||
return groups
|
||||
|
||||
@staticmethod
|
||||
def _build_chunk_notes(extracted: _ExtractedNotes, pages: list[int]) -> ChunkNotes:
|
||||
"""Build a public ChunkNotes from the model's output and the wrapper's pages."""
|
||||
return ChunkNotes(
|
||||
pages=pages,
|
||||
summary=extracted.summary,
|
||||
relevant_excerpts=extracted.relevant_excerpts,
|
||||
facts=extracted.facts,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_extraction_prompt(content: str, question: str) -> str:
|
||||
"""Single prompt shape used for every round.
|
||||
|
||||
The system prompt explains the role; the user prompt just hands over
|
||||
the question and the content. Whether ``content`` is raw page text
|
||||
with ``[Page N]`` markers or formatted notes with ``[Notes from
|
||||
pages A-B]`` markers, the same instructions apply.
|
||||
"""
|
||||
return f"User question:\n{question}\n\nContent:\n{content}"
|
||||
|
||||
@staticmethod
|
||||
def _rendered_notes_size(notes: list[ChunkNotes]) -> int:
|
||||
"""Length in characters of what :meth:`format_notes` would produce."""
|
||||
return len(ChunkedReasoner.format_notes(notes))
|
||||
|
||||
async def reason[T: BaseModel](
|
||||
self,
|
||||
*,
|
||||
pages: list[Page],
|
||||
question: str,
|
||||
answer_prompt: str,
|
||||
answer_type: type[T],
|
||||
) -> T:
|
||||
"""Map over pages, then synthesise a structured answer of type ``T``.
|
||||
|
||||
Args:
|
||||
pages: Document pages in order.
|
||||
question: The user's question, passed to both workers and the
|
||||
synthesiser. Workers use it to decide what's relevant.
|
||||
answer_prompt: System prompt for the synthesis stage. Should
|
||||
instruct the model to answer ``question`` from the notes
|
||||
supplied. Owned by the caller because the answer's tone,
|
||||
format, and grounding rules are domain-specific.
|
||||
answer_type: Pydantic model describing the structured answer.
|
||||
|
||||
Returns:
|
||||
An instance of ``answer_type`` produced by the synthesis stage.
|
||||
"""
|
||||
notes = await self.gather_notes(pages, question)
|
||||
if not notes:
|
||||
raise RuntimeError("All chunked-reasoning workers failed; no notes to synthesise from")
|
||||
return await self._synthesise(question, notes, answer_prompt, answer_type)
|
||||
|
||||
@staticmethod
|
||||
def format_notes(notes: list[ChunkNotes]) -> str:
|
||||
"""Render notes as readable text for inclusion in another agent's tool result.
|
||||
|
||||
Order is preserved. Page numbers, summary, excerpts and facts are all
|
||||
emitted; empty sections are omitted.
|
||||
"""
|
||||
sections: list[str] = []
|
||||
for n in notes:
|
||||
page_label = (
|
||||
f"pages {n.pages[0]}-{n.pages[-1]}"
|
||||
if len(n.pages) > 1
|
||||
else f"page {n.pages[0]}"
|
||||
if n.pages
|
||||
else "unknown pages"
|
||||
)
|
||||
block = [f"[Notes from {page_label}]", f"Summary: {n.summary}"]
|
||||
if n.relevant_excerpts:
|
||||
block.append("Relevant excerpts:")
|
||||
block.extend(f"- {e}" for e in n.relevant_excerpts)
|
||||
if n.facts:
|
||||
block.append("Facts:")
|
||||
block.extend(f"- {f}" for f in n.facts)
|
||||
sections.append("\n".join(block))
|
||||
return "\n\n".join(sections)
|
||||
|
||||
def _slice_pages(self, pages: list[Page]) -> list[list[Page]]:
|
||||
"""Group consecutive pages into character-budgeted slices.
|
||||
|
||||
Page boundaries are preserved: a single page is never split across
|
||||
slices. If one page exceeds the budget on its own, it becomes its
|
||||
own slice.
|
||||
"""
|
||||
slices: list[list[Page]] = []
|
||||
current: list[Page] = []
|
||||
current_chars = 0
|
||||
for page in pages:
|
||||
if current and current_chars + page.char_count > self._chars_per_slice:
|
||||
slices.append(current)
|
||||
current = []
|
||||
current_chars = 0
|
||||
current.append(page)
|
||||
current_chars += page.char_count
|
||||
if current:
|
||||
slices.append(current)
|
||||
return slices
|
||||
|
||||
async def _synthesise[T: BaseModel](
|
||||
self,
|
||||
question: str,
|
||||
notes: list[ChunkNotes],
|
||||
answer_prompt: str,
|
||||
answer_type: type[T],
|
||||
) -> T:
|
||||
agent: Agent[None, T] = Agent(
|
||||
model=self._runtime.smart_model,
|
||||
output_type=NativeOutput(answer_type),
|
||||
system_prompt=answer_prompt,
|
||||
model_settings=self._runtime.smart_model_settings,
|
||||
)
|
||||
prompt = f"User question:\n{question}\n\nNotes from across the document:\n\n{self.format_notes(notes)}"
|
||||
result = await agent.run(prompt)
|
||||
return result.output
|
||||
@@ -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)
|
||||
@@ -19,13 +19,13 @@ from stirling.agents.pdf_comment import PdfCommentAgent
|
||||
from stirling.api.middleware import UserIdMiddleware
|
||||
from stirling.api.routes import (
|
||||
agent_draft_router,
|
||||
document_router,
|
||||
execution_router,
|
||||
ledger_router,
|
||||
orchestrator_router,
|
||||
pdf_comments_router,
|
||||
pdf_edit_router,
|
||||
pdf_question_router,
|
||||
rag_router,
|
||||
)
|
||||
from stirling.config import AppSettings, load_settings
|
||||
from stirling.contracts import HealthResponse
|
||||
@@ -57,7 +57,7 @@ async def lifespan(fast_api: FastAPI):
|
||||
if tracer_provider:
|
||||
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
|
||||
yield
|
||||
await runtime.rag_service.close()
|
||||
await runtime.documents.close()
|
||||
if tracer_provider:
|
||||
tracer_provider.shutdown()
|
||||
|
||||
@@ -69,7 +69,7 @@ app.include_router(pdf_edit_router)
|
||||
app.include_router(pdf_question_router)
|
||||
app.include_router(agent_draft_router)
|
||||
app.include_router(execution_router)
|
||||
app.include_router(rag_router)
|
||||
app.include_router(document_router)
|
||||
app.include_router(ledger_router)
|
||||
app.include_router(pdf_comments_router)
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ from stirling.agents import (
|
||||
)
|
||||
from stirling.agents.ledger import MathAuditorAgent
|
||||
from stirling.agents.pdf_comment import PdfCommentAgent
|
||||
from stirling.rag import RagService
|
||||
from stirling.documents import DocumentService
|
||||
from stirling.services import AppRuntime
|
||||
|
||||
|
||||
@@ -39,8 +39,8 @@ def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
|
||||
return request.app.state.execution_planning_agent
|
||||
|
||||
|
||||
def get_rag_service(request: Request) -> RagService:
|
||||
return request.app.state.runtime.rag_service
|
||||
def get_document_service(request: Request) -> DocumentService:
|
||||
return request.app.state.runtime.documents
|
||||
|
||||
|
||||
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
|
||||
|
||||
@@ -1,19 +1,19 @@
|
||||
from .agent_drafts import router as agent_draft_router
|
||||
from .documents import router as document_router
|
||||
from .execution import router as execution_router
|
||||
from .ledger import router as ledger_router
|
||||
from .orchestrator import router as orchestrator_router
|
||||
from .pdf_comments import router as pdf_comments_router
|
||||
from .pdf_edit import router as pdf_edit_router
|
||||
from .pdf_questions import router as pdf_question_router
|
||||
from .rag import router as rag_router
|
||||
|
||||
__all__ = [
|
||||
"agent_draft_router",
|
||||
"document_router",
|
||||
"execution_router",
|
||||
"ledger_router",
|
||||
"orchestrator_router",
|
||||
"pdf_comments_router",
|
||||
"pdf_edit_router",
|
||||
"pdf_question_router",
|
||||
"rag_router",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.api.dependencies import get_document_service
|
||||
from stirling.contracts import (
|
||||
DeleteDocumentResponse,
|
||||
IngestDocumentRequest,
|
||||
IngestDocumentResponse,
|
||||
)
|
||||
from stirling.documents import DocumentService
|
||||
from stirling.models import FileId
|
||||
|
||||
router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
|
||||
|
||||
|
||||
@router.post("", response_model=IngestDocumentResponse)
|
||||
async def ingest_document(
|
||||
request: IngestDocumentRequest,
|
||||
documents: Annotated[DocumentService, Depends(get_document_service)],
|
||||
) -> IngestDocumentResponse:
|
||||
"""Replace-ingest a document's content under ``document_id``.
|
||||
|
||||
Stores both representations in one shot:
|
||||
* embedded chunks for RAG search,
|
||||
* ordered page text for whole-document reading.
|
||||
Any previously-stored content for this document is removed first.
|
||||
"""
|
||||
pages = request.page_text or []
|
||||
chunks_indexed = await documents.ingest(
|
||||
collection=request.document_id,
|
||||
pages=pages,
|
||||
source=request.source,
|
||||
)
|
||||
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed)
|
||||
|
||||
|
||||
@router.delete("/{document_id}", response_model=DeleteDocumentResponse)
|
||||
async def delete_document(
|
||||
document_id: FileId,
|
||||
documents: Annotated[DocumentService, Depends(get_document_service)],
|
||||
) -> DeleteDocumentResponse:
|
||||
"""Remove a document's content. Idempotent."""
|
||||
existed = await documents.has_collection(document_id)
|
||||
if existed:
|
||||
await documents.delete_collection(document_id)
|
||||
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
|
||||
@@ -1,19 +1,170 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, assert_never
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from stirling.agents import OrchestratorAgent
|
||||
from stirling.api.dependencies import get_orchestrator_agent
|
||||
from stirling.contracts import OrchestratorRequest, OrchestratorResponse
|
||||
from stirling.contracts import OrchestratorRequest, OrchestratorResponse, ProgressEvent
|
||||
from stirling.services import reset_progress_emitter, set_progress_emitter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Cadence for keep-alive heartbeats on the streaming endpoint. Java forwards
|
||||
# them to the frontend as SSE comments; their job is to make every layer of
|
||||
# the connection visibly alive at this rhythm so disconnects surface within a
|
||||
# bounded window instead of waiting for the next progress event.
|
||||
HEARTBEAT_INTERVAL_SECONDS = 10.0
|
||||
|
||||
router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
|
||||
|
||||
|
||||
@router.post("", response_model=OrchestratorResponse)
|
||||
@router.post("")
|
||||
async def orchestrate(
|
||||
request: OrchestratorRequest,
|
||||
agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
|
||||
) -> OrchestratorResponse:
|
||||
return await agent.handle(request)
|
||||
) -> StreamingResponse:
|
||||
"""Run the orchestrator and stream NDJSON events.
|
||||
|
||||
Each output line is a JSON object with an ``event`` field. ``progress``
|
||||
events arrive whenever an inner agent reports work (e.g. each
|
||||
chunked-reasoner slice completing); the final ``result`` event carries the
|
||||
typed orchestrator response. ``error`` events surface failures without
|
||||
breaking the connection. ``heartbeat`` events fire on a fixed cadence to
|
||||
keep idle connections visibly alive so disconnects propagate.
|
||||
|
||||
The stream itself is the liveness signal: as long as events flow, work is
|
||||
alive. Java consumes this with a long total timeout and treats line
|
||||
arrival as forward progress.
|
||||
"""
|
||||
return StreamingResponse(
|
||||
_OrchestratorStream(
|
||||
agent=agent,
|
||||
request=request,
|
||||
heartbeat_interval_seconds=HEARTBEAT_INTERVAL_SECONDS,
|
||||
).iterate(),
|
||||
media_type="application/x-ndjson",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProgressFrame:
|
||||
event: ProgressEvent
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ResultFrame:
|
||||
response: OrchestratorResponse
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ErrorFrame:
|
||||
message: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _HeartbeatFrame:
|
||||
"""No payload: a heartbeat exists only to push bytes through the pipe.
|
||||
|
||||
Without periodic traffic, a slow workflow phase (e.g. all extractor
|
||||
workers busy on long calls) leaves the engine writer, Java's SSE
|
||||
forwarder, and the frontend's fetch all silently waiting. A closed
|
||||
connection at any layer wouldn't surface until the next real event,
|
||||
which could be many tens of seconds away. Heartbeats bound that window
|
||||
to :data:`HEARTBEAT_INTERVAL_SECONDS`.
|
||||
"""
|
||||
|
||||
|
||||
type _StreamFrame = _ProgressFrame | _ResultFrame | _ErrorFrame | _HeartbeatFrame
|
||||
|
||||
|
||||
def _serialize_frame(frame: _StreamFrame) -> bytes:
|
||||
"""Render a frame as one NDJSON line."""
|
||||
match frame:
|
||||
case _ProgressFrame(event=event):
|
||||
body = {"event": "progress", **event.model_dump(mode="json")}
|
||||
case _ResultFrame(response=response):
|
||||
body = {"event": "result", "response": response.model_dump(mode="json")}
|
||||
case _ErrorFrame(message=message):
|
||||
body = {"event": "error", "message": message}
|
||||
case _HeartbeatFrame():
|
||||
body = {"event": "heartbeat"}
|
||||
case _:
|
||||
assert_never(frame)
|
||||
return (json.dumps(body) + "\n").encode("utf-8")
|
||||
|
||||
|
||||
class _OrchestratorStream:
|
||||
"""Drives one streaming orchestrator request.
|
||||
|
||||
Owns the per-request queue and pumps progress events through it; the agent
|
||||
runs as a child task so its emissions and the streaming response interleave.
|
||||
A heartbeat task pushes keep-alive messages onto the same queue at a fixed
|
||||
cadence so the connection stays visibly alive between progress events.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
agent: OrchestratorAgent,
|
||||
request: OrchestratorRequest,
|
||||
heartbeat_interval_seconds: float,
|
||||
) -> None:
|
||||
self._agent = agent
|
||||
self._request = request
|
||||
self._heartbeat_interval_seconds = heartbeat_interval_seconds
|
||||
self._queue: asyncio.Queue[_StreamFrame | None] = asyncio.Queue()
|
||||
|
||||
async def iterate(self) -> AsyncIterator[bytes]:
|
||||
token = set_progress_emitter(self._emit_progress)
|
||||
agent_task = asyncio.create_task(self._run_agent())
|
||||
heartbeat_task = asyncio.create_task(self._emit_heartbeats())
|
||||
try:
|
||||
while True:
|
||||
frame = await self._queue.get()
|
||||
if frame is None:
|
||||
break
|
||||
yield _serialize_frame(frame)
|
||||
finally:
|
||||
reset_progress_emitter(token)
|
||||
await self._cancel_task(heartbeat_task)
|
||||
await self._cancel_task(agent_task)
|
||||
|
||||
async def _emit_progress(self, event: ProgressEvent) -> None:
|
||||
await self._queue.put(_ProgressFrame(event=event))
|
||||
|
||||
async def _emit_heartbeats(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(self._heartbeat_interval_seconds)
|
||||
await self._queue.put(_HeartbeatFrame())
|
||||
|
||||
async def _run_agent(self) -> None:
|
||||
try:
|
||||
response = await self._agent.handle(self._request)
|
||||
await self._queue.put(_ResultFrame(response=response))
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.exception("orchestrator stream failed")
|
||||
await self._queue.put(_ErrorFrame(message=str(exc)))
|
||||
finally:
|
||||
await self._queue.put(None)
|
||||
|
||||
@staticmethod
|
||||
async def _cancel_task(task: asyncio.Task[None]) -> None:
|
||||
if task.done():
|
||||
return
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception("background task failed during cancellation", exc_info=True)
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
|
||||
from stirling.api.dependencies import get_rag_service
|
||||
from stirling.contracts import (
|
||||
DeleteDocumentResponse,
|
||||
IngestDocumentRequest,
|
||||
IngestDocumentResponse,
|
||||
PdfContentType,
|
||||
)
|
||||
from stirling.models import FileId
|
||||
from stirling.rag import Document, RagService
|
||||
|
||||
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
|
||||
|
||||
|
||||
@router.post("/documents", response_model=IngestDocumentResponse)
|
||||
async def ingest_document(
|
||||
request: IngestDocumentRequest,
|
||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
||||
) -> IngestDocumentResponse:
|
||||
"""Replace-ingest a document's content under ``document_id``.
|
||||
|
||||
Any previously-stored content for this document is removed and the
|
||||
provided content replaces it wholesale. All pages are chunked up front
|
||||
and then embedded in a single batched call so large documents (e.g. a
|
||||
500-page book) don't fan out into hundreds of embedding requests.
|
||||
"""
|
||||
await rag.delete_collection(request.document_id)
|
||||
|
||||
chunks: list[Document] = []
|
||||
if request.page_text:
|
||||
for page in request.page_text:
|
||||
if not page.text.strip():
|
||||
continue
|
||||
chunks.extend(
|
||||
rag.chunk_text(
|
||||
text=page.text,
|
||||
source=f"{request.source}:page:{page.page_number}",
|
||||
base_metadata={
|
||||
"page_number": str(page.page_number),
|
||||
"content_type": PdfContentType.PAGE_TEXT.value,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
indexed = await rag.index_documents(request.document_id, chunks) if chunks else 0
|
||||
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=indexed)
|
||||
|
||||
|
||||
@router.delete("/documents/{document_id}", response_model=DeleteDocumentResponse)
|
||||
async def delete_document(
|
||||
document_id: FileId,
|
||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
||||
) -> DeleteDocumentResponse:
|
||||
"""Remove a document's content from RAG. Idempotent."""
|
||||
existed = await rag.has_collection(document_id)
|
||||
if existed:
|
||||
await rag.delete_collection(document_id)
|
||||
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
|
||||
@@ -38,18 +38,38 @@ class AppSettings(BaseSettings):
|
||||
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
|
||||
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
|
||||
|
||||
# Chunked reasoner settings (whole-document map-reduce).
|
||||
chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE")
|
||||
chunked_reasoner_concurrency: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CONCURRENCY")
|
||||
chunked_reasoner_worker_timeout_seconds: float = Field(
|
||||
validation_alias="STIRLING_CHUNKED_REASONER_WORKER_TIMEOUT_SECONDS"
|
||||
)
|
||||
# Maximum size, in characters, of the rendered notes block before the
|
||||
# reasoner folds slice notes hierarchically. The Anthropic context limit
|
||||
# is 200k tokens (~880k chars); we leave a generous margin for the
|
||||
# downstream agent's system prompt, history, tool definitions, and
|
||||
# response budget.
|
||||
chunked_reasoner_notes_char_budget: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET")
|
||||
|
||||
max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
|
||||
max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS")
|
||||
|
||||
log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
|
||||
log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
|
||||
# When true, raises httpx + httpcore logger levels so every outgoing
|
||||
# model SDK call is logged with timing. Use to diagnose worker stalls:
|
||||
# a hung request shows the "Request: POST ..." line with no matching
|
||||
# response line, confirming the hang is transport-layer (not in our
|
||||
# code or the Anthropic SDK itself). Off by default — DEBUG-level
|
||||
# output is high-volume.
|
||||
http_debug: bool = Field(default=False, validation_alias="STIRLING_HTTP_DEBUG")
|
||||
|
||||
posthog_enabled: bool = Field(validation_alias="STIRLING_POSTHOG_ENABLED")
|
||||
posthog_api_key: str = Field(validation_alias="STIRLING_POSTHOG_API_KEY")
|
||||
posthog_host: str = Field(validation_alias="STIRLING_POSTHOG_HOST")
|
||||
|
||||
|
||||
def _configure_logging(level_name: str, log_file: str) -> None:
|
||||
def _configure_logging(level_name: str, log_file: str, http_debug: bool) -> None:
|
||||
"""Configure the ``stirling`` logger hierarchy."""
|
||||
level = logging.getLevelNamesMapping().get(level_name.upper())
|
||||
if level is None:
|
||||
@@ -83,11 +103,43 @@ def _configure_logging(level_name: str, log_file: str) -> None:
|
||||
fh.setLevel(level)
|
||||
root.addHandler(fh)
|
||||
|
||||
if http_debug:
|
||||
_enable_http_debug(formatter)
|
||||
|
||||
|
||||
def _enable_http_debug(formatter: logging.Formatter) -> None:
|
||||
"""Surface every httpx/httpcore call against the Anthropic API.
|
||||
|
||||
httpx emits one INFO line per request with the URL and final status,
|
||||
which is the most useful signal for diagnosing hung worker calls: a
|
||||
successful call shows "Request" then "Response" within a second or two;
|
||||
a hung one shows "Request" with no matching response until it's
|
||||
cancelled. httpcore at DEBUG drills down to TCP / HTTP/2 stream events
|
||||
if the user wants to see exactly where bytes stop flowing.
|
||||
|
||||
The ``stirling`` console handler is scoped to its own logger tree, so
|
||||
we attach a dedicated stream handler here. Without it, httpx records
|
||||
propagate to the root logger which has no handler in our setup and the
|
||||
output is silently dropped.
|
||||
"""
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(formatter)
|
||||
handler.setLevel(logging.DEBUG)
|
||||
|
||||
for name, level in (("httpx", logging.INFO), ("httpcore", logging.DEBUG)):
|
||||
lg = logging.getLogger(name)
|
||||
lg.setLevel(level)
|
||||
# Idempotent: avoid stacking handlers on settings reload.
|
||||
if not any(getattr(h, "_stirling_http_debug", False) for h in lg.handlers):
|
||||
handler._stirling_http_debug = True # type: ignore[attr-defined]
|
||||
lg.addHandler(handler)
|
||||
lg.propagate = False
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def load_settings() -> AppSettings:
|
||||
load_dotenv(ENV_FILE)
|
||||
load_dotenv(ENV_LOCAL_FILE, override=True)
|
||||
settings = AppSettings.model_validate({})
|
||||
_configure_logging(settings.log_level, settings.log_file)
|
||||
_configure_logging(settings.log_level, settings.log_file, settings.http_debug)
|
||||
return settings
|
||||
|
||||
@@ -28,6 +28,14 @@ from .common import (
|
||||
format_conversation_history,
|
||||
format_file_names,
|
||||
)
|
||||
from .documents import (
|
||||
DeleteDocumentResponse,
|
||||
IngestDocumentRequest,
|
||||
IngestDocumentResponse,
|
||||
Page,
|
||||
PageRange,
|
||||
PageText,
|
||||
)
|
||||
from .execution import (
|
||||
AgentExecutionRequest,
|
||||
CannotContinueExecutionAction,
|
||||
@@ -79,11 +87,12 @@ from .pdf_questions import (
|
||||
PdfQuestionResponse,
|
||||
PdfQuestionTerminalResponse,
|
||||
)
|
||||
from .rag import (
|
||||
DeleteDocumentResponse,
|
||||
IngestDocumentRequest,
|
||||
IngestDocumentResponse,
|
||||
IngestedPageText,
|
||||
from .progress import (
|
||||
ProgressEvent,
|
||||
WholeDocCompressionRound,
|
||||
WholeDocReadDone,
|
||||
WholeDocReadStarted,
|
||||
WholeDocSliceDone,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
@@ -123,7 +132,6 @@ __all__ = [
|
||||
"HealthResponse",
|
||||
"IngestDocumentRequest",
|
||||
"IngestDocumentResponse",
|
||||
"IngestedPageText",
|
||||
"MathAuditorToolReportArtifact",
|
||||
"NeedContentFileRequest",
|
||||
"NeedContentResponse",
|
||||
@@ -131,6 +139,9 @@ __all__ = [
|
||||
"NextExecutionAction",
|
||||
"OrchestratorRequest",
|
||||
"OrchestratorResponse",
|
||||
"Page",
|
||||
"PageRange",
|
||||
"PageText",
|
||||
"PdfCommentInstruction",
|
||||
"PdfCommentReport",
|
||||
"PdfCommentRequest",
|
||||
@@ -146,6 +157,7 @@ __all__ = [
|
||||
"PdfQuestionResponse",
|
||||
"PdfQuestionTerminalResponse",
|
||||
"PdfTextSelection",
|
||||
"ProgressEvent",
|
||||
"Requisition",
|
||||
"Severity",
|
||||
"StepKind",
|
||||
@@ -156,6 +168,10 @@ __all__ = [
|
||||
"ToolReportArtifact",
|
||||
"UnsupportedCapabilityResponse",
|
||||
"Verdict",
|
||||
"WholeDocCompressionRound",
|
||||
"WholeDocReadDone",
|
||||
"WholeDocReadStarted",
|
||||
"WholeDocSliceDone",
|
||||
"WorkflowArtifact",
|
||||
"WorkflowOutcome",
|
||||
]
|
||||
|
||||
@@ -167,10 +167,10 @@ ToolReportArtifact = MathAuditorToolReportArtifact
|
||||
|
||||
|
||||
class NeedIngestResponse(ApiModel):
|
||||
"""Signal that the listed files must be ingested into RAG before the agent can continue.
|
||||
"""Signal that the listed files must be ingested before the agent can continue.
|
||||
|
||||
Java's handling: for each file, extract the requested content types, POST to
|
||||
``/api/v1/rag/documents`` keyed by ``file.id``, then retry the original request.
|
||||
``/api/v1/documents`` keyed by ``file.id``, then retry the original request.
|
||||
"""
|
||||
|
||||
outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .common import FileId
|
||||
|
||||
|
||||
class PageText(ApiModel):
|
||||
"""A single page of extracted text on the ingest wire."""
|
||||
|
||||
page_number: int = Field(ge=1)
|
||||
text: str
|
||||
|
||||
|
||||
class Page(ApiModel):
|
||||
"""A single page of a document, retrieved from storage.
|
||||
|
||||
``char_count`` is precomputed at ingest time and reported here so callers
|
||||
can budget how much content they want to read without first concatenating
|
||||
the text of every page.
|
||||
"""
|
||||
|
||||
page_number: int = Field(ge=1)
|
||||
text: str
|
||||
char_count: int = Field(ge=0)
|
||||
|
||||
|
||||
class PageRange(ApiModel):
|
||||
"""Inclusive page range for partial reads. Both bounds are 1-indexed."""
|
||||
|
||||
start: int = Field(ge=1)
|
||||
end: int = Field(ge=1)
|
||||
|
||||
|
||||
class IngestDocumentRequest(ApiModel):
|
||||
"""Replace-ingest a document's content under the given ``document_id``.
|
||||
|
||||
Each call wipes any previously-stored content for the document and writes
|
||||
both the vector-chunk and ordered-page representations from the supplied
|
||||
pages.
|
||||
|
||||
``source`` is a human-readable label (typically the original filename)
|
||||
that flows into chunk metadata so search results are readable when
|
||||
``document_id`` is a hash.
|
||||
"""
|
||||
|
||||
document_id: FileId = Field(min_length=1)
|
||||
source: str = Field(min_length=1)
|
||||
page_text: list[PageText] | None = None
|
||||
|
||||
|
||||
class IngestDocumentResponse(ApiModel):
|
||||
document_id: FileId
|
||||
chunks_indexed: int
|
||||
|
||||
|
||||
class DeleteDocumentResponse(ApiModel):
|
||||
document_id: FileId
|
||||
deleted: bool
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Progress events emitted by deep callees during a streaming orchestrator run.
|
||||
|
||||
Each subclass models one engine-side phase. The Java side forwards the JSON
|
||||
verbatim into ``AiWorkflowProgressEvent.engineDetail``; the frontend switches
|
||||
on ``phase`` and renders the typed fields.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
|
||||
class WholeDocReadStarted(ApiModel):
|
||||
phase: Literal["whole_doc_read_started"] = "whole_doc_read_started"
|
||||
question: str
|
||||
pages: int
|
||||
slices: int
|
||||
|
||||
|
||||
class WholeDocSliceDone(ApiModel):
|
||||
"""Emitted as each chunked-reasoner worker completes.
|
||||
|
||||
``completed`` is a monotonically increasing counter (1..total) reflecting
|
||||
the order in which workers finished, NOT the slice's position in the
|
||||
document. Callers showing "Read X of Y" should use this directly so X
|
||||
increments by one with each event.
|
||||
"""
|
||||
|
||||
phase: Literal["whole_doc_slice_done"] = "whole_doc_slice_done"
|
||||
completed: int
|
||||
total: int
|
||||
pages: str
|
||||
duration_ms: int
|
||||
excerpts: int
|
||||
facts: int
|
||||
|
||||
|
||||
class WholeDocCompressionRound(ApiModel):
|
||||
"""Emitted when the gathered slice notes exceed the synthesis context
|
||||
budget and the reasoner consolidates them with a fast-model fold pass.
|
||||
|
||||
Long documents (a 3000-page novel produces ~900k chars of raw notes)
|
||||
would otherwise overflow the smart-model's prompt. ``notes_in`` is the
|
||||
count entering the round; ``groups`` is the number of fold calls fired
|
||||
(each producing one consolidated note). One or two rounds usually fit;
|
||||
the event fires per round so callers can render "Consolidating notes
|
||||
(round N)..." rather than going silent through the fold.
|
||||
"""
|
||||
|
||||
phase: Literal["whole_doc_compression_round"] = "whole_doc_compression_round"
|
||||
round_number: int
|
||||
notes_in: int
|
||||
groups: int
|
||||
|
||||
|
||||
class WholeDocReadDone(ApiModel):
|
||||
phase: Literal["whole_doc_read_done"] = "whole_doc_read_done"
|
||||
completed: int
|
||||
slices: int
|
||||
duration_seconds: float
|
||||
|
||||
|
||||
type ProgressEvent = Annotated[
|
||||
WholeDocReadStarted | WholeDocSliceDone | WholeDocCompressionRound | WholeDocReadDone,
|
||||
Field(discriminator="phase"),
|
||||
]
|
||||
@@ -1,39 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from stirling.models import ApiModel
|
||||
|
||||
from .common import FileId
|
||||
|
||||
|
||||
class IngestedPageText(ApiModel):
|
||||
page_number: int = Field(ge=1)
|
||||
text: str
|
||||
|
||||
|
||||
class IngestDocumentRequest(ApiModel):
|
||||
"""Replace-ingest a document's content into RAG under the given document_id.
|
||||
|
||||
Each content-type field is optional; the endpoint replaces the document's entire
|
||||
stored content with whatever is provided. To add a content type later, call again
|
||||
with all content types the document should have (incremental-add-without-replace
|
||||
will be a separate endpoint if/when we need it).
|
||||
|
||||
``source`` is a human-readable label (typically the original filename) that flows
|
||||
into chunk metadata so search results are readable when document_id is a hash.
|
||||
"""
|
||||
|
||||
document_id: FileId = Field(min_length=1)
|
||||
source: str = Field(min_length=1)
|
||||
page_text: list[IngestedPageText] | None = None
|
||||
|
||||
|
||||
class IngestDocumentResponse(ApiModel):
|
||||
document_id: FileId
|
||||
chunks_indexed: int
|
||||
|
||||
|
||||
class DeleteDocumentResponse(ApiModel):
|
||||
document_id: FileId
|
||||
deleted: bool
|
||||
@@ -1,4 +1,14 @@
|
||||
# RAG Integration Guide
|
||||
# Document Storage
|
||||
|
||||
The `documents` package owns all stored content for a document under a single
|
||||
`collection` (file id):
|
||||
|
||||
* **Vector chunks** — small, embedded chunks for RAG-style retrieval.
|
||||
* **Ordered pages** — the original page text retained in document order, used
|
||||
for whole-document reading.
|
||||
|
||||
Both representations are populated by a single `ingest()` call and removed
|
||||
together by `delete_collection()`.
|
||||
|
||||
## Adding RAG to an Agent
|
||||
|
||||
@@ -22,19 +32,21 @@ That's it. The agent gets a `search_knowledge` tool it can call autonomously.
|
||||
|
||||
## Scoping to Specific Collections
|
||||
|
||||
Collections are named buckets of indexed documents — think folders. By default an agent searches everything in the store. Pass `collections=` to restrict it to only the docs indexed under those names.
|
||||
Collections are named buckets of indexed documents - think folders. By default
|
||||
an agent searches everything in the store. Pass `collections=` to restrict it
|
||||
to only the docs indexed under those names.
|
||||
|
||||
```python
|
||||
from stirling.rag import RagCapability
|
||||
from stirling.documents import RagCapability
|
||||
|
||||
# Only searches docs indexed under "company-docs" — ignores everything else
|
||||
scoped = RagCapability(runtime.rag_service, collections=["company-docs"], top_k=3)
|
||||
# Only searches docs indexed under "company-docs"
|
||||
scoped = RagCapability(runtime.documents, collections=["company-docs"], top_k=3)
|
||||
|
||||
# Searches multiple collections
|
||||
multi = RagCapability(runtime.rag_service, collections=["company-docs", "product-specs"])
|
||||
multi = RagCapability(runtime.documents, collections=["company-docs", "product-specs"])
|
||||
|
||||
# No collections arg = searches all collections in the store
|
||||
everything = RagCapability(runtime.rag_service)
|
||||
everything = RagCapability(runtime.documents)
|
||||
```
|
||||
|
||||
## Config
|
||||
@@ -51,7 +63,8 @@ STIRLING_RAG_CHUNK_OVERLAP=64
|
||||
STIRLING_RAG_TOP_K=5
|
||||
```
|
||||
|
||||
Provider credentials (and any local overrides) go in the uncommitted `engine/.env.local`:
|
||||
Provider credentials (and any local overrides) go in the uncommitted
|
||||
`engine/.env.local`:
|
||||
|
||||
```
|
||||
VOYAGE_API_KEY=your-key
|
||||
@@ -59,13 +72,17 @@ VOYAGE_API_KEY=your-key
|
||||
|
||||
## Backends
|
||||
|
||||
**`sqlite`** — Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev and self-hosted deployments.
|
||||
**`sqlite`** - Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev
|
||||
and self-hosted deployments.
|
||||
|
||||
**`pgvector`** — External PostgreSQL with the `vector` extension. Point `STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
|
||||
**`pgvector`** - External PostgreSQL with the `vector` extension. Point
|
||||
`STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
|
||||
|
||||
Both backends implement the same `VectorStore` interface, so agents and the RAG service work identically regardless of which you pick.
|
||||
Both backends implement the same `DocumentStore` interface, so agents and the
|
||||
service work identically regardless of which you pick.
|
||||
|
||||
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model string accordingly and point at the server via its native env var:
|
||||
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model
|
||||
string accordingly and point at the server via its native env var:
|
||||
|
||||
```
|
||||
# Ollama running on another machine
|
||||
@@ -81,8 +98,5 @@ OPENAI_BASE_URL=http://192.168.1.50:8080/v1
|
||||
|
||||
| Method | Endpoint | Purpose |
|
||||
|--------|----------|---------|
|
||||
| GET | `/api/v1/rag/status` | Report embedding model and existing collections |
|
||||
| POST | `/api/v1/rag/index` | Index text into a collection |
|
||||
| POST | `/api/v1/rag/search` | Search a collection |
|
||||
| GET | `/api/v1/rag/collections` | List collections |
|
||||
| DELETE | `/api/v1/rag/collections/{name}` | Delete a collection |
|
||||
| POST | `/api/v1/documents` | Replace-ingest a document's pages |
|
||||
| DELETE | `/api/v1/documents/{document_id}` | Delete a document's stored content |
|
||||
@@ -0,0 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.documents.embedder import EmbeddingService
|
||||
from stirling.documents.pgvector_store import PgVectorStore
|
||||
from stirling.documents.rag_capability import RagCapability
|
||||
from stirling.documents.service import DocumentService
|
||||
from stirling.documents.sqlite_vec_store import SqliteVecStore
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
__all__ = [
|
||||
"Document",
|
||||
"DocumentService",
|
||||
"DocumentStore",
|
||||
"EmbeddingService",
|
||||
"PgVectorStore",
|
||||
"RagCapability",
|
||||
"SearchResult",
|
||||
"SqliteVecStore",
|
||||
"StoredPage",
|
||||
]
|
||||
@@ -2,8 +2,8 @@ from __future__ import annotations
|
||||
|
||||
from pydantic_ai import Embedder
|
||||
|
||||
from stirling.rag.chunker import chunk_text
|
||||
from stirling.rag.store import Document
|
||||
from stirling.documents.chunker import chunk_text
|
||||
from stirling.documents.store import Document
|
||||
|
||||
# Keep each upstream embed request under every major provider's per-call limit while
|
||||
# still batching large enough that a book-sized document ingests in a reasonable number
|
||||
+88
-6
@@ -5,14 +5,20 @@ import json
|
||||
import psycopg
|
||||
from pgvector.psycopg import register_vector_async
|
||||
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
|
||||
class PgVectorStore(VectorStore):
|
||||
class PgVectorStore(DocumentStore):
|
||||
"""PostgreSQL + pgvector backed store.
|
||||
|
||||
Connects to an external Postgres instance (DSN provided via config) and uses the
|
||||
`vector` extension for similarity search. The schema is created on first use.
|
||||
|
||||
Holds two tables under the same connection:
|
||||
|
||||
* ``rag_documents`` - vector chunks for RAG search.
|
||||
* ``document_pages`` - ordered page text for whole-document reading.
|
||||
"""
|
||||
|
||||
def __init__(self, dsn: str) -> None:
|
||||
@@ -32,11 +38,20 @@ class PgVectorStore(VectorStore):
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||
await cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents_meta (
|
||||
collection TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS rag_documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
collection TEXT NOT NULL
|
||||
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
text TEXT NOT NULL,
|
||||
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||||
embedding vector NOT NULL,
|
||||
@@ -45,9 +60,36 @@ class PgVectorStore(VectorStore):
|
||||
"""
|
||||
)
|
||||
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
|
||||
await cur.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS document_pages (
|
||||
collection TEXT NOT NULL
|
||||
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
page_number INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
char_count INTEGER NOT NULL,
|
||||
PRIMARY KEY (collection, page_number)
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||
await conn.commit()
|
||||
self._initialized = True
|
||||
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
INSERT INTO documents_meta (collection, source)
|
||||
VALUES (%s, %s)
|
||||
ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source
|
||||
""",
|
||||
(collection, source),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
@@ -106,18 +148,58 @@ class PgVectorStore(VectorStore):
|
||||
for r in rows
|
||||
]
|
||||
|
||||
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,))
|
||||
if pages:
|
||||
await cur.executemany(
|
||||
"""
|
||||
INSERT INTO document_pages (collection, page_number, text, char_count)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
""",
|
||||
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
if page_range is None:
|
||||
await cur.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = %s ORDER BY page_number",
|
||||
(collection,),
|
||||
)
|
||||
else:
|
||||
await cur.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = %s AND page_number BETWEEN %s AND %s "
|
||||
"ORDER BY page_number",
|
||||
(collection, page_range.start, page_range.end),
|
||||
)
|
||||
rows = await cur.fetchall()
|
||||
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
|
||||
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("DELETE FROM rag_documents WHERE collection = %s", (collection,))
|
||||
# Cascade FKs handle rag_documents and document_pages.
|
||||
await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,))
|
||||
await conn.commit()
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
await self._ensure_schema()
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("SELECT DISTINCT collection FROM rag_documents ORDER BY collection")
|
||||
await cur.execute("SELECT collection FROM documents_meta ORDER BY collection")
|
||||
rows = await cur.fetchall()
|
||||
return [r[0] for r in rows]
|
||||
|
||||
@@ -126,7 +208,7 @@ class PgVectorStore(VectorStore):
|
||||
async with await self._connect() as conn:
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"SELECT 1 FROM rag_documents WHERE collection = %s LIMIT 1",
|
||||
"SELECT 1 FROM documents_meta WHERE collection = %s",
|
||||
(collection,),
|
||||
)
|
||||
row = await cur.fetchone()
|
||||
+7
-7
@@ -6,9 +6,9 @@ from collections.abc import Awaitable, Callable
|
||||
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
||||
from pydantic_ai.toolsets import AbstractToolset
|
||||
|
||||
from stirling.documents.service import DocumentService
|
||||
from stirling.documents.store import SearchResult
|
||||
from stirling.models import FileId
|
||||
from stirling.rag.service import RagService
|
||||
from stirling.rag.store import SearchResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -34,12 +34,12 @@ class RagCapability:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rag_service: RagService,
|
||||
documents: DocumentService,
|
||||
collections: list[FileId] | None = None,
|
||||
top_k: int = 5,
|
||||
max_searches: int = 5,
|
||||
) -> None:
|
||||
self._rag_service = rag_service
|
||||
self._documents = documents
|
||||
self._collections = collections
|
||||
self._top_k = top_k
|
||||
self._max_searches = max_searches
|
||||
@@ -74,7 +74,7 @@ class RagCapability:
|
||||
)
|
||||
|
||||
async def _dynamic_instructions(self) -> str:
|
||||
collections = await self._rag_service.list_collections()
|
||||
collections = await self._documents.list_collections()
|
||||
if collections:
|
||||
names = ", ".join(collections)
|
||||
collection_desc = f"the following knowledge base collections: {names}"
|
||||
@@ -115,12 +115,12 @@ class RagCapability:
|
||||
if self._collections:
|
||||
all_results = []
|
||||
for col in self._collections:
|
||||
col_results = await self._rag_service.search(query, collection=col, top_k=k)
|
||||
col_results = await self._documents.search(query, collection=col, top_k=k)
|
||||
all_results.extend(col_results)
|
||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||
results = all_results[:k]
|
||||
else:
|
||||
results = await self._rag_service.search(query, top_k=k)
|
||||
results = await self._documents.search(query, top_k=k)
|
||||
|
||||
if not results:
|
||||
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
|
||||
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange, PageText
|
||||
from stirling.documents.embedder import EmbeddingService
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
from stirling.models import FileId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PAGE_NUMBER_METADATA_KEY = "page_number"
|
||||
CONTENT_TYPE_METADATA_KEY = "content_type"
|
||||
PAGE_TEXT_CONTENT_TYPE = "page_text"
|
||||
|
||||
|
||||
class DocumentService:
|
||||
"""Top-level facade for stored document content.
|
||||
|
||||
Holds two representations of every document under a single ``collection``:
|
||||
|
||||
* **Vector chunks** for RAG-style semantic retrieval (``search``).
|
||||
* **Ordered pages** for whole-document reading (``read_pages``).
|
||||
|
||||
Both are populated by :meth:`ingest` from a single ``pages`` payload. Agents
|
||||
pick the strategy that fits the question; they don't need to know which
|
||||
storage they're hitting.
|
||||
"""
|
||||
|
||||
def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
|
||||
self._embedder = embedder
|
||||
self._store = store
|
||||
self._default_top_k = default_top_k
|
||||
|
||||
async def ingest(
|
||||
self,
|
||||
collection: FileId,
|
||||
pages: list[PageText],
|
||||
source: str,
|
||||
) -> int:
|
||||
"""Replace-ingest a document. Returns the number of vector chunks indexed.
|
||||
|
||||
This wipes any previously-stored content for ``collection`` and writes
|
||||
both the vector-chunk and page-text representations from the same
|
||||
``pages`` payload. Pages with empty/whitespace-only text are skipped
|
||||
for chunking but still written to the page store so page numbering is
|
||||
preserved end-to-end.
|
||||
"""
|
||||
await self._store.delete_collection(collection)
|
||||
await self._store.ensure_collection(collection, source)
|
||||
|
||||
stored_pages = [StoredPage(page_number=p.page_number, text=p.text, char_count=len(p.text)) for p in pages]
|
||||
await self._store.add_pages(collection, stored_pages)
|
||||
|
||||
chunks: list[Document] = []
|
||||
for page in pages:
|
||||
if not page.text.strip():
|
||||
continue
|
||||
chunks.extend(
|
||||
self._embedder.chunk_and_prepare(
|
||||
text=page.text,
|
||||
source=f"{source}:page:{page.page_number}",
|
||||
base_metadata={
|
||||
PAGE_NUMBER_METADATA_KEY: str(page.page_number),
|
||||
CONTENT_TYPE_METADATA_KEY: PAGE_TEXT_CONTENT_TYPE,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
if not chunks:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in chunks])
|
||||
await self._store.add_documents(collection, chunks, embeddings)
|
||||
return len(chunks)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
collection: FileId | None = None,
|
||||
top_k: int | None = None,
|
||||
) -> list[SearchResult]:
|
||||
"""Embed query and search across one or all collections.
|
||||
|
||||
If collection is None, searches all available collections and merges results.
|
||||
"""
|
||||
k = top_k if top_k is not None else self._default_top_k
|
||||
query_embedding = await self._embedder.embed_query(query)
|
||||
|
||||
if collection is not None:
|
||||
if not await self._store.has_collection(collection):
|
||||
return []
|
||||
return await self._store.search(collection, query_embedding, k)
|
||||
|
||||
# Search all collections, skipping any that error (e.g. dimension mismatch)
|
||||
collections = await self._store.list_collections()
|
||||
all_results: list[SearchResult] = []
|
||||
for col_name in collections:
|
||||
try:
|
||||
results = await self._store.search(col_name, query_embedding, k)
|
||||
all_results.extend(results)
|
||||
except Exception: # noqa: BLE001 - any backend error on one collection should not stop the others
|
||||
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
|
||||
|
||||
# Sort by score descending, return top_k across all collections
|
||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||
return all_results[:k]
|
||||
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: FileId,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
"""Return ordered page text for ``collection``.
|
||||
|
||||
Empty list if the collection has no stored pages.
|
||||
"""
|
||||
return await self._store.read_pages(collection, page_range)
|
||||
|
||||
async def delete_collection(self, collection: FileId) -> None:
|
||||
"""Remove a collection's chunks and pages."""
|
||||
await self._store.delete_collection(collection)
|
||||
|
||||
async def has_collection(self, collection: FileId) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
return await self._store.has_collection(collection)
|
||||
|
||||
async def list_collections(self) -> list[FileId]:
|
||||
"""List all available collections."""
|
||||
return [FileId(name) for name in await self._store.list_collections()]
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Release the underlying store's resources."""
|
||||
await self._store.close()
|
||||
+90
-12
@@ -9,10 +9,11 @@ from pathlib import Path
|
||||
|
||||
import sqlite_vec
|
||||
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||
|
||||
|
||||
class SqliteVecStore(VectorStore):
|
||||
class SqliteVecStore(DocumentStore):
|
||||
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
|
||||
|
||||
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
|
||||
@@ -32,6 +33,8 @@ class SqliteVecStore(VectorStore):
|
||||
conn.enable_load_extension(True)
|
||||
sqlite_vec.load(conn)
|
||||
conn.enable_load_extension(False)
|
||||
# Required so cascade deletes from documents_meta clean up child tables.
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
if self._db_path is not None:
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
@@ -45,10 +48,18 @@ class SqliteVecStore(VectorStore):
|
||||
return cls(":memory:")
|
||||
|
||||
def _init_schema(self) -> None:
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents_meta (
|
||||
collection TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS collections (
|
||||
name TEXT PRIMARY KEY,
|
||||
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
dim INTEGER NOT NULL,
|
||||
table_name TEXT NOT NULL
|
||||
)
|
||||
@@ -58,7 +69,7 @@ class SqliteVecStore(VectorStore):
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL,
|
||||
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
text TEXT NOT NULL,
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
vec_rowid INTEGER NOT NULL,
|
||||
@@ -67,6 +78,32 @@ class SqliteVecStore(VectorStore):
|
||||
"""
|
||||
)
|
||||
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS document_pages (
|
||||
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
page_number INTEGER NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
char_count INTEGER NOT NULL,
|
||||
PRIMARY KEY (collection, page_number)
|
||||
)
|
||||
"""
|
||||
)
|
||||
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||
self._conn.commit()
|
||||
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_ensure_collection, collection, source)
|
||||
|
||||
def _sync_ensure_collection(self, collection: str, source: str) -> None:
|
||||
self._conn.execute(
|
||||
"""
|
||||
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
|
||||
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
|
||||
""",
|
||||
(collection, source),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
@staticmethod
|
||||
@@ -196,18 +233,56 @@ class SqliteVecStore(VectorStore):
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_add_pages, collection, pages)
|
||||
|
||||
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,))
|
||||
if pages:
|
||||
self._conn.executemany(
|
||||
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)",
|
||||
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_read_pages, collection, page_range)
|
||||
|
||||
def _sync_read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None,
|
||||
) -> list[Page]:
|
||||
if page_range is None:
|
||||
rows = self._conn.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number",
|
||||
(collection,),
|
||||
).fetchall()
|
||||
else:
|
||||
rows = self._conn.execute(
|
||||
"SELECT page_number, text, char_count FROM document_pages "
|
||||
"WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number",
|
||||
(collection, page_range.start, page_range.end),
|
||||
).fetchall()
|
||||
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
|
||||
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_delete_collection, collection)
|
||||
|
||||
def _sync_delete_collection(self, collection: str) -> None:
|
||||
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
|
||||
# (collections, documents, document_pages) when documents_meta is deleted.
|
||||
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is None:
|
||||
return
|
||||
table_name = row[0]
|
||||
self._conn.execute(f"DROP TABLE IF EXISTS {table_name}")
|
||||
self._conn.execute("DELETE FROM documents WHERE collection = ?", (collection,))
|
||||
self._conn.execute("DELETE FROM collections WHERE name = ?", (collection,))
|
||||
if row is not None:
|
||||
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
|
||||
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,))
|
||||
self._conn.commit()
|
||||
|
||||
async def list_collections(self) -> list[str]:
|
||||
@@ -215,7 +290,7 @@ class SqliteVecStore(VectorStore):
|
||||
return await asyncio.to_thread(self._sync_list_collections)
|
||||
|
||||
def _sync_list_collections(self) -> list[str]:
|
||||
rows = self._conn.execute("SELECT name FROM collections ORDER BY name").fetchall()
|
||||
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall()
|
||||
return [r[0] for r in rows]
|
||||
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
@@ -223,7 +298,10 @@ class SqliteVecStore(VectorStore):
|
||||
return await asyncio.to_thread(self._sync_has_collection, collection)
|
||||
|
||||
def _sync_has_collection(self, collection: str) -> bool:
|
||||
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
row = self._conn.execute(
|
||||
"SELECT 1 FROM documents_meta WHERE collection = ?",
|
||||
(collection,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
async def close(self) -> None:
|
||||
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from stirling.contracts.documents import Page, PageRange
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""A chunk of text with metadata, ready for embedding and storage."""
|
||||
|
||||
id: str
|
||||
text: str
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""A document returned from a vector search with its relevance score."""
|
||||
|
||||
document: Document
|
||||
score: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class StoredPage:
|
||||
"""A page as written to the store. ``char_count`` is precomputed at ingest."""
|
||||
|
||||
page_number: int
|
||||
text: str
|
||||
char_count: int
|
||||
|
||||
|
||||
class DocumentStore(ABC):
|
||||
"""Abstract interface for document storage backends.
|
||||
|
||||
Backends hold two representations of every document:
|
||||
|
||||
* **Vector chunks** - small, embedded chunks used for RAG search.
|
||||
* **Ordered pages** - the original page text retained in document order,
|
||||
used for whole-document reading.
|
||||
|
||||
Both representations live under the same ``collection`` (file id) and are
|
||||
rooted at a single parent row in ``documents_meta``. Removing that parent
|
||||
row cascades to both child representations, so :meth:`delete_collection`
|
||||
is one logical delete.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
"""Upsert the top-level ``documents_meta`` row for this collection.
|
||||
|
||||
Must be called before :meth:`add_pages` or :meth:`add_documents`. Both
|
||||
of those write into child tables that hold a foreign key to the parent
|
||||
row, so it must exist first.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
"""Store vector chunks with their embeddings in the named collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 5,
|
||||
) -> list[SearchResult]:
|
||||
"""Return the top_k most similar vector chunks from the collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
"""Replace the stored pages for ``collection`` with the supplied pages.
|
||||
|
||||
Implementations must remove any previously-stored pages for the
|
||||
collection before writing, so callers can re-ingest by calling this
|
||||
method again.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def read_pages(
|
||||
self,
|
||||
collection: str,
|
||||
page_range: PageRange | None = None,
|
||||
) -> list[Page]:
|
||||
"""Return ordered pages for ``collection``.
|
||||
|
||||
If ``page_range`` is ``None`` all pages are returned. Otherwise only
|
||||
pages whose ``page_number`` falls within the inclusive range are
|
||||
returned. Pages are always ordered by ``page_number`` ascending.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
"""Remove a collection's chunks and pages."""
|
||||
|
||||
@abstractmethod
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""Return names of all existing collections."""
|
||||
|
||||
@abstractmethod
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Release any resources held by the store (connections, handles, etc.)."""
|
||||
@@ -12,9 +12,12 @@ FileId = NewType("FileId", str)
|
||||
|
||||
|
||||
class ApiModel(BaseModel):
|
||||
"""Base for every contract model crossing a service boundary."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
alias_generator=to_camel,
|
||||
extra="forbid",
|
||||
validate_by_name=True,
|
||||
validate_by_alias=True,
|
||||
serialize_by_alias=True,
|
||||
)
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from stirling.rag.capability import RagCapability
|
||||
from stirling.rag.embedder import EmbeddingService
|
||||
from stirling.rag.pgvector_store import PgVectorStore
|
||||
from stirling.rag.service import RagService
|
||||
from stirling.rag.sqlite_vec_store import SqliteVecStore
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
__all__ = [
|
||||
"Document",
|
||||
"EmbeddingService",
|
||||
"PgVectorStore",
|
||||
"RagCapability",
|
||||
"RagService",
|
||||
"SearchResult",
|
||||
"SqliteVecStore",
|
||||
"VectorStore",
|
||||
]
|
||||
@@ -1,102 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from stirling.models import FileId
|
||||
from stirling.rag.embedder import EmbeddingService
|
||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RagService:
|
||||
"""Orchestrates embedding and vector storage for RAG workflows."""
|
||||
|
||||
def __init__(self, embedder: EmbeddingService, store: VectorStore, default_top_k: int = 5) -> None:
|
||||
self._embedder = embedder
|
||||
self._store = store
|
||||
self._default_top_k = default_top_k
|
||||
|
||||
async def index_text(
|
||||
self,
|
||||
collection: FileId,
|
||||
text: str,
|
||||
source: str = "",
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> int:
|
||||
"""Chunk, embed, and store text. Returns the number of chunks indexed."""
|
||||
documents = self._embedder.chunk_and_prepare(text, source=source, base_metadata=metadata)
|
||||
if not documents:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
||||
await self._store.add_documents(collection, documents, embeddings)
|
||||
return len(documents)
|
||||
|
||||
async def index_documents(self, collection: FileId, documents: list[Document]) -> int:
|
||||
"""Embed and store pre-chunked documents. Returns the number stored."""
|
||||
if not documents:
|
||||
return 0
|
||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
||||
await self._store.add_documents(collection, documents, embeddings)
|
||||
return len(documents)
|
||||
|
||||
def chunk_text(
|
||||
self,
|
||||
text: str,
|
||||
source: str = "",
|
||||
base_metadata: dict[str, str] | None = None,
|
||||
) -> list[Document]:
|
||||
"""Chunk text into Document objects ready for indexing. Does NOT embed.
|
||||
|
||||
Exposed so callers that ingest many chunks can accumulate them across calls
|
||||
and then pass the full batch to ``index_documents`` for a single embedding pass.
|
||||
"""
|
||||
return self._embedder.chunk_and_prepare(text, source=source, base_metadata=base_metadata)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: str,
|
||||
collection: FileId | None = None,
|
||||
top_k: int | None = None,
|
||||
) -> list[SearchResult]:
|
||||
"""Embed query and search across one or all collections.
|
||||
|
||||
If collection is None, searches all available collections and merges results.
|
||||
"""
|
||||
k = top_k if top_k is not None else self._default_top_k
|
||||
query_embedding = await self._embedder.embed_query(query)
|
||||
|
||||
if collection is not None:
|
||||
if not await self._store.has_collection(collection):
|
||||
return []
|
||||
return await self._store.search(collection, query_embedding, k)
|
||||
|
||||
# Search all collections, skipping any that error (e.g. dimension mismatch)
|
||||
collections = await self._store.list_collections()
|
||||
all_results: list[SearchResult] = []
|
||||
for col_name in collections:
|
||||
try:
|
||||
results = await self._store.search(col_name, query_embedding, k)
|
||||
all_results.extend(results)
|
||||
except Exception: # noqa: BLE001 — any backend error on one collection should not stop the others
|
||||
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
|
||||
|
||||
# Sort by score descending, return top_k across all collections
|
||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||
return all_results[:k]
|
||||
|
||||
async def delete_collection(self, collection: FileId) -> None:
|
||||
"""Remove a collection and all its documents."""
|
||||
await self._store.delete_collection(collection)
|
||||
|
||||
async def has_collection(self, collection: FileId) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
return await self._store.has_collection(collection)
|
||||
|
||||
async def list_collections(self) -> list[FileId]:
|
||||
"""List all available collections."""
|
||||
return [FileId(name) for name in await self._store.list_collections()]
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Release the underlying vector store's resources."""
|
||||
await self._store.close()
|
||||
@@ -1,63 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class Document:
|
||||
"""A chunk of text with metadata, ready for embedding and storage."""
|
||||
|
||||
id: str
|
||||
text: str
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""A document returned from a vector search with its relevance score."""
|
||||
|
||||
document: Document
|
||||
score: float
|
||||
|
||||
|
||||
class VectorStore(ABC):
|
||||
"""Abstract interface for vector storage backends.
|
||||
|
||||
Implementations must handle persistence, collection management,
|
||||
and nearest-neighbor search over pre-computed embeddings.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def add_documents(
|
||||
self,
|
||||
collection: str,
|
||||
documents: list[Document],
|
||||
embeddings: list[list[float]],
|
||||
) -> None:
|
||||
"""Store documents with their embeddings in the named collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
collection: str,
|
||||
query_embedding: list[float],
|
||||
top_k: int = 5,
|
||||
) -> list[SearchResult]:
|
||||
"""Return the top_k most similar documents from the collection."""
|
||||
|
||||
@abstractmethod
|
||||
async def delete_collection(self, collection: str) -> None:
|
||||
"""Remove a collection and all its documents."""
|
||||
|
||||
@abstractmethod
|
||||
async def list_collections(self) -> list[str]:
|
||||
"""Return names of all existing collections."""
|
||||
|
||||
@abstractmethod
|
||||
async def has_collection(self, collection: str) -> bool:
|
||||
"""Check whether a collection exists."""
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Release any resources held by the store (connections, handles, etc.)."""
|
||||
@@ -1,11 +1,21 @@
|
||||
"""Shared services used by the Stirling AI runtime."""
|
||||
|
||||
from .progress import (
|
||||
ProgressEmitter,
|
||||
emit_progress,
|
||||
reset_progress_emitter,
|
||||
set_progress_emitter,
|
||||
)
|
||||
from .runtime import AppRuntime, build_model_settings, build_runtime
|
||||
from .tracking import setup_posthog_tracking
|
||||
|
||||
__all__ = [
|
||||
"AppRuntime",
|
||||
"ProgressEmitter",
|
||||
"build_model_settings",
|
||||
"build_runtime",
|
||||
"emit_progress",
|
||||
"reset_progress_emitter",
|
||||
"set_progress_emitter",
|
||||
"setup_posthog_tracking",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Per-request progress emission, plumbed via a ContextVar so deep call stacks
|
||||
can publish typed events to the streaming orchestrator endpoint without every
|
||||
intermediate layer knowing about it.
|
||||
|
||||
Outside a streaming request no emitter is bound and ``emit_progress`` is a
|
||||
no-op, so callers in agents/services can emit unconditionally.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextvars import ContextVar, Token
|
||||
|
||||
from stirling.contracts import ProgressEvent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
type ProgressEmitter = Callable[[ProgressEvent], Awaitable[None]]
|
||||
|
||||
_emitter: ContextVar[ProgressEmitter | None] = ContextVar("stirling_progress_emitter", default=None)
|
||||
|
||||
|
||||
def set_progress_emitter(emitter: ProgressEmitter | None) -> Token[ProgressEmitter | None]:
|
||||
return _emitter.set(emitter)
|
||||
|
||||
|
||||
def reset_progress_emitter(token: Token[ProgressEmitter | None]) -> None:
|
||||
_emitter.reset(token)
|
||||
|
||||
|
||||
async def emit_progress(event: ProgressEvent) -> None:
|
||||
"""Publish ``event`` to the current request's emitter, if any.
|
||||
|
||||
Failures inside the emitter are logged and swallowed so progress emission
|
||||
can never break the work it's reporting on.
|
||||
"""
|
||||
emitter = _emitter.get()
|
||||
if emitter is None:
|
||||
return
|
||||
try:
|
||||
await emitter(event)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("progress emitter raised; dropping event %r", event.phase)
|
||||
@@ -4,28 +4,49 @@ import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import assert_never
|
||||
|
||||
import httpx
|
||||
from pydantic_ai.models import Model, infer_model
|
||||
from pydantic_ai.models.anthropic import AnthropicModel
|
||||
from pydantic_ai.providers.anthropic import AnthropicProvider
|
||||
from pydantic_ai.settings import ModelSettings
|
||||
|
||||
from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
|
||||
from stirling.rag import (
|
||||
from stirling.documents import (
|
||||
DocumentService,
|
||||
DocumentStore,
|
||||
EmbeddingService,
|
||||
PgVectorStore,
|
||||
RagCapability,
|
||||
RagService,
|
||||
SqliteVecStore,
|
||||
VectorStore,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _build_anthropic_http_client() -> httpx.AsyncClient:
|
||||
"""Build the httpx client used for Anthropic API calls.
|
||||
|
||||
We disable connection-pool keepalive so every request opens a fresh
|
||||
TCP+TLS connection. The default HTTP/1.1 pool reuses connections that
|
||||
Anthropic's front door (Cloudflare) sometimes closes silently between
|
||||
requests; the next request that picks up a stale connection sends its
|
||||
body into a black hole and never gets a response, hanging until our
|
||||
chunked-reasoner timeout fires.
|
||||
|
||||
A fresh handshake costs ~150ms — rounding error against a 5-15s LLM
|
||||
call. The trade is determinism: we never reuse a connection that might
|
||||
have died in the pool. See ``STIRLING_HTTP_DEBUG`` traces of slice 6
|
||||
on 2026-05-06 for the concrete failure mode this addresses.
|
||||
"""
|
||||
return httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=0))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AppRuntime:
|
||||
settings: AppSettings
|
||||
fast_model: Model
|
||||
smart_model: Model
|
||||
rag_service: RagService
|
||||
documents: DocumentService
|
||||
rag_capability: RagCapability
|
||||
|
||||
@property
|
||||
@@ -53,47 +74,62 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
|
||||
raise ValueError(f"Unsupported model {model_name}. This model does not support structured outputs.")
|
||||
|
||||
|
||||
def _build_vector_store(settings: AppSettings) -> VectorStore:
|
||||
"""Build the configured vector store backend."""
|
||||
def _build_document_store(settings: AppSettings) -> DocumentStore:
|
||||
"""Build the configured document store backend."""
|
||||
if settings.rag_backend == RagBackend.SQLITE:
|
||||
store_path = settings.rag_store_path
|
||||
# Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
|
||||
if str(store_path) != ":memory:" and not store_path.is_absolute():
|
||||
store_path = ENGINE_ROOT / store_path
|
||||
logger.info("RAG backend=sqlite, db_path=%s", store_path)
|
||||
logger.info("Document store backend=sqlite, db_path=%s", store_path)
|
||||
return SqliteVecStore(db_path=store_path)
|
||||
if settings.rag_backend == RagBackend.PGVECTOR:
|
||||
logger.info("RAG backend=pgvector, dsn=<configured>")
|
||||
logger.info("Document store backend=pgvector, dsn=<configured>")
|
||||
return PgVectorStore(dsn=settings.rag_pgvector_dsn)
|
||||
assert_never(settings.rag_backend)
|
||||
|
||||
|
||||
def _build_rag(settings: AppSettings) -> tuple[RagService, RagCapability]:
|
||||
"""Build the RAG service and capability."""
|
||||
logger.info("RAG: embedding_model=%s", settings.rag_embedding_model)
|
||||
def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]:
|
||||
"""Build the document service and the RAG-search capability that wraps it."""
|
||||
logger.info("Documents: embedding_model=%s", settings.rag_embedding_model)
|
||||
embedder = EmbeddingService(
|
||||
model_name=settings.rag_embedding_model,
|
||||
chunk_size=settings.rag_chunk_size,
|
||||
chunk_overlap=settings.rag_chunk_overlap,
|
||||
)
|
||||
store = _build_vector_store(settings)
|
||||
service = RagService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
|
||||
capability = RagCapability(rag_service=service, top_k=settings.rag_default_top_k)
|
||||
store = _build_document_store(settings)
|
||||
service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
|
||||
capability = RagCapability(documents=service, top_k=settings.rag_default_top_k)
|
||||
return service, capability
|
||||
|
||||
|
||||
def build_runtime(settings: AppSettings) -> AppRuntime:
|
||||
fast_model = infer_model(settings.fast_model_name)
|
||||
smart_model = infer_model(settings.smart_model_name)
|
||||
fast_model = _build_model(settings.fast_model_name)
|
||||
smart_model = _build_model(settings.smart_model_name)
|
||||
validate_structured_output_support(fast_model, settings.fast_model_name)
|
||||
validate_structured_output_support(smart_model, settings.smart_model_name)
|
||||
|
||||
rag_service, rag_capability = _build_rag(settings)
|
||||
documents, rag_capability = _build_documents(settings)
|
||||
|
||||
return AppRuntime(
|
||||
settings=settings,
|
||||
fast_model=fast_model,
|
||||
smart_model=smart_model,
|
||||
rag_service=rag_service,
|
||||
documents=documents,
|
||||
rag_capability=rag_capability,
|
||||
)
|
||||
|
||||
|
||||
def _build_model(model_name: str) -> Model:
|
||||
"""Construct a model, injecting our keepalive-free httpx client for
|
||||
Anthropic models so workers don't pick up stale pooled connections.
|
||||
|
||||
Other providers fall back to ``infer_model`` defaults; the stale-pool
|
||||
issue is specific to the Cloudflare-fronted Anthropic API in our
|
||||
observations and the fix doesn't necessarily apply elsewhere.
|
||||
"""
|
||||
if model_name.startswith("anthropic:"):
|
||||
bare_name = model_name.removeprefix("anthropic:")
|
||||
provider = AnthropicProvider(http_client=_build_anthropic_http_client())
|
||||
return AnthropicModel(bare_name, provider=provider)
|
||||
return infer_model(model_name)
|
||||
|
||||
Reference in New Issue
Block a user