From 672e81d28643860fe7990dc9d2bcddbe76577ae8 Mon Sep 17 00:00:00 2001 From: James Brunton Date: Thu, 14 May 2026 14:19:38 +0100 Subject: [PATCH] Add ability for Stirling engine to reason across large documents (#6314) # Description of Changes Adds storage in the database for full document content alongside the RAG content (and changes the service to `DocumentService` instead of `RagService`). Then adds a generic capability that should be usable by any agent (currently just used by the Question Agent) which allows the agent to pull out the full contents of the doc, chunks it into various sections that will fit in the context window, and then processes them in parallel to create an intermediate result, and then processes the intermediate result into a final answer. It will re-chunk as many times as necessary to get the content small enough for the actual answer to be analysed (I've tested on PDFs ~3500 pages long, which is well above the context limit and requires maybe 3 rounds of compression to get an answer). The new full doc analysis stuff is heavier than the RAG lookup so both remain. The agents should use RAG for targeted info and the chunked reasoner for info that requires reading the full doc. --- .../controller/api/AiEngineController.java | 42 +- .../model/api/ai/AiDocumentIngestRequest.java | 24 + .../model/api/ai/AiEngineProgressDetail.java | 59 ++ .../{AiRagPageText.java => AiPageText.java} | 4 +- .../model/api/ai/AiRagIngestRequest.java | 24 - .../model/api/ai/AiWorkflowPhase.java | 8 +- .../model/api/ai/AiWorkflowProgressEvent.java | 24 +- .../proprietary/service/AiEngineClient.java | 66 ++ .../service/AiWorkflowService.java | 82 ++- .../service/AiWorkflowServiceTest.java | 67 +- engine/.env | 21 + engine/src/stirling/agents/pdf_questions.py | 70 +- engine/src/stirling/agents/shared/__init__.py | 10 + .../agents/shared/chunked_reasoner.py | 616 ++++++++++++++++++ .../agents/shared/whole_doc_reader.py | 140 ++++ engine/src/stirling/api/app.py | 6 +- engine/src/stirling/api/dependencies.py | 6 +- engine/src/stirling/api/routes/__init__.py | 4 +- engine/src/stirling/api/routes/documents.py | 49 ++ .../src/stirling/api/routes/orchestrator.py | 161 ++++- engine/src/stirling/api/routes/rag.py | 63 -- engine/src/stirling/config/settings.py | 56 +- engine/src/stirling/contracts/__init__.py | 28 +- engine/src/stirling/contracts/common.py | 4 +- engine/src/stirling/contracts/documents.py | 61 ++ engine/src/stirling/contracts/progress.py | 70 ++ engine/src/stirling/contracts/rag.py | 39 -- .../src/stirling/{rag => documents}/README.md | 48 +- engine/src/stirling/documents/__init__.py | 20 + .../stirling/{rag => documents}/chunker.py | 0 .../stirling/{rag => documents}/embedder.py | 4 +- .../{rag => documents}/pgvector_store.py | 94 ++- .../rag_capability.py} | 14 +- engine/src/stirling/documents/service.py | 133 ++++ .../{rag => documents}/sqlite_vec_store.py | 102 ++- engine/src/stirling/documents/store.py | 113 ++++ engine/src/stirling/models/base.py | 3 + engine/src/stirling/rag/__init__.py | 19 - engine/src/stirling/rag/service.py | 102 --- engine/src/stirling/rag/store.py | 63 -- engine/src/stirling/services/__init__.py | 10 + engine/src/stirling/services/progress.py | 47 ++ engine/src/stirling/services/runtime.py | 72 +- engine/tests/agents/test_chunked_reasoner.py | 511 +++++++++++++++ engine/tests/agents/test_whole_doc_reader.py | 220 +++++++ engine/tests/conftest.py | 4 + engine/tests/pdf_comment/test_routes.py | 4 + .../tests/{test_rag.py => test_documents.py} | 195 ++++-- ...rag_routes.py => test_documents_routes.py} | 52 +- engine/tests/test_pdf_question_agent.py | 29 +- engine/tests/test_stirling_api.py | 120 +++- engine/tests/test_stirling_contracts.py | 4 + .../public/locales/en-GB/translation.toml | 4 + .../components/chat/ChatContext.tsx | 80 +++ .../prototypes/components/chat/ChatPanel.tsx | 34 + 55 files changed, 3327 insertions(+), 578 deletions(-) create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiEngineProgressDetail.java rename app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/{AiRagPageText.java => AiPageText.java} (71%) delete mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagIngestRequest.java create mode 100644 engine/src/stirling/agents/shared/__init__.py create mode 100644 engine/src/stirling/agents/shared/chunked_reasoner.py create mode 100644 engine/src/stirling/agents/shared/whole_doc_reader.py create mode 100644 engine/src/stirling/api/routes/documents.py delete mode 100644 engine/src/stirling/api/routes/rag.py create mode 100644 engine/src/stirling/contracts/documents.py create mode 100644 engine/src/stirling/contracts/progress.py delete mode 100644 engine/src/stirling/contracts/rag.py rename engine/src/stirling/{rag => documents}/README.md (50%) create mode 100644 engine/src/stirling/documents/__init__.py rename engine/src/stirling/{rag => documents}/chunker.py (100%) rename engine/src/stirling/{rag => documents}/embedder.py (96%) rename engine/src/stirling/{rag => documents}/pgvector_store.py (53%) rename engine/src/stirling/{rag/capability.py => documents/rag_capability.py} (93%) create mode 100644 engine/src/stirling/documents/service.py rename engine/src/stirling/{rag => documents}/sqlite_vec_store.py (67%) create mode 100644 engine/src/stirling/documents/store.py delete mode 100644 engine/src/stirling/rag/__init__.py delete mode 100644 engine/src/stirling/rag/service.py delete mode 100644 engine/src/stirling/rag/store.py create mode 100644 engine/src/stirling/services/progress.py create mode 100644 engine/tests/agents/test_chunked_reasoner.py create mode 100644 engine/tests/agents/test_whole_doc_reader.py rename engine/tests/{test_rag.py => test_documents.py} (56%) rename engine/tests/{test_rag_routes.py => test_documents_routes.py} (83%) 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,