diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java index 068314903..28fc7168d 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java @@ -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); } } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java new file mode 100644 index 000000000..47c431807 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java @@ -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 pageText; +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiEngineProgressDetail.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiEngineProgressDetail.java new file mode 100644 index 000000000..6dce13854 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiEngineProgressDetail.java @@ -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. + * + *

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 {} +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagPageText.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPageText.java similarity index 71% rename from app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagPageText.java rename to app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPageText.java index 417b04a38..6b6fa8720 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagPageText.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPageText.java @@ -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; diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagIngestRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagIngestRequest.java deleted file mode 100644 index ddb9081b6..000000000 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagIngestRequest.java +++ /dev/null @@ -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 pageText; -} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowPhase.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowPhase.java index b1ab9fff2..be2e88c70 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowPhase.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowPhase.java @@ -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; diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowProgressEvent.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowProgressEvent.java index 92c15a000..abd3c0920 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowProgressEvent.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowProgressEvent.java @@ -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); } } diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java index fb9963ba8..f5585ea5b 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java @@ -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. + * + *

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 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> 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 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()) { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java index 64de88e4d..b0352fe20 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java @@ -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 pages = new ArrayList<>(); + List 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 diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java index 3b450bc6c..b3693aa58 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java @@ -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 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 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) { diff --git a/engine/.env b/engine/.env index 18dcfa99d..e334a9bf9 100644 --- a/engine/.env +++ b/engine/.env @@ -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 diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index c27b7e612..04296e897 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -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." ) diff --git a/engine/src/stirling/agents/shared/__init__.py b/engine/src/stirling/agents/shared/__init__.py new file mode 100644 index 000000000..2b7a48e7f --- /dev/null +++ b/engine/src/stirling/agents/shared/__init__.py @@ -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", +] diff --git a/engine/src/stirling/agents/shared/chunked_reasoner.py b/engine/src/stirling/agents/shared/chunked_reasoner.py new file mode 100644 index 000000000..c2c19ef44 --- /dev/null +++ b/engine/src/stirling/agents/shared/chunked_reasoner.py @@ -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 diff --git a/engine/src/stirling/agents/shared/whole_doc_reader.py b/engine/src/stirling/agents/shared/whole_doc_reader.py new file mode 100644 index 000000000..324a5f535 --- /dev/null +++ b/engine/src/stirling/agents/shared/whole_doc_reader.py @@ -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) diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index 97cc700da..3539379bb 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -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) diff --git a/engine/src/stirling/api/dependencies.py b/engine/src/stirling/api/dependencies.py index 158f67cd6..70cbad877 100644 --- a/engine/src/stirling/api/dependencies.py +++ b/engine/src/stirling/api/dependencies.py @@ -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: diff --git a/engine/src/stirling/api/routes/__init__.py b/engine/src/stirling/api/routes/__init__.py index 62be04904..37572c5e7 100644 --- a/engine/src/stirling/api/routes/__init__.py +++ b/engine/src/stirling/api/routes/__init__.py @@ -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", ] diff --git a/engine/src/stirling/api/routes/documents.py b/engine/src/stirling/api/routes/documents.py new file mode 100644 index 000000000..4887b7ac3 --- /dev/null +++ b/engine/src/stirling/api/routes/documents.py @@ -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) diff --git a/engine/src/stirling/api/routes/orchestrator.py b/engine/src/stirling/api/routes/orchestrator.py index bfd5779e7..b015c2bcf 100644 --- a/engine/src/stirling/api/routes/orchestrator.py +++ b/engine/src/stirling/api/routes/orchestrator.py @@ -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) diff --git a/engine/src/stirling/api/routes/rag.py b/engine/src/stirling/api/routes/rag.py deleted file mode 100644 index 8e36ca8a3..000000000 --- a/engine/src/stirling/api/routes/rag.py +++ /dev/null @@ -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) diff --git a/engine/src/stirling/config/settings.py b/engine/src/stirling/config/settings.py index 9087307a6..3b9fd2479 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -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 diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index d6657ed3e..6858ba67e 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -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", ] diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 73df8b5b2..3cd895e14 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -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 diff --git a/engine/src/stirling/contracts/documents.py b/engine/src/stirling/contracts/documents.py new file mode 100644 index 000000000..23e2ee8f4 --- /dev/null +++ b/engine/src/stirling/contracts/documents.py @@ -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 diff --git a/engine/src/stirling/contracts/progress.py b/engine/src/stirling/contracts/progress.py new file mode 100644 index 000000000..84d2d0436 --- /dev/null +++ b/engine/src/stirling/contracts/progress.py @@ -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"), +] diff --git a/engine/src/stirling/contracts/rag.py b/engine/src/stirling/contracts/rag.py deleted file mode 100644 index 9f655bbb4..000000000 --- a/engine/src/stirling/contracts/rag.py +++ /dev/null @@ -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 diff --git a/engine/src/stirling/rag/README.md b/engine/src/stirling/documents/README.md similarity index 50% rename from engine/src/stirling/rag/README.md rename to engine/src/stirling/documents/README.md index 309a3a7c9..e1c759044 100644 --- a/engine/src/stirling/rag/README.md +++ b/engine/src/stirling/documents/README.md @@ -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 | diff --git a/engine/src/stirling/documents/__init__.py b/engine/src/stirling/documents/__init__.py new file mode 100644 index 000000000..3df0eaaa0 --- /dev/null +++ b/engine/src/stirling/documents/__init__.py @@ -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", +] diff --git a/engine/src/stirling/rag/chunker.py b/engine/src/stirling/documents/chunker.py similarity index 100% rename from engine/src/stirling/rag/chunker.py rename to engine/src/stirling/documents/chunker.py diff --git a/engine/src/stirling/rag/embedder.py b/engine/src/stirling/documents/embedder.py similarity index 96% rename from engine/src/stirling/rag/embedder.py rename to engine/src/stirling/documents/embedder.py index c97c28759..13e9c7018 100644 --- a/engine/src/stirling/rag/embedder.py +++ b/engine/src/stirling/documents/embedder.py @@ -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 diff --git a/engine/src/stirling/rag/pgvector_store.py b/engine/src/stirling/documents/pgvector_store.py similarity index 53% rename from engine/src/stirling/rag/pgvector_store.py rename to engine/src/stirling/documents/pgvector_store.py index b73596e49..c79afafe4 100644 --- a/engine/src/stirling/rag/pgvector_store.py +++ b/engine/src/stirling/documents/pgvector_store.py @@ -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() diff --git a/engine/src/stirling/rag/capability.py b/engine/src/stirling/documents/rag_capability.py similarity index 93% rename from engine/src/stirling/rag/capability.py rename to engine/src/stirling/documents/rag_capability.py index 494000f08..bed22e94d 100644 --- a/engine/src/stirling/rag/capability.py +++ b/engine/src/stirling/documents/rag_capability.py @@ -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) diff --git a/engine/src/stirling/documents/service.py b/engine/src/stirling/documents/service.py new file mode 100644 index 000000000..e542ffa58 --- /dev/null +++ b/engine/src/stirling/documents/service.py @@ -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() diff --git a/engine/src/stirling/rag/sqlite_vec_store.py b/engine/src/stirling/documents/sqlite_vec_store.py similarity index 67% rename from engine/src/stirling/rag/sqlite_vec_store.py rename to engine/src/stirling/documents/sqlite_vec_store.py index 25bbffcd3..49f5041a2 100644 --- a/engine/src/stirling/rag/sqlite_vec_store.py +++ b/engine/src/stirling/documents/sqlite_vec_store.py @@ -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: diff --git a/engine/src/stirling/documents/store.py b/engine/src/stirling/documents/store.py new file mode 100644 index 000000000..d59fd9d63 --- /dev/null +++ b/engine/src/stirling/documents/store.py @@ -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.).""" diff --git a/engine/src/stirling/models/base.py b/engine/src/stirling/models/base.py index 815c09854..12c32b62c 100644 --- a/engine/src/stirling/models/base.py +++ b/engine/src/stirling/models/base.py @@ -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, ) diff --git a/engine/src/stirling/rag/__init__.py b/engine/src/stirling/rag/__init__.py deleted file mode 100644 index 3abc1f818..000000000 --- a/engine/src/stirling/rag/__init__.py +++ /dev/null @@ -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", -] diff --git a/engine/src/stirling/rag/service.py b/engine/src/stirling/rag/service.py deleted file mode 100644 index e818cdccd..000000000 --- a/engine/src/stirling/rag/service.py +++ /dev/null @@ -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() diff --git a/engine/src/stirling/rag/store.py b/engine/src/stirling/rag/store.py deleted file mode 100644 index 1a73300d7..000000000 --- a/engine/src/stirling/rag/store.py +++ /dev/null @@ -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.).""" diff --git a/engine/src/stirling/services/__init__.py b/engine/src/stirling/services/__init__.py index 8894f4114..32189312c 100644 --- a/engine/src/stirling/services/__init__.py +++ b/engine/src/stirling/services/__init__.py @@ -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", ] diff --git a/engine/src/stirling/services/progress.py b/engine/src/stirling/services/progress.py new file mode 100644 index 000000000..d9300faca --- /dev/null +++ b/engine/src/stirling/services/progress.py @@ -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) diff --git a/engine/src/stirling/services/runtime.py b/engine/src/stirling/services/runtime.py index 3b91abbd7..4d5841489 100644 --- a/engine/src/stirling/services/runtime.py +++ b/engine/src/stirling/services/runtime.py @@ -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=") + logger.info("Document store backend=pgvector, dsn=") 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) diff --git a/engine/tests/agents/test_chunked_reasoner.py b/engine/tests/agents/test_chunked_reasoner.py new file mode 100644 index 000000000..3c373304b --- /dev/null +++ b/engine/tests/agents/test_chunked_reasoner.py @@ -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"] diff --git a/engine/tests/agents/test_whole_doc_reader.py b/engine/tests/agents/test_whole_doc_reader.py new file mode 100644 index 000000000..588a43f37 --- /dev/null +++ b/engine/tests/agents/test_whole_doc_reader.py @@ -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 diff --git a/engine/tests/conftest.py b/engine/tests/conftest.py index 36194b37f..91cce3fb2 100644 --- a/engine/tests/conftest.py +++ b/engine/tests/conftest.py @@ -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, diff --git a/engine/tests/pdf_comment/test_routes.py b/engine/tests/pdf_comment/test_routes.py index 9d35a615e..6fce78f25 100644 --- a/engine/tests/pdf_comment/test_routes.py +++ b/engine/tests/pdf_comment/test_routes.py @@ -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, diff --git a/engine/tests/test_rag.py b/engine/tests/test_documents.py similarity index 56% rename from engine/tests/test_rag.py rename to engine/tests/test_documents.py index 9616ae834..a5febf078 100644 --- a/engine/tests/test_rag.py +++ b/engine/tests/test_documents.py @@ -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] diff --git a/engine/tests/test_rag_routes.py b/engine/tests/test_documents_routes.py similarity index 83% rename from engine/tests/test_rag_routes.py rename to engine/tests/test_documents_routes.py index f9f61c996..844d78518 100644 --- a/engine/tests/test_rag_routes.py +++ b/engine/tests/test_documents_routes.py @@ -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")) diff --git a/engine/tests/test_pdf_question_agent.py b/engine/tests/test_pdf_question_agent.py index 714406339..26bd5c0ff 100644 --- a/engine/tests/test_pdf_question_agent.py +++ b/engine/tests/test_pdf_question_agent.py @@ -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( diff --git a/engine/tests/test_stirling_api.py b/engine/tests/test_stirling_api.py index 547b1e2a5..4eaecf90f 100644 --- a/engine/tests/test_stirling_api.py +++ b/engine/tests/test_stirling_api.py @@ -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: diff --git a/engine/tests/test_stirling_contracts.py b/engine/tests/test_stirling_contracts.py index 7e4a2416d..ccc057adb 100644 --- a/engine/tests/test_stirling_contracts.py +++ b/engine/tests/test_stirling_contracts.py @@ -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, diff --git a/frontend/public/locales/en-GB/translation.toml b/frontend/public/locales/en-GB/translation.toml index 05be4b505..9dc9a8915 100644 --- a/frontend/public/locales/en-GB/translation.toml +++ b/frontend/public/locales/en-GB/translation.toml @@ -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" diff --git a/frontend/src/prototypes/components/chat/ChatContext.tsx b/frontend/src/prototypes/components/chat/ChatContext.tsx index 027bd5695..88e2b2b81 100644 --- a/frontend/src/prototypes/components/chat/ChatContext.tsx +++ b/frontend/src/prototypes/components/chat/ChatContext.tsx @@ -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([ + "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, }, }); }, diff --git a/frontend/src/prototypes/components/chat/ChatPanel.tsx b/frontend/src/prototypes/components/chat/ChatPanel.tsx index 342f74bfa..8d916b29a 100644 --- a/frontend/src/prototypes/components/chat/ChatPanel.tsx +++ b/frontend/src/prototypes/components/chat/ChatPanel.tsx @@ -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,