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

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

The new full doc analysis stuff is heavier than the RAG lookup so both
remain. The agents should use RAG for targeted info and the chunked
reasoner for info that requires reading the full doc.
This commit is contained in:
James Brunton
2026-05-14 13:19:38 +00:00
committed by GitHub
parent 8abe734f0b
commit 672e81d286
55 changed files with 3327 additions and 578 deletions
@@ -30,6 +30,7 @@ import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.job.ResultFile;
import stirling.software.common.service.JobOwnershipService;
import stirling.software.common.service.TaskManager;
import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
@@ -140,13 +141,32 @@ public class AiEngineController {
}
private void runOrchestrationStream(AiWorkflowRequest request, SseEmitter emitter) {
AiWorkflowService.ProgressListener listener =
new AiWorkflowService.ProgressListener() {
@Override
public void onProgress(AiWorkflowProgressEvent event) {
sendEvent(emitter, "progress", event);
}
@Override
public void onHeartbeat() {
// Forward upstream heartbeats so the SSE pipe stays visibly alive between
// real progress events; if the frontend has gone away, sendEvent throws,
// which propagates up through the stream consumer and closes our upstream
// engine connection so the engine can cancel its in-flight workflow.
sendEvent(emitter, "heartbeat", Map.of());
}
};
try {
AiWorkflowResponse result =
aiWorkflowService.orchestrate(
request, progress -> sendEvent(emitter, "progress", progress));
AiWorkflowResponse result = aiWorkflowService.orchestrate(request, listener);
registerFileResultAsJob(result);
sendEvent(emitter, "result", result);
emitter.complete();
} catch (ClientDisconnectedException e) {
// The frontend gave up mid-stream. The exception unwinding through orchestrate()
// already closed the upstream engine connection (engine sees disconnect and cancels).
// The emitter is already toast; nothing useful left to send.
log.debug("Client disconnected mid-stream; aborting workflow", e);
} catch (Exception e) {
log.error("AI orchestration stream failed", e);
// Emit an error frame for the frontend and then complete normally. Using
@@ -192,7 +212,21 @@ public class AiEngineController {
try {
emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
} catch (IOException e) {
log.debug("Failed to send SSE event (client may have disconnected)", e);
// Surface the disconnect so the streaming pipeline unwinds: callers higher up close
// the upstream engine connection, which lets the engine cancel its in-flight workflow.
// Without this, the engine would keep producing (and billing for) tokens whose results
// nobody is reading.
throw new ClientDisconnectedException("Client disconnected from SSE stream", e);
}
}
/**
* Thrown by {@link #sendEvent} when the SSE emitter's underlying connection is gone. Treated as
* a signal to abort the workflow, not as an error to report.
*/
private static final class ClientDisconnectedException extends RuntimeException {
ClientDisconnectedException(String message, Throwable cause) {
super(message, cause);
}
}
@@ -0,0 +1,24 @@
package stirling.software.proprietary.model.api.ai;
import java.util.List;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
/**
* Body for {@code POST /api/v1/documents} on the AI engine. Sent by Java when the engine reports
* {@code need_ingest} and the requested document's extracted content must be stored before the
* workflow can continue.
*/
@Data
@NoArgsConstructor
@AllArgsConstructor
public class AiDocumentIngestRequest {
private String documentId;
private String source;
private List<AiPageText> pageText;
}
@@ -0,0 +1,59 @@
package stirling.software.proprietary.model.api.ai;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
/**
* Typed engine-emitted progress detail, mirroring the Python {@code ProgressEvent} discriminated
* union (see {@code engine/src/stirling/contracts/progress.py}). Carried inside {@link
* AiWorkflowProgressEvent#getEngineDetail()} for {@link AiWorkflowPhase#ENGINE_PROGRESS} events.
*
* <p>Sealed so adding a new engine-side phase forces a matching subtype on the Java side instead of
* silently passing through as an opaque map. The {@code phase} string is the discriminator and
* stays on the wire so the frontend (which doesn't know about Java's class hierarchy) can switch on
* it.
*/
@JsonTypeInfo(
use = JsonTypeInfo.Id.NAME,
include = JsonTypeInfo.As.EXISTING_PROPERTY,
property = "phase",
visible = true)
@JsonSubTypes({
@JsonSubTypes.Type(
value = AiEngineProgressDetail.WholeDocReadStarted.class,
name = "whole_doc_read_started"),
@JsonSubTypes.Type(
value = AiEngineProgressDetail.WholeDocSliceDone.class,
name = "whole_doc_slice_done"),
@JsonSubTypes.Type(
value = AiEngineProgressDetail.WholeDocCompressionRound.class,
name = "whole_doc_compression_round"),
@JsonSubTypes.Type(
value = AiEngineProgressDetail.WholeDocReadDone.class,
name = "whole_doc_read_done"),
})
@JsonIgnoreProperties(ignoreUnknown = true)
public sealed interface AiEngineProgressDetail {
String phase();
record WholeDocReadStarted(String phase, String question, int pages, int slices)
implements AiEngineProgressDetail {}
record WholeDocSliceDone(
String phase,
int completed,
int total,
String pages,
int durationMs,
int excerpts,
int facts)
implements AiEngineProgressDetail {}
record WholeDocCompressionRound(String phase, int roundNumber, int notesIn, int groups)
implements AiEngineProgressDetail {}
record WholeDocReadDone(String phase, int completed, int slices, double durationSeconds)
implements AiEngineProgressDetail {}
}
@@ -4,11 +4,11 @@ import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
/** A single page of extracted text for RAG ingest requests. */
/** A single page of extracted text for document ingest requests. */
@Data
@NoArgsConstructor
@AllArgsConstructor
public class AiRagPageText {
public class AiPageText {
private int pageNumber;
@@ -1,24 +0,0 @@
package stirling.software.proprietary.model.api.ai;
import java.util.List;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
/**
* Body for {@code POST /api/v1/rag/documents} on the AI engine. Sent by Java when the engine
* reports {@code need_ingest} and the requested document's extracted content must be stored before
* the workflow can continue.
*/
@Data
@NoArgsConstructor
@AllArgsConstructor
public class AiRagIngestRequest {
private String documentId;
private String source;
private List<AiRagPageText> pageText;
}
@@ -9,7 +9,13 @@ public enum AiWorkflowPhase {
CALLING_ENGINE("calling_engine"),
EXTRACTING_CONTENT("extracting_content"),
EXECUTING_TOOL("executing_tool"),
PROCESSING("processing");
PROCESSING("processing"),
/**
* Generic engine-emitted progress event (e.g. chunked-reasoner slice progress). The original
* engine event JSON is carried in {@link AiWorkflowProgressEvent#getEngineDetail()}, including
* a specific {@code phase} string the frontend can switch on.
*/
ENGINE_PROGRESS("engine_progress");
private final String value;
@@ -23,8 +23,17 @@ public class AiWorkflowProgressEvent {
/** Total number of plan steps, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
private Integer stepCount;
/**
* Engine-emitted event payload, for {@link AiWorkflowPhase#ENGINE_PROGRESS} events. The payload
* is a typed subtype keyed on its {@code phase} string (e.g. {@code "whole_doc_slice_done"})
* carrying phase-specific fields (slice index, page range, durations, etc.) that the frontend
* can render as detailed progress.
*/
private AiEngineProgressDetail engineDetail;
public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
return new AiWorkflowProgressEvent(phase, System.currentTimeMillis(), null, null, null);
return new AiWorkflowProgressEvent(
phase, System.currentTimeMillis(), null, null, null, null);
}
public static AiWorkflowProgressEvent executingTool(String tool, int stepIndex, int stepCount) {
@@ -33,6 +42,17 @@ public class AiWorkflowProgressEvent {
System.currentTimeMillis(),
tool,
stepIndex,
stepCount);
stepCount,
null);
}
public static AiWorkflowProgressEvent engineProgress(AiEngineProgressDetail detail) {
return new AiWorkflowProgressEvent(
AiWorkflowPhase.ENGINE_PROGRESS,
System.currentTimeMillis(),
null,
null,
null,
detail);
}
}
@@ -7,6 +7,8 @@ import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpTimeoutException;
import java.time.Duration;
import java.util.function.Consumer;
import java.util.stream.Stream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
@@ -83,6 +85,70 @@ public class AiEngineClient {
return response.body();
}
/**
* POST a JSON body and consume the response as a stream of NDJSON lines. Each line is passed to
* {@code lineConsumer} in arrival order; the call returns when the engine closes the stream.
*
* <p>This is the right shape for long-running orchestrator calls that emit incremental
* progress. The total HTTP timeout is the long-running timeout (typically 600s+), but in
* practice line arrival keeps the connection logically alive: as long as the engine emits
* events, the work is progressing. Genuine engine hangs still hit the total timeout.
*/
public void streamPost(String path, String jsonBody, Consumer<String> lineConsumer)
throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
throw new ResponseStatusException(
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
}
String url = config.getUrl().stripTrailing() + path;
Duration timeout = Duration.ofSeconds(config.getLongRunningTimeoutSeconds());
log.debug(
"Proxying AI engine streaming request to {} (timeout {}s)",
url,
timeout.toSeconds());
HttpRequest request =
HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Content-Type", "application/json")
.header("Accept", "application/x-ndjson")
.timeout(timeout)
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
.build();
HttpResponse<Stream<String>> response;
try {
response = httpClient.send(request, HttpResponse.BodyHandlers.ofLines());
} catch (HttpTimeoutException e) {
throw new ResponseStatusException(HttpStatus.GATEWAY_TIMEOUT, "AI engine timed out", e);
} catch (IOException e) {
throw new ResponseStatusException(
HttpStatus.SERVICE_UNAVAILABLE, "AI engine unreachable: " + e.getMessage(), e);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new ResponseStatusException(
HttpStatus.SERVICE_UNAVAILABLE, "AI engine request was interrupted");
}
int status = response.statusCode();
if (status >= 400) {
throw new ResponseStatusException(
HttpStatus.valueOf(status >= 500 ? 502 : status),
"AI engine returned error: " + status);
}
try (Stream<String> lines = response.body()) {
lines.forEach(
line -> {
if (!line.isEmpty()) {
lineConsumer.accept(line);
}
});
}
}
public String get(String path) throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
@@ -37,9 +37,10 @@ import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.ZipExtractionUtils;
import stirling.software.proprietary.model.api.ai.AiConversationMessage;
import stirling.software.proprietary.model.api.ai.AiDocumentIngestRequest;
import stirling.software.proprietary.model.api.ai.AiEngineProgressDetail;
import stirling.software.proprietary.model.api.ai.AiFile;
import stirling.software.proprietary.model.api.ai.AiRagIngestRequest;
import stirling.software.proprietary.model.api.ai.AiRagPageText;
import stirling.software.proprietary.model.api.ai.AiPageText;
import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
@@ -61,7 +62,7 @@ import tools.jackson.databind.ObjectMapper;
@RequiredArgsConstructor
public class AiWorkflowService {
private static final String RAG_DOCUMENTS_ENDPOINT = "/api/v1/rag/documents";
private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents";
private final CustomPDFDocumentFactory pdfDocumentFactory;
private final AiEngineClient aiEngineClient;
@@ -77,6 +78,14 @@ public class AiWorkflowService {
@FunctionalInterface
public interface ProgressListener {
void onProgress(AiWorkflowProgressEvent event);
/**
* Called when the engine emits a keep-alive heartbeat. Default is a no-op; consumers that
* forward to a downstream connection (e.g. an SSE emitter) override this to push a
* heartbeat through, so the next downstream-disconnect surfaces immediately rather than
* waiting for the next real progress event.
*/
default void onHeartbeat() {}
}
private static final ProgressListener NOOP_LISTENER = event -> {};
@@ -144,7 +153,7 @@ public class AiWorkflowService {
ProgressListener listener)
throws IOException {
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
AiWorkflowResponse response = invokeOrchestrator(request);
AiWorkflowResponse response = invokeOrchestrator(request, listener);
return switch (response.getOutcome()) {
case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
@@ -278,22 +287,22 @@ public class AiWorkflowService {
}
private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException {
List<AiRagPageText> pages = new ArrayList<>();
List<AiPageText> pages = new ArrayList<>();
try (PDDocument document = pdfDocumentFactory.load(multipartFile, true)) {
int pageCount = document.getNumberOfPages();
for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
String pageText = pdfContentExtractor.extractPageTextRaw(document, pageNumber);
if (pageText != null && !pageText.isBlank()) {
pages.add(new AiRagPageText(pageNumber, pageText));
pages.add(new AiPageText(pageNumber, pageText));
}
}
}
AiRagIngestRequest ingestRequest =
new AiRagIngestRequest(file.getId(), file.getName(), pages);
AiDocumentIngestRequest ingestRequest =
new AiDocumentIngestRequest(file.getId(), file.getName(), pages);
String body = objectMapper.writeValueAsString(ingestRequest);
aiEngineClient.postLongRunning(RAG_DOCUMENTS_ENDPOINT, body);
aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body);
log.debug(
"Ingested file into RAG: id={}, name={}, pages={}",
"Ingested document: id={}, name={}, pages={}",
file.getId(),
file.getName(),
pages.size());
@@ -651,10 +660,57 @@ public class AiWorkflowService {
return response;
}
private AiWorkflowResponse invokeOrchestrator(WorkflowTurnRequest request) throws IOException {
/**
* Drive the engine's streaming orchestrator endpoint. Progress events are forwarded to {@code
* listener} as they arrive (each one keeps the SSE connection to the frontend alive too). The
* final {@code result} event carries the full {@link AiWorkflowResponse}; an {@code error}
* event surfaces engine-side failures.
*/
private AiWorkflowResponse invokeOrchestrator(
WorkflowTurnRequest request, ProgressListener listener) throws IOException {
String requestBody = objectMapper.writeValueAsString(request);
String responseBody = aiEngineClient.post("/api/v1/orchestrator", requestBody);
return objectMapper.readValue(responseBody, AiWorkflowResponse.class);
AiWorkflowResponse[] resultHolder = new AiWorkflowResponse[1];
String[] errorHolder = new String[1];
aiEngineClient.streamPost(
"/api/v1/orchestrator",
requestBody,
line -> handleStreamLine(line, listener, resultHolder, errorHolder));
if (errorHolder[0] != null) {
throw new IOException("AI engine returned error: " + errorHolder[0]);
}
if (resultHolder[0] == null) {
throw new IOException("AI engine stream ended without a result");
}
return resultHolder[0];
}
private void handleStreamLine(
String line,
ProgressListener listener,
AiWorkflowResponse[] resultHolder,
String[] errorHolder) {
try {
JsonNode node = objectMapper.readTree(line);
String event = node.path("event").asText();
switch (event) {
case "progress" -> {
AiEngineProgressDetail detail =
objectMapper.treeToValue(node, AiEngineProgressDetail.class);
listener.onProgress(AiWorkflowProgressEvent.engineProgress(detail));
}
case "result" -> {
JsonNode response = node.path("response");
resultHolder[0] = objectMapper.treeToValue(response, AiWorkflowResponse.class);
}
case "error" -> errorHolder[0] = node.path("message").asText("unknown error");
case "heartbeat" -> listener.onHeartbeat();
default -> log.warn("Ignoring unknown engine stream event: {}", event);
}
} catch (JacksonException e) {
log.warn("Failed to parse engine stream line: {}", line, e);
}
}
@Data
@@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
@@ -22,6 +23,7 @@ import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
@@ -419,7 +421,7 @@ class AiWorkflowServiceTest {
}
@Test
void needIngestExtractsPageTextAndPostsToRagThenRetries() throws IOException {
void needIngestExtractsPageTextAndPostsThenRetries() throws IOException {
MockMultipartFile input = pdf("report.pdf", "bytes");
when(fileIdStrategy.idFor(any())).thenReturn("report-id");
@@ -431,37 +433,64 @@ class AiWorkflowServiceTest {
.thenReturn("page content");
int[] orchestratorCalls = {0};
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString()))
.thenAnswer(
doAnswer(
inv -> {
orchestratorCalls[0]++;
String responseJson;
if (orchestratorCalls[0] == 1) {
return """
{
"outcome":"need_ingest",
"resumeWith":"pdf_question",
"reason":"ingest first",
"filesToIngest":[{"id":"report-id","name":"report.pdf"}],
"contentTypes":["page_text"]
}
""";
responseJson =
"""
{
"outcome":"need_ingest",
"resumeWith":"pdf_question",
"reason":"ingest first",
"filesToIngest":[{"id":"report-id","name":"report.pdf"}],
"contentTypes":["page_text"]
}
""";
} else {
responseJson =
"""
{"outcome":"answer","answer":"done","evidence":[]}
""";
}
return """
{"outcome":"answer","answer":"done","evidence":[]}
""";
});
Consumer<String> consumer = inv.getArgument(2);
consumer.accept(wrapAsResultEvent(responseJson));
return null;
})
.when(aiEngineClient)
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/rag/documents"), anyString());
verify(aiEngineClient, times(2)).post(eq("/api/v1/orchestrator"), anyString());
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString());
verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any());
}
// --- helpers ---
private void stubOrchestrator(String responseJson) throws IOException {
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString())).thenReturn(responseJson);
doAnswer(
inv -> {
Consumer<String> consumer = inv.getArgument(2);
consumer.accept(wrapAsResultEvent(responseJson));
return null;
})
.when(aiEngineClient)
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
}
/**
* Wrap a unary orchestrator response JSON as the NDJSON "result" event the streaming endpoint
* emits, so existing tests can keep their per-outcome JSON literals.
*/
private String wrapAsResultEvent(String responseJson) throws IOException {
return objectMapper
.createObjectNode()
.put("event", "result")
.set("response", objectMapper.readTree(responseJson))
.toString();
}
private void stubEndpoint(String endpoint, Resource body) {
+21
View File
@@ -37,6 +37,22 @@ STIRLING_RAG_TOP_K=20
# rather than chain more searches.
STIRLING_RAG_MAX_SEARCHES=5
# Chunked reasoner settings: how big each per-worker slice is (in characters),
# how many workers may run in parallel against the fast model, and how long
# any single worker is allowed to wait for a response before being abandoned.
# Worker timeouts protect gather_notes from upstream model stalls (which
# otherwise hang at the provider's ~10 minute HTTP default); the affected
# slice is dropped and the rest of the document still answers.
STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE=16000
STIRLING_CHUNKED_REASONER_CONCURRENCY=10
STIRLING_CHUNKED_REASONER_WORKER_TIMEOUT_SECONDS=60
# When the rendered slice notes would exceed this many characters, the
# reasoner folds them hierarchically with fast-model calls until they fit.
# This keeps the synthesis prompt under the model's context limit on long
# documents (a 3000-page novel produces ~900k chars of raw notes).
STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET=250000
# Upper bounds on PDF page text the engine will request per extraction round.
STIRLING_MAX_PAGES=200
STIRLING_MAX_CHARACTERS=200000
@@ -51,3 +67,8 @@ STIRLING_LOG_LEVEL=INFO
# Path to log file. Rolls daily, keeps 1 backup. Leave empty for console only.
STIRLING_LOG_FILE=
# Set true to log every outgoing httpx / Anthropic SDK request with timing.
# Use when diagnosing worker stalls: a hung call shows a "Request" line with
# no matching "Response" line. Noisy; leave off in normal use.
STIRLING_HTTP_DEBUG=false
+45 -25
View File
@@ -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)
+3 -3
View File
@@ -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)
+3 -3
View File
@@ -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:
+2 -2
View File
@@ -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)
+156 -5
View File
@@ -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)
-63
View File
@@ -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)
+54 -2
View File
@@ -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
+22 -6
View File
@@ -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",
]
+2 -2
View File
@@ -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
+70
View File
@@ -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"),
]
-39
View File
@@ -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 |
+20
View File
@@ -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
@@ -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()
@@ -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)
+133
View File
@@ -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()
@@ -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:
+113
View File
@@ -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.)."""
+3
View File
@@ -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,
)
-19
View File
@@ -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",
]
-102
View File
@@ -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()
-63
View File
@@ -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.)."""
+10
View File
@@ -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",
]
+47
View File
@@ -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)
+54 -18
View File
@@ -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)
@@ -0,0 +1,511 @@
"""Tests for ``ChunkedReasoner``: chunking, fan-out, and synthesis wiring.
LLM calls are stubbed at the agent boundary; the runtime fixture supplies a
``test`` model so construction succeeds without provider credentials.
"""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from unittest.mock import AsyncMock, patch
import pytest
from pydantic import BaseModel
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
from stirling.contracts import WholeDocSliceDone
from stirling.contracts.documents import Page
from stirling.services import reset_progress_emitter, set_progress_emitter
from stirling.services.runtime import AppRuntime
@dataclass
class _StubAgentResult[T]:
output: T
class _Answer(BaseModel):
answer: str
def _page(n: int, text: str) -> Page:
return Page(page_number=n, text=text, char_count=len(text))
# Slicing logic
class TestSlicePages:
"""``_slice_pages`` is pure: no model calls, no I/O."""
def test_groups_consecutive_pages_under_budget(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=20)
pages = [_page(1, "abc"), _page(2, "defgh"), _page(3, "ij")]
slices = reasoner._slice_pages(pages)
assert len(slices) == 1
assert [p.page_number for p in slices[0]] == [1, 2, 3]
def test_starts_new_slice_when_budget_exceeded(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(1, "a" * 6), _page(2, "b" * 6), _page(3, "c" * 6)]
slices = reasoner._slice_pages(pages)
# Each page is 6 chars, budget 10 -> two pages would be 12 (over),
# so the slicer breaks after each page.
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
def test_oversized_page_becomes_its_own_slice(self, runtime: AppRuntime) -> None:
"""A single page larger than the budget is never split. The reasoner
accepts that this slice exceeds the budget rather than breaking page
boundaries."""
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(1, "small"), _page(2, "x" * 100), _page(3, "tiny")]
slices = reasoner._slice_pages(pages)
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
def test_preserves_input_order(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
pages = [_page(i, f"page-{i}") for i in range(1, 11)]
slices = reasoner._slice_pages(pages)
flattened = [p.page_number for s in slices for p in s]
assert flattened == list(range(1, 11))
# End-to-end orchestration
class TestReason:
@pytest.mark.anyio
async def test_runs_one_chunk_per_slice_and_synthesises(self, runtime: AppRuntime) -> None:
"""Three small pages with a generous budget produce one chunk and one extractor call;
the synthesis stage receives notes from all chunks and returns the final answer."""
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
pages = [_page(1, "alpha"), _page(2, "beta"), _page(3, "gamma")]
canned_notes = ChunkNotes(pages=[1, 2, 3], summary="all three pages", facts=["fact-1"])
canned_answer = _Answer(answer="final answer")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
):
result = await reasoner.reason(
pages=pages,
question="summarise this",
answer_prompt="answer the question from the notes",
answer_type=_Answer,
)
assert result == canned_answer
assert chunk_mock.await_count == 1
synth_mock.assert_awaited_once()
synth_args = synth_mock.await_args
assert synth_args is not None
# _synthesise(question, notes, answer_prompt, answer_type)
_, notes_arg, _, type_arg = synth_args.args
assert notes_arg == [canned_notes]
assert type_arg is _Answer
@pytest.mark.anyio
async def test_fans_out_when_pages_exceed_slice_budget(self, runtime: AppRuntime) -> None:
"""Pages that don't fit into a single slice produce one extractor call per slice."""
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 6)]
canned_notes = ChunkNotes(pages=[0], summary="placeholder")
canned_answer = _Answer(answer="ok")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)),
):
await reasoner.reason(
pages=pages,
question="aggregate",
answer_prompt="answer",
answer_type=_Answer,
)
# 5 pages, budget 10, each page 8 chars -> 5 slices -> 5 chunk calls.
assert chunk_mock.await_count == 5
@pytest.mark.anyio
async def test_skips_first_round_chunks_that_raise_and_continues(self, runtime: AppRuntime) -> None:
"""First-round chunks have no fallback notes, so a failure is dropped
rather than preserving anything; the surviving notes still flow into
synthesis."""
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
good = ChunkNotes(pages=[1], summary="ok")
async_results = [good, RuntimeError("chunk boom"), good]
async def _chunk(*_args: object, **_kwargs: object) -> tuple[ChunkNotes, float]:
value = async_results.pop(0)
if isinstance(value, BaseException):
raise value
return value, 0.0
canned_answer = _Answer(answer="resilient")
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=_chunk)),
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
):
result = await reasoner.reason(
pages=pages,
question="aggregate",
answer_prompt="answer",
answer_type=_Answer,
)
assert result == canned_answer
synth_args = synth_mock.await_args
assert synth_args is not None
_, notes_arg, _, _ = synth_args.args
assert len(notes_arg) == 2
@pytest.mark.anyio
async def test_raises_when_every_first_round_chunk_fails(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 3)]
with (
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=RuntimeError("boom"))),
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
pytest.raises(RuntimeError, match="no notes"),
):
await reasoner.reason(
pages=pages,
question="anything",
answer_prompt="answer",
answer_type=_Answer,
)
synth_mock.assert_not_awaited()
@pytest.mark.anyio
async def test_rejects_empty_pages(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime)
with pytest.raises(ValueError, match="at least one page"):
await reasoner.reason(
pages=[],
question="x",
answer_prompt="y",
answer_type=_Answer,
)
@pytest.mark.anyio
async def test_progress_events_carry_monotonic_completion_counter(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
# Each chunk's worker awaits a different release event; we release them in
# reverse order so completion order is the inverse of slice order.
release_events = [asyncio.Event() for _ in pages]
next_call_index = 0
async def _gated_worker(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
nonlocal next_call_index
this_call = next_call_index
next_call_index += 1
await release_events[this_call].wait()
return _StubAgentResult(output=ChunkNotes(pages=[this_call + 1], summary=f"slice-{this_call}"))
emitted: list[WholeDocSliceDone] = []
async def _capture_emitter(event: object) -> None:
if isinstance(event, WholeDocSliceDone):
emitted.append(event)
async def _release_in_reverse() -> None:
# Wait briefly so all three worker tasks are blocked on their events
# before we start releasing them.
await asyncio.sleep(0)
for ev in reversed(release_events):
ev.set()
# Yield so the just-released worker can run to completion before
# we release the next one — keeps ordering deterministic.
await asyncio.sleep(0)
await asyncio.sleep(0)
token = set_progress_emitter(_capture_emitter)
try:
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_gated_worker)):
gather_task = asyncio.create_task(reasoner.gather_notes(pages, "anything"))
await _release_in_reverse()
notes = await gather_task
finally:
reset_progress_emitter(token)
assert len(notes) == 3
assert [event.completed for event in emitted] == [1, 2, 3]
assert all(event.total == 3 for event in emitted)
@pytest.mark.anyio
async def test_worker_timeout_is_terminal_for_the_chunk(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
pages = [_page(1, "x" * 8), _page(2, "y" * 8)]
attempts = 0
async def _hang_forever(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
nonlocal attempts
attempts += 1
await asyncio.sleep(10)
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
with (
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang_forever)),
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
pytest.raises(RuntimeError, match="no notes"),
):
await reasoner.reason(
pages=pages,
question="anything",
answer_prompt="answer",
answer_type=_Answer,
)
# One attempt per slice; no retry path.
assert attempts == len(pages)
synth_mock.assert_not_awaited()
@pytest.mark.anyio
async def test_worker_timeout_drops_stalled_chunks(self, runtime: AppRuntime) -> None:
"""A worker that exceeds ``worker_timeout_seconds`` is abandoned, not awaited.
Without this guard one stuck upstream call would pin gather_notes to its
provider HTTP timeout (~10 minutes), starving the orchestrator request.
"""
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
pages = [_page(i, "x" * 8) for i in range(1, 4)]
async def _hang(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
await asyncio.sleep(10)
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
with (
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang)),
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
pytest.raises(RuntimeError, match="no notes"),
):
await reasoner.reason(
pages=pages,
question="anything",
answer_prompt="answer",
answer_type=_Answer,
)
synth_mock.assert_not_awaited()
# Prompt construction
class TestPromptConstruction:
def test_extraction_prompt_includes_question_and_page_markers(self, runtime: AppRuntime) -> None:
"""A first-round chunk's content carries ``[Page N]`` markers; the
extraction prompt prepends the user question."""
reasoner = ChunkedReasoner(runtime)
chunk = reasoner._chunk_from_pages([_page(2, "page two body"), _page(3, "page three body")])
prompt = reasoner._build_extraction_prompt(chunk.content, "what is on page two?")
assert "what is on page two?" in prompt
assert "[Page 2]" in prompt
assert "[Page 3]" in prompt
assert "page two body" in prompt
def test_format_notes_groups_by_page_label(self) -> None:
notes = [
ChunkNotes(pages=[1], summary="single", facts=["f-1"]),
ChunkNotes(pages=[2, 3, 4], summary="range", relevant_excerpts=["quote-1"]),
]
rendered = ChunkedReasoner.format_notes(notes)
assert "[Notes from page 1]" in rendered
assert "[Notes from pages 2-4]" in rendered
assert "f-1" in rendered
assert "quote-1" in rendered
# Hierarchical compression
#
# The compression loop is part of ``_run_chunks`` and isn't exposed
# directly, so these tests drive it end-to-end via ``gather_notes`` with a
# stubbed extractor that controls per-call output (and per-call failure
# patterns) by counting calls.
class TestCompression:
@pytest.mark.anyio
async def test_no_compression_when_under_budget(self, runtime: AppRuntime) -> None:
"""First-round notes that already fit the budget result in zero
compression rounds: the only extractor calls are one per slice."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=10_000)
pages = [_page(i, "x" * 150) for i in range(1, 5)] # each page exceeds slice budget alone -> 4 slices
canned = _ExtractedNotes(summary="ok")
with patch.object(
reasoner._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
) as ext_mock:
notes = await reasoner.gather_notes(pages, "anything")
assert ext_mock.await_count == 4
assert len(notes) == 4
@pytest.mark.anyio
async def test_runs_compression_when_over_budget(self, runtime: AppRuntime) -> None:
"""When first-round notes overflow the budget, the loop regroups them
and runs the extractor again. Output is shorter than input; pages from
every input slice survive in the consolidated notes."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=200)
pages = [_page(i, "x" * 150) for i in range(1, 5)]
call_count = 0
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
nonlocal call_count
call_count += 1
if call_count <= 4:
# Round 1: each note ~80 chars rendered. 4 * 80 = 320 chars, over the 200 budget.
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
# Round 2: smaller note so the post-round set fits the budget.
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)) as ext_mock:
notes = await reasoner.gather_notes(pages, "anything")
# 4 first-round + 2 compression-round calls = 6 total.
assert ext_mock.await_count == 6
# Compressed from 4 notes to 2.
assert len(notes) == 2
# Pages from every original slice are preserved through compression.
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
@pytest.mark.anyio
async def test_compression_preserves_input_notes_when_a_group_fails(self, runtime: AppRuntime) -> None:
"""A compression chunk that raises has its input notes carried forward
rather than dropped, so page coverage isn't silently lost. The
succeeding chunk is replaced by its consolidated note.
Budget is sized so the post-round survivors (2 preserved + 1
consolidated) fit, leaving a single compression round as the
observable interaction."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=300)
pages = [_page(i, "x" * 150) for i in range(1, 5)]
call_count = 0
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
nonlocal call_count
call_count += 1
if call_count <= 4:
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
if call_count == 5:
# The first compression call (covering 2 of the round-1 notes) fails.
raise RuntimeError("compression group fails")
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
notes = await reasoner.gather_notes(pages, "anything")
# 2 preserved round-1 notes + 1 consolidated note = 3 notes total. Pages
# from every original slice are still covered (preservation worked).
assert len(notes) == 3
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
consolidated = [n for n in notes if n.summary == "ok"]
preserved = [n for n in notes if n.summary.startswith("x")]
assert len(consolidated) == 1
assert len(preserved) == 2
@pytest.mark.anyio
async def test_compression_bails_when_every_group_fails(self, runtime: AppRuntime) -> None:
"""If every chunk in a compression round fails, every input note is
preserved (none consolidated). The loop exits rather than retrying
the same shape forever."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=200)
pages = [_page(i, "x" * 150) for i in range(1, 5)]
call_count = 0
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
nonlocal call_count
call_count += 1
if call_count <= 4:
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
raise RuntimeError("compression always fails")
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
notes = await reasoner.gather_notes(pages, "anything")
# All 4 round-1 notes preserved through the bailed compression round.
assert len(notes) == 4
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
class TestGroupNotesForCompression:
"""``_group_notes_for_compression`` is pure and packs by rendered char count."""
def test_packs_consecutive_notes_under_budget(self, runtime: AppRuntime) -> None:
reasoner = ChunkedReasoner(runtime, chars_per_slice=10_000)
notes = [ChunkNotes(pages=[i], summary=f"s-{i}") for i in range(1, 5)]
groups = reasoner._group_notes_for_compression(notes)
assert len(groups) == 1
assert [n.pages[0] for n in groups[0]] == [1, 2, 3, 4]
def test_starts_new_group_when_budget_exceeded(self, runtime: AppRuntime) -> None:
"""Each note already exceeds the per-group budget, so each becomes its
own group; this matches how slice-pages handles oversize pages."""
reasoner = ChunkedReasoner(runtime, chars_per_slice=5)
notes = [ChunkNotes(pages=[i], summary=f"slice-{i}-with-prose-to-fill-space") for i in range(1, 5)]
groups = reasoner._group_notes_for_compression(notes)
assert [[n.pages[0] for n in g] for g in groups] == [[1], [2], [3], [4]]
class TestExtractChunk:
@pytest.mark.anyio
async def test_pages_are_unioned_for_compression_chunks(self, runtime: AppRuntime) -> None:
"""A compression chunk's resulting note carries the union of input pages.
The model output schema doesn't include pages, so the wrapper is the
single source of truth."""
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
reasoner = ChunkedReasoner(runtime)
group = [
ChunkNotes(pages=[1, 2], summary="a"),
ChunkNotes(pages=[3], summary="b"),
ChunkNotes(pages=[4, 5], summary="c"),
]
chunk = reasoner._chunk_from_notes(group)
canned = _ExtractedNotes(summary="merged", facts=["x"], relevant_excerpts=["y"])
with patch.object(
reasoner._extractor,
"run",
AsyncMock(return_value=_StubAgentResult(output=canned)),
):
note, _ = await reasoner._extract_chunk(chunk, "anything")
assert note.pages == [1, 2, 3, 4, 5]
assert note.summary == "merged"
assert note.facts == ["x"]
assert note.relevant_excerpts == ["y"]
@@ -0,0 +1,220 @@
"""Tests for ``WholeDocReaderCapability``: tool dispatch, multi-file iteration,
budget enforcement, and graceful handling of missing pages.
The map-phase LLM call is patched at the reasoner boundary so tests don't hit
any model.
"""
from __future__ import annotations
from dataclasses import replace
from unittest.mock import AsyncMock, patch
import pytest
from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability
from stirling.contracts import AiFile, PageText
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.services.runtime import AppRuntime
class StubEmbedder:
"""Deterministic embeddings so tests don't need a real provider."""
def __init__(self, dim: int = 8) -> None:
self._dim = dim
async def embed_query(self, text: str) -> list[float]:
h = hash(text) % 1000
return [(h + i) / 1000.0 for i in range(self._dim)]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
return [await self.embed_query(t) for t in texts]
def chunk_and_prepare(
self,
text: str,
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs: list[Document] = []
for i, chunk in enumerate(chunks):
meta = dict(base_metadata) if base_metadata else {}
meta["source"] = source
meta["chunk_index"] = str(i)
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
return docs
@pytest.fixture
def runtime_with_stub_docs(runtime: AppRuntime) -> AppRuntime:
stub = DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=runtime.settings.rag_default_top_k,
)
return replace(runtime, documents=stub)
def _ai_file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name)
@pytest.mark.anyio
async def test_read_full_document_returns_formatted_notes_for_single_file(
runtime_with_stub_docs: AppRuntime,
) -> None:
"""The tool reads the file's stored pages, calls the reasoner's map phase,
and returns the formatted notes prefixed by the file name."""
pages = [
PageText(page_number=1, text="Chapter one prose."),
PageText(page_number=2, text="Chapter two prose."),
]
await runtime_with_stub_docs.documents.ingest(FileId("doc-id"), pages, source="doc.pdf")
reasoner = ChunkedReasoner(runtime_with_stub_docs)
canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])]
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned_notes)) as gather_mock:
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-id", "doc.pdf")],
reasoner=reasoner,
)
result = await capability._read_full_document("what is in the document?")
gather_mock.assert_awaited_once()
call = gather_mock.await_args
assert call is not None
pages_arg = call.args[0]
assert [p.page_number for p in pages_arg] == [1, 2]
assert "=== doc.pdf ===" in result
assert "fact-A" in result
assert "[Notes from pages 1-2]" in result
@pytest.mark.anyio
async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs: AppRuntime) -> None:
"""Multi-file requests run the map phase per file and return one section
per file in the formatted output."""
for cid, source in (("doc-a", "a.pdf"), ("doc-b", "b.pdf")):
await runtime_with_stub_docs.documents.ingest(
FileId(cid),
[PageText(page_number=1, text=f"contents of {cid}")],
source=source,
)
reasoner = ChunkedReasoner(runtime_with_stub_docs)
notes_by_call = [
[ChunkNotes(pages=[1], summary="a-summary")],
[ChunkNotes(pages=[1], summary="b-summary")],
]
async def _gather(*_args: object, **_kwargs: object) -> list[ChunkNotes]:
return notes_by_call.pop(0)
with patch.object(reasoner, "gather_notes", AsyncMock(side_effect=_gather)) as gather_mock:
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")],
reasoner=reasoner,
)
result = await capability._read_full_document("compare them")
assert gather_mock.await_count == 2
assert "=== a.pdf ===" in result
assert "a-summary" in result
assert "=== b.pdf ===" in result
assert "b-summary" in result
@pytest.mark.anyio
async def test_read_full_document_skips_files_without_pages(runtime_with_stub_docs: AppRuntime) -> None:
"""Files with no stored pages are quietly skipped; the tool still runs
the map phase for files that do have pages."""
await runtime_with_stub_docs.documents.ingest(
FileId("present"),
[PageText(page_number=1, text="real content")],
source="present.pdf",
)
# 'missing' is never ingested -> read_pages returns [].
reasoner = ChunkedReasoner(runtime_with_stub_docs)
canned = [ChunkNotes(pages=[1], summary="present summary")]
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned)) as gather_mock:
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")],
reasoner=reasoner,
)
result = await capability._read_full_document("anything")
gather_mock.assert_awaited_once()
assert "=== present.pdf ===" in result
assert "missing.pdf" not in result
@pytest.mark.anyio
async def test_read_full_document_returns_empty_message_when_no_pages_anywhere(
runtime_with_stub_docs: AppRuntime,
) -> None:
reasoner = ChunkedReasoner(runtime_with_stub_docs)
with patch.object(reasoner, "gather_notes", AsyncMock()) as gather_mock:
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("nope", "nope.pdf")],
reasoner=reasoner,
)
result = await capability._read_full_document("anything")
gather_mock.assert_not_awaited()
assert result == "Could not read any document content."
@pytest.mark.anyio
async def test_read_full_document_budget_hides_tool_when_exhausted(
runtime_with_stub_docs: AppRuntime,
) -> None:
"""The prepare callback returns None once max_reads is reached so the
agent can no longer call the tool on subsequent turns. Mirrors
RagCapability's per-run budget."""
await runtime_with_stub_docs.documents.ingest(
FileId("doc-id"),
[PageText(page_number=1, text="content")],
source="doc.pdf",
)
reasoner = ChunkedReasoner(runtime_with_stub_docs)
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])):
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-id", "doc.pdf")],
reasoner=reasoner,
max_reads=1,
)
sentinel: object = object()
# Budget intact -> prepare returns the tool.
assert await capability._prepare_read_full_document(None, sentinel) is sentinel # type: ignore[arg-type]
# Spend the budget.
await capability._read_full_document("anything")
# Budget spent -> prepare returns None.
assert await capability._prepare_read_full_document(None, sentinel) is None # type: ignore[arg-type]
@pytest.mark.anyio
async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRuntime) -> None:
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")],
)
text = capability.instructions
assert "alpha.pdf" in text
assert "beta.pdf" in text
assert "read_full_document" in text
+4
View File
@@ -31,6 +31,10 @@ def build_app_settings() -> AppSettings:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_notes_char_budget=250_000,
chunked_reasoner_worker_timeout_seconds=60.0,
max_pages=200,
max_characters=200_000,
posthog_enabled=False,
+4
View File
@@ -43,6 +43,10 @@ class StubSettingsProvider:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_worker_timeout_seconds=60.0,
chunked_reasoner_notes_char_budget=250_000,
max_pages=100,
max_characters=100_000,
posthog_enabled=False,
@@ -2,14 +2,15 @@ from __future__ import annotations
import pytest
from stirling.contracts import PageText
from stirling.documents.chunker import chunk_text
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, SearchResult
from stirling.models import FileId
from stirling.rag.capability import RagCapability
from stirling.rag.chunker import chunk_text
from stirling.rag.service import RagService
from stirling.rag.sqlite_vec_store import SqliteVecStore
from stirling.rag.store import Document, SearchResult
# ── chunk_text ──────────────────────────────────────────────────────────
# chunk_text
class TestChunkText:
@@ -26,7 +27,6 @@ class TestChunkText:
def test_splits_on_paragraph_boundaries(self) -> None:
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
chunks = chunk_text(text, chunk_size=30, overlap=0)
# Each paragraph fits in 30 chars, so they should be split
assert len(chunks) >= 2
assert "First paragraph." in chunks[0]
@@ -35,22 +35,19 @@ class TestChunkText:
chunks = chunk_text(text, chunk_size=100, overlap=10)
assert len(chunks) > 1
for chunk in chunks:
# Chunks may slightly exceed due to sentence boundary snapping
assert len(chunk) <= 200 # generous upper bound
assert len(chunk) <= 200
def test_overlap_produces_shared_content(self) -> None:
sentences = [f"Sentence number {i}." for i in range(20)]
text = " ".join(sentences)
chunks = chunk_text(text, chunk_size=100, overlap=30)
if len(chunks) >= 2:
# After word-boundary snapping, the second chunk should share
# some content with the tail of the first chunk
words_in_first_tail = chunks[0].split()[-3:] # last 3 words
words_in_first_tail = chunks[0].split()[-3:]
overlap_text = " ".join(words_in_first_tail)
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
# ── SqliteVecStore ──────────────────────────────────────────────────────
# SqliteVecStore
class TestSqliteVecStore:
@@ -59,12 +56,12 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("test-col", "test.pdf")
docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
]
# Simple 3-dimensional embeddings for testing
embeddings = [
[1.0, 0.0, 0.0],
[0.9, 0.1, 0.0],
@@ -72,17 +69,16 @@ class TestSqliteVecStore:
]
await store.add_documents("test-col", docs, embeddings)
# Search with a query close to the programming-related docs
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
assert len(results) == 2
assert isinstance(results[0], SearchResult)
# The closest should be doc "1" (exact match on first dimension)
assert results[0].document.id == "1"
assert results[0].score > 0.5
@pytest.mark.anyio
async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("my-collection", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
@@ -94,6 +90,7 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("to-delete", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("to-delete", docs, [[1.0]])
@@ -104,6 +101,7 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("empty-test", "test.pdf")
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
@@ -117,7 +115,7 @@ class TestSqliteVecStore:
await store.add_documents("bad", docs, [[1.0], [2.0]])
# ── RagService (with stub embedder) ────────────────────────────────────
# DocumentService (with stub embedder)
class StubEmbeddingService:
@@ -127,7 +125,6 @@ class StubEmbeddingService:
self._dim = dim
async def embed_query(self, text: str) -> list[float]:
# Deterministic embedding based on hash of text
h = hash(text) % 1000
return [(h + i) / 1000.0 for i in range(self._dim)]
@@ -140,8 +137,6 @@ class StubEmbeddingService:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
for i, chunk in enumerate(chunks):
@@ -154,53 +149,113 @@ class StubEmbeddingService:
@pytest.fixture
def rag_service() -> RagService:
"""Each RagService test gets its own fresh ephemeral store to avoid dimension conflicts."""
def documents() -> DocumentService:
"""Each DocumentService test gets its own fresh ephemeral store to avoid dimension conflicts."""
store = SqliteVecStore.ephemeral()
return RagService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
return DocumentService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
class TestRagService:
def _pages(text: str) -> list[PageText]:
return [PageText(page_number=1, text=text)]
class TestDocumentService:
@pytest.mark.anyio
async def test_index_and_search(self, rag_service: RagService) -> None:
async def test_ingest_and_search(self, documents: DocumentService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await rag_service.index_text(FileId("docs"), text, source="guide.pdf")
count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf")
assert count > 0
results = await rag_service.search("Python libraries", collection=FileId("docs"))
results = await documents.search("Python libraries", collection=FileId("docs"))
assert len(results) > 0
assert results[0].document.text # non-empty text
assert results[0].document.text
@pytest.mark.anyio
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
count = await rag_service.index_text(FileId("docs"), "", source="empty.pdf")
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None:
count = await documents.ingest(FileId("docs"), _pages(""), source="empty.pdf")
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
results = await rag_service.search("anything", collection=FileId("nonexistent"))
async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None:
results = await documents.search("anything", collection=FileId("nonexistent"))
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("col-a"), "Machine learning overview.", source="ml.pdf")
await rag_service.index_text(FileId("col-b"), "Deep learning with neural networks.", source="dl.pdf")
async def test_search_all_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf")
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf")
results = await rag_service.search("neural networks")
results = await documents.search("neural networks")
assert len(results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf")
collections = await rag_service.list_collections()
async def test_delete_collection(self, documents: DocumentService) -> None:
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf")
collections = await documents.list_collections()
assert "temp" in collections
await rag_service.delete_collection(FileId("temp"))
collections = await rag_service.list_collections()
await documents.delete_collection(FileId("temp"))
collections = await documents.list_collections()
assert "temp" not in collections
@pytest.mark.anyio
async def test_ingest_stores_pages_in_order(self, documents: DocumentService) -> None:
pages = [
PageText(page_number=1, text="First page text."),
PageText(page_number=2, text="Second page text."),
PageText(page_number=3, text="Third page text."),
]
await documents.ingest(FileId("ordered"), pages, source="ordered.pdf")
# ── RagCapability ──────────────────────────────────────────────────────
stored = await documents.read_pages(FileId("ordered"))
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[0].text == "First page text."
assert stored[0].char_count == len("First page text.")
@pytest.mark.anyio
async def test_read_pages_with_range(self, documents: DocumentService) -> None:
from stirling.contracts import PageRange
pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)]
await documents.ingest(FileId("ranged"), pages, source="r.pdf")
subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4))
assert [p.page_number for p in subset] == [2, 3, 4]
@pytest.mark.anyio
async def test_ingest_replaces_previous_pages(self, documents: DocumentService) -> None:
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
source="v1.pdf",
)
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="new")],
source="v2.pdf",
)
stored = await documents.read_pages(FileId("doc"))
assert [p.page_number for p in stored] == [1]
assert stored[0].text == "new"
@pytest.mark.anyio
async def test_ingest_keeps_blank_pages_in_page_store(self, documents: DocumentService) -> None:
"""Blank pages are skipped for embedding but retained in the page store
so page numbering stays continuous when reading back."""
pages = [
PageText(page_number=1, text="Real text on page 1."),
PageText(page_number=2, text=" "),
PageText(page_number=3, text="Real text on page 3."),
]
await documents.ingest(FileId("with-blanks"), pages, source="blanks.pdf")
stored = await documents.read_pages(FileId("with-blanks"))
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[1].text.strip() == ""
# RagCapability
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
@@ -210,27 +265,27 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
toolset = capability.toolset
assert isinstance(toolset, FunctionToolset)
tool = toolset.tools["search_knowledge"]
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg] — pyright can't infer the generic tool function's kwargs
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg]
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None:
cap = RagCapability(documents, collections=[FileId("docs"), FileId("manuals")])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
instructions = cap.instructions
assert callable(instructions)
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("col-a"), "Alpha content.", source="a.pdf")
await rag_service.index_text(FileId("col-b"), "Beta content.", source="b.pdf")
cap = RagCapability(rag_service)
async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf")
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf")
cap = RagCapability(documents)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -238,23 +293,23 @@ class TestRagCapability:
assert "col-b" in text
@pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
assert "empty" in text.lower()
@pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, rag_service: RagService) -> None:
cap = RagCapability(rag_service)
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf")
cap = RagCapability(rag_service)
async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None:
await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf")
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
assert "source: guide.pdf" in output
@@ -262,43 +317,39 @@ class TestRagCapability:
assert "relevance:" in output
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
await rag_service.index_text(FileId("pinned"), "Pinned collection content.", source="pinned.pdf")
await rag_service.index_text(FileId("other"), "Content in another collection.", source="other.pdf")
async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf")
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf")
cap = RagCapability(rag_service, collections=[FileId("pinned")])
cap = RagCapability(documents, collections=[FileId("pinned")])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await rag_service.index_text(FileId("bulk"), paragraphs, source="bulk.pdf")
await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf")
cap = RagCapability(rag_service)
cap = RagCapability(documents)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
# Only two results requested, so only Result 1 and Result 2 should appear
assert "[Result 1" in output
assert "[Result 2" in output
assert "[Result 3" not in output
@pytest.mark.anyio
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, rag_service: RagService) -> None:
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None:
"""The prepare callback must return None once max_searches has been reached
so the agent can no longer call the tool on subsequent turns."""
await rag_service.index_text(FileId("docs"), "Some content.", source="x.pdf")
cap = RagCapability(rag_service, max_searches=2)
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf")
cap = RagCapability(documents, max_searches=2)
tool_def = _dummy_tool_def()
# Budget intact: prepare returns the tool definition.
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
# Use the budget.
await _invoke_search_knowledge(cap, "content")
await _invoke_search_knowledge(cap, "content")
# Budget spent: prepare returns None, removing the tool from the agent's next turn.
assert await cap._prepare_search_knowledge(None, tool_def) is None # type: ignore[arg-type]
@@ -6,9 +6,9 @@ import pytest
from fastapi.testclient import TestClient
from stirling.api import app
from stirling.api.dependencies import get_rag_service
from stirling.api.dependencies import get_document_service
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.rag import Document, RagService, SqliteVecStore
class StubEmbedder:
@@ -30,7 +30,7 @@ class StubEmbedder:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs = []
@@ -43,8 +43,8 @@ class StubEmbedder:
return docs
def _build_service() -> RagService:
return RagService(
def _build_service() -> DocumentService:
return DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=3,
@@ -52,25 +52,25 @@ def _build_service() -> RagService:
@pytest.fixture
def service() -> RagService:
def service() -> DocumentService:
return _build_service()
@pytest.fixture
def client(service: RagService) -> Iterator[TestClient]:
app.dependency_overrides[get_rag_service] = lambda: service
def client(service: DocumentService) -> Iterator[TestClient]:
app.dependency_overrides[get_document_service] = lambda: service
try:
yield TestClient(app)
finally:
app.dependency_overrides.pop(get_rag_service, None)
app.dependency_overrides.pop(get_document_service, None)
# ── POST /documents ─────────────────────────────────────────────────────
def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None:
def test_ingest_document_indexes_page_text(client: TestClient, service: DocumentService) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "doc-123",
"source": "report.pdf",
@@ -87,9 +87,9 @@ def test_ingest_document_indexes_page_text(client: TestClient, service: RagServi
@pytest.mark.anyio
async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None:
async def test_ingest_document_replaces_existing_content(client: TestClient, service: DocumentService) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "replace-me",
"source": "replace-me.pdf",
@@ -98,7 +98,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
)
# Second ingest with different content should replace the first entirely
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "replace-me",
"source": "replace-me.pdf",
@@ -115,7 +115,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "mixed",
"source": "mixed.pdf",
@@ -130,14 +130,14 @@ def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
response = client.post("/api/v1/rag/documents", json={"documentId": "empty", "source": "empty.pdf"})
response = client.post("/api/v1/documents", json={"documentId": "empty", "source": "empty.pdf"})
assert response.status_code == 200
assert response.json()["chunksIndexed"] == 0
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -145,7 +145,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -153,7 +153,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
)
assert response.status_code == 422
@@ -161,7 +161,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
response = client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "bad-page",
"source": "bad-page.pdf",
@@ -176,30 +176,30 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) ->
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={
"documentId": "to-delete",
"source": "to-delete.pdf",
"pageText": [{"pageNumber": 1, "text": "Text."}],
},
)
response = client.delete("/api/v1/rag/documents/to-delete")
response = client.delete("/api/v1/documents/to-delete")
assert response.status_code == 200
assert response.json() == {"documentId": "to-delete", "deleted": True}
def test_delete_document_is_idempotent(client: TestClient) -> None:
response = client.delete("/api/v1/rag/documents/never-existed")
response = client.delete("/api/v1/documents/never-existed")
assert response.status_code == 200
assert response.json() == {"documentId": "never-existed", "deleted": False}
@pytest.mark.anyio
async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None:
async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
client.post(
"/api/v1/rag/documents",
"/api/v1/documents",
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
)
assert await service.has_collection(FileId("gone"))
client.delete("/api/v1/rag/documents/gone")
client.delete("/api/v1/documents/gone")
assert not await service.has_collection(FileId("gone"))
+15 -14
View File
@@ -9,6 +9,7 @@ from stirling.contracts import (
AiFile,
ExtractedFileText,
NeedIngestResponse,
PageText,
PdfContentType,
PdfQuestionAnswerResponse,
PdfQuestionNotFoundResponse,
@@ -17,8 +18,8 @@ from stirling.contracts import (
PdfTextSelection,
SupportedCapability,
)
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.rag import Document, RagService, SqliteVecStore
from stirling.services.runtime import AppRuntime
@@ -41,7 +42,7 @@ class StubEmbedder:
source: str = "",
base_metadata: dict[str, str] | None = None,
) -> list[Document]:
from stirling.rag.chunker import chunk_text
from stirling.documents.chunker import chunk_text
chunks = chunk_text(text, 100, 10)
docs: list[Document] = []
@@ -65,13 +66,13 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
@pytest.fixture
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
"""A runtime whose RAG service uses a stub embedder + ephemeral store."""
stub = RagService(
"""A runtime whose document service uses a stub embedder + ephemeral store."""
stub = DocumentService(
embedder=StubEmbedder(), # type: ignore[arg-type]
store=SqliteVecStore.ephemeral(),
default_top_k=runtime.settings.rag_default_top_k,
)
return replace(runtime, rag_service=stub)
return replace(runtime, documents=stub)
@pytest.mark.anyio
@@ -89,9 +90,9 @@ async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag:
@pytest.mark.anyio
async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("present-id"),
text="Invoice total: 120.00.",
await runtime_with_stub_rag.documents.ingest(
FileId("present-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="present.pdf",
)
agent = PdfQuestionAgent(runtime_with_stub_rag)
@@ -106,9 +107,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) ->
@pytest.mark.anyio
async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("invoice-id"),
text="Invoice total: 120.00.",
await runtime_with_stub_rag.documents.ingest(
FileId("invoice-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="invoice.pdf",
)
agent = StubPdfQuestionAgent(
@@ -137,9 +138,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub
@pytest.mark.anyio
async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
await runtime_with_stub_rag.rag_service.index_text(
collection=FileId("shipping-id"),
text="This page contains only a shipping address.",
await runtime_with_stub_rag.documents.ingest(
FileId("shipping-id"),
[PageText(page_number=1, text="This page contains only a shipping address.")],
source="shipping.pdf",
)
agent = StubPdfQuestionAgent(
+115 -5
View File
@@ -1,3 +1,7 @@
import asyncio
import json
from unittest.mock import patch
from conftest import build_app_settings
from fastapi.testclient import TestClient
@@ -25,8 +29,11 @@ from stirling.contracts import (
PdfQuestionNotFoundResponse,
PdfQuestionRequest,
SupportedCapability,
WholeDocReadStarted,
WholeDocSliceDone,
)
from stirling.models.tool_models import Angle, RotatePdfParams
from stirling.services import emit_progress
class StubOrchestratorAgent:
@@ -40,6 +47,34 @@ class StubOrchestratorAgent:
)
class StubProgressOrchestratorAgent:
"""Orchestrator stub that emits two progress events before returning.
Used to verify the streaming endpoint plumbs the ContextVar emitter through
to deep callees and forwards events as NDJSON in order.
"""
async def handle(self, request: OrchestratorRequest) -> NeedContentResponse:
await emit_progress(WholeDocReadStarted(question="x", pages=10, slices=2))
await emit_progress(
WholeDocSliceDone(
completed=1,
total=2,
pages="pages=1-5",
duration_ms=42,
excerpts=2,
facts=3,
)
)
return NeedContentResponse(
resume_with=SupportedCapability.PDF_QUESTION,
reason=request.user_message,
files=[],
max_pages=1,
max_characters=1000,
)
class StubPdfEditAgent:
async def handle(self, request: PdfEditRequest) -> EditCannotDoResponse:
return EditCannotDoResponse(reason=request.user_message)
@@ -87,14 +122,89 @@ def test_health_route() -> None:
assert response.json()["status"] == "ok"
def test_orchestrator_route() -> None:
response = client.post(
def test_orchestrator_route_streams_result_only_when_no_progress() -> None:
"""The orchestrator endpoint always streams NDJSON. An agent that emits no
progress events still produces a single ``result`` frame with the typed
response body."""
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
)
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
assert response.status_code == 200
assert response.json()["outcome"] == "need_content"
assert [e["event"] for e in events] == ["result"]
body = events[0]["response"]
assert body["outcome"] == "need_content"
def test_orchestrator_route_streams_progress_then_result() -> None:
"""When an agent emits progress via the ContextVar emitter, those frames
arrive on the wire before the final result frame."""
app.dependency_overrides[get_orchestrator_agent] = lambda: StubProgressOrchestratorAgent()
try:
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "stream this", "files": [{"id": "test-id", "name": "test.pdf"}]},
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
finally:
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
progress = [e for e in events if e["event"] == "progress"]
results = [e for e in events if e["event"] == "result"]
assert [p["phase"] for p in progress] == ["whole_doc_read_started", "whole_doc_slice_done"]
assert progress[1]["completed"] == 1
assert len(results) == 1
response = results[0]["response"]
assert response["outcome"] == "need_content"
# Wire format must be camelCase: Java's Jackson deserializer expects camelCase
# field names. ``maxPages`` here doubles as a regression guard against the
# snake_case bug that surfaced as "need_ingest without listing any files to ingest".
assert "maxPages" in response
assert "max_pages" not in response
def test_orchestrator_route_emits_heartbeats_while_agent_is_busy() -> None:
"""While the agent is in flight, the streaming endpoint emits heartbeat
frames at the configured cadence so each layer of the connection stays
visibly alive and disconnects propagate within bounded latency."""
class _SlowAgent:
async def handle(self, _request: OrchestratorRequest) -> NeedContentResponse:
# Sleep long enough for several heartbeats at the patched cadence.
await asyncio.sleep(0.2)
return NeedContentResponse(
resume_with=SupportedCapability.PDF_QUESTION,
reason="ok",
files=[],
max_pages=1,
max_characters=1000,
)
app.dependency_overrides[get_orchestrator_agent] = lambda: _SlowAgent()
try:
with patch("stirling.api.routes.orchestrator.HEARTBEAT_INTERVAL_SECONDS", 0.03):
with client.stream(
"POST",
"/api/v1/orchestrator",
json={"userMessage": "wait", "files": [{"id": "test-id", "name": "test.pdf"}]},
) as response:
assert response.status_code == 200
events = [json.loads(line) for line in response.iter_lines() if line]
finally:
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
heartbeats = [e for e in events if e["event"] == "heartbeat"]
results = [e for e in events if e["event"] == "result"]
# At least a couple of heartbeats fired during the 0.2s agent sleep at 0.03s cadence.
assert len(heartbeats) >= 2
# The result still arrives after the agent finishes.
assert len(results) == 1
assert results[0]["response"]["outcome"] == "need_content"
def test_pdf_edit_route() -> None:
+4
View File
@@ -92,6 +92,10 @@ def test_app_settings_accepts_model_configuration() -> None:
rag_chunk_overlap=64,
rag_default_top_k=5,
rag_max_searches=5,
chunked_reasoner_chars_per_slice=16_000,
chunked_reasoner_concurrency=10,
chunked_reasoner_worker_timeout_seconds=60.0,
chunked_reasoner_notes_char_budget=250_000,
max_pages=200,
max_characters=200_000,
posthog_enabled=False,
@@ -2682,6 +2682,10 @@ executing_tool_step = "Running {{tool}} (step {{step}} of {{total}})..."
extracting_content = "Extracting content from your documents..."
processing = "Processing extracted content..."
thinking = "Thinking..."
whole_doc_compression_round = "Consolidating notes..."
whole_doc_read_done = "Finished reading the document..."
whole_doc_read_started = "Reading the document..."
whole_doc_slice_done = "Reading the document... ({{percent}}% complete)"
[chat.toolsUsed]
summary = "Ran {{count}} tools"
@@ -37,6 +37,79 @@ export enum AiWorkflowPhase {
EXTRACTING_CONTENT = "extracting_content",
EXECUTING_TOOL = "executing_tool",
PROCESSING = "processing",
ENGINE_PROGRESS = "engine_progress",
}
/**
* Engine-side progress detail for ENGINE_PROGRESS events. Mirrors the Python
* {@code ProgressEvent} discriminated union (engine/src/stirling/contracts/progress.py)
* and the Java {@code AiEngineProgressDetail} sealed interface; the {@code phase}
* string is the discriminator. Field names are camelCase because the engine
* serialises by alias.
*/
export interface WholeDocReadStartedDetail {
phase: "whole_doc_read_started";
question: string;
pages: number;
slices: number;
}
export interface WholeDocSliceDoneDetail {
phase: "whole_doc_slice_done";
completed: number;
total: number;
/** Page-range label, e.g. "pages=1-5". */
pages: string;
durationMs: number;
excerpts: number;
facts: number;
}
export interface WholeDocCompressionRoundDetail {
phase: "whole_doc_compression_round";
roundNumber: number;
notesIn: number;
groups: number;
}
export interface WholeDocReadDoneDetail {
phase: "whole_doc_read_done";
completed: number;
slices: number;
durationSeconds: number;
}
export type EngineProgressDetail =
| WholeDocReadStartedDetail
| WholeDocSliceDoneDetail
| WholeDocCompressionRoundDetail
| WholeDocReadDoneDetail;
/**
* What we actually carry across the wire boundary: a known typed variant, or a
* forward-compat shape with just the discriminator string. The "unknown" arm
* exists so a new engine-side phase rolling out before a frontend update keeps
* rendering the generic processing message instead of crashing the union.
*/
export interface UnknownEngineProgressDetail {
phase: string;
}
export type AnyEngineProgressDetail =
| EngineProgressDetail
| UnknownEngineProgressDetail;
const KNOWN_ENGINE_PHASES = new Set<string>([
"whole_doc_read_started",
"whole_doc_slice_done",
"whole_doc_compression_round",
"whole_doc_read_done",
]);
export function isKnownEngineProgressDetail(
detail: AnyEngineProgressDetail,
): detail is EngineProgressDetail {
return KNOWN_ENGINE_PHASES.has(detail.phase);
}
export interface AiWorkflowProgress {
@@ -47,6 +120,11 @@ export interface AiWorkflowProgress {
stepIndex?: number;
/** Total number of plan steps, for EXECUTING_TOOL events. */
stepCount?: number;
/**
* Engine-side event payload, for ENGINE_PROGRESS events. Typed sub-phase
* record (e.g. {@link WholeDocSliceDoneDetail}) the UI can render in detail.
*/
engineDetail?: AnyEngineProgressDetail;
}
type AiWorkflowOutcome =
@@ -159,6 +237,7 @@ interface ProgressEvent {
tool?: string;
stepIndex?: number;
stepCount?: number;
engineDetail?: AnyEngineProgressDetail;
}
async function consumeSSEStream(
@@ -378,6 +457,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
tool: data.tool,
stepIndex: data.stepIndex,
stepCount: data.stepCount,
engineDetail: data.engineDetail,
},
});
},
@@ -29,7 +29,9 @@ import ExpandLessIcon from "@mui/icons-material/ExpandLess";
import {
useChat,
AiWorkflowPhase,
isKnownEngineProgressDetail,
type AiWorkflowProgress,
type AnyEngineProgressDetail,
} from "@app/components/chat/ChatContext";
import { useTranslatedToolCatalog } from "@app/data/useTranslatedToolRegistry";
import "@app/components/chat/ChatPanel.css";
@@ -90,9 +92,41 @@ function formatProgress(
})
: t("chat.progress.executing_tool_generic");
}
if (progress.phase === AiWorkflowPhase.ENGINE_PROGRESS) {
return formatEngineProgress(progress.engineDetail, t);
}
return t(`chat.progress.${progress.phase}`);
}
/**
* Render an engine-side progress event (e.g. chunked-reasoner slice progress) into a user-facing
* message. Falls through to the generic processing label for unknown sub-phases so adding new
* engine events doesn't break the UI before the frontend learns about them.
*/
function formatEngineProgress(
detail: AnyEngineProgressDetail | undefined,
t: TranslateFn,
): string {
if (!detail || !isKnownEngineProgressDetail(detail)) {
return t("chat.progress.processing");
}
switch (detail.phase) {
case "whole_doc_read_started":
return t("chat.progress.whole_doc_read_started");
case "whole_doc_slice_done": {
const percent =
detail.total > 0
? Math.round((detail.completed / detail.total) * 100)
: 0;
return t("chat.progress.whole_doc_slice_done", { percent });
}
case "whole_doc_compression_round":
return t("chat.progress.whole_doc_compression_round");
case "whole_doc_read_done":
return t("chat.progress.whole_doc_read_done");
}
}
function ToolsUsedBlock({
tools,
resolveToolName,