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.
2026-07-16 11:23:10 +02:00 · 2026-05-14 13:19:38 +00:00
parent 8abe734f0b
commit 672e81d286
55 changed files with 3327 additions and 578 deletions
@@ -30,6 +30,7 @@ import lombok.extern.slf4j.Slf4j;
 import stirling.software.common.model.job.ResultFile;
 import stirling.software.common.service.JobOwnershipService;
 import stirling.software.common.service.TaskManager;
+import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
 import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
 import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
 import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
@@ -140,13 +141,32 @@ public class AiEngineController {
    }

    private void runOrchestrationStream(AiWorkflowRequest request, SseEmitter emitter) {
+        AiWorkflowService.ProgressListener listener =
+                new AiWorkflowService.ProgressListener() {
+                    @Override
+                    public void onProgress(AiWorkflowProgressEvent event) {
+                        sendEvent(emitter, "progress", event);
+                    }
+
+                    @Override
+                    public void onHeartbeat() {
+                        // Forward upstream heartbeats so the SSE pipe stays visibly alive between
+                        // real progress events; if the frontend has gone away, sendEvent throws,
+                        // which propagates up through the stream consumer and closes our upstream
+                        // engine connection so the engine can cancel its in-flight workflow.
+                        sendEvent(emitter, "heartbeat", Map.of());
+                    }
+                };
        try {
-            AiWorkflowResponse result =
-                    aiWorkflowService.orchestrate(
-                            request, progress -> sendEvent(emitter, "progress", progress));
+            AiWorkflowResponse result = aiWorkflowService.orchestrate(request, listener);
            registerFileResultAsJob(result);
            sendEvent(emitter, "result", result);
            emitter.complete();
+        } catch (ClientDisconnectedException e) {
+            // The frontend gave up mid-stream. The exception unwinding through orchestrate()
+            // already closed the upstream engine connection (engine sees disconnect and cancels).
+            // The emitter is already toast; nothing useful left to send.
+            log.debug("Client disconnected mid-stream; aborting workflow", e);
        } catch (Exception e) {
            log.error("AI orchestration stream failed", e);
            // Emit an error frame for the frontend and then complete normally. Using
@@ -192,7 +212,21 @@ public class AiEngineController {
        try {
            emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
        } catch (IOException e) {
-            log.debug("Failed to send SSE event (client may have disconnected)", e);
+            // Surface the disconnect so the streaming pipeline unwinds: callers higher up close
+            // the upstream engine connection, which lets the engine cancel its in-flight workflow.
+            // Without this, the engine would keep producing (and billing for) tokens whose results
+            // nobody is reading.
+            throw new ClientDisconnectedException("Client disconnected from SSE stream", e);
+        }
+    }
+
+    /**
+     * Thrown by {@link #sendEvent} when the SSE emitter's underlying connection is gone. Treated as
+     * a signal to abort the workflow, not as an error to report.
+     */
+    private static final class ClientDisconnectedException extends RuntimeException {
+        ClientDisconnectedException(String message, Throwable cause) {
+            super(message, cause);
        }
    }

@@ -0,0 +1,24 @@
+package stirling.software.proprietary.model.api.ai;
+
+import java.util.List;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+/**
+ * Body for {@code POST /api/v1/documents} on the AI engine. Sent by Java when the engine reports
+ * {@code need_ingest} and the requested document's extracted content must be stored before the
+ * workflow can continue.
+ */
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+public class AiDocumentIngestRequest {
+
+    private String documentId;
+
+    private String source;
+
+    private List<AiPageText> pageText;
+}
@@ -0,0 +1,59 @@
+package stirling.software.proprietary.model.api.ai;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+
+/**
+ * Typed engine-emitted progress detail, mirroring the Python {@code ProgressEvent} discriminated
+ * union (see {@code engine/src/stirling/contracts/progress.py}). Carried inside {@link
+ * AiWorkflowProgressEvent#getEngineDetail()} for {@link AiWorkflowPhase#ENGINE_PROGRESS} events.
+ *
+ * <p>Sealed so adding a new engine-side phase forces a matching subtype on the Java side instead of
+ * silently passing through as an opaque map. The {@code phase} string is the discriminator and
+ * stays on the wire so the frontend (which doesn't know about Java's class hierarchy) can switch on
+ * it.
+ */
+@JsonTypeInfo(
+        use = JsonTypeInfo.Id.NAME,
+        include = JsonTypeInfo.As.EXISTING_PROPERTY,
+        property = "phase",
+        visible = true)
+@JsonSubTypes({
+    @JsonSubTypes.Type(
+            value = AiEngineProgressDetail.WholeDocReadStarted.class,
+            name = "whole_doc_read_started"),
+    @JsonSubTypes.Type(
+            value = AiEngineProgressDetail.WholeDocSliceDone.class,
+            name = "whole_doc_slice_done"),
+    @JsonSubTypes.Type(
+            value = AiEngineProgressDetail.WholeDocCompressionRound.class,
+            name = "whole_doc_compression_round"),
+    @JsonSubTypes.Type(
+            value = AiEngineProgressDetail.WholeDocReadDone.class,
+            name = "whole_doc_read_done"),
+})
+@JsonIgnoreProperties(ignoreUnknown = true)
+public sealed interface AiEngineProgressDetail {
+
+    String phase();
+
+    record WholeDocReadStarted(String phase, String question, int pages, int slices)
+            implements AiEngineProgressDetail {}
+
+    record WholeDocSliceDone(
+            String phase,
+            int completed,
+            int total,
+            String pages,
+            int durationMs,
+            int excerpts,
+            int facts)
+            implements AiEngineProgressDetail {}
+
+    record WholeDocCompressionRound(String phase, int roundNumber, int notesIn, int groups)
+            implements AiEngineProgressDetail {}
+
+    record WholeDocReadDone(String phase, int completed, int slices, double durationSeconds)
+            implements AiEngineProgressDetail {}
+}
@@ -4,11 +4,11 @@ import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;

-/** A single page of extracted text for RAG ingest requests. */
+/** A single page of extracted text for document ingest requests. */
@Data
@NoArgsConstructor
@AllArgsConstructor
-public class AiRagPageText {
+public class AiPageText {

    private int pageNumber;

@@ -1,24 +0,0 @@
-package stirling.software.proprietary.model.api.ai;
-
-import java.util.List;
-
-import lombok.AllArgsConstructor;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-
-/**
- * Body for {@code POST /api/v1/rag/documents} on the AI engine. Sent by Java when the engine
- * reports {@code need_ingest} and the requested document's extracted content must be stored before
- * the workflow can continue.
- */
-@Data
-@NoArgsConstructor
-@AllArgsConstructor
-public class AiRagIngestRequest {
-
-    private String documentId;
-
-    private String source;
-
-    private List<AiRagPageText> pageText;
-}
@@ -9,7 +9,13 @@ public enum AiWorkflowPhase {
    CALLING_ENGINE("calling_engine"),
    EXTRACTING_CONTENT("extracting_content"),
    EXECUTING_TOOL("executing_tool"),
-    PROCESSING("processing");
+    PROCESSING("processing"),
+    /**
+     * Generic engine-emitted progress event (e.g. chunked-reasoner slice progress). The original
+     * engine event JSON is carried in {@link AiWorkflowProgressEvent#getEngineDetail()}, including
+     * a specific {@code phase} string the frontend can switch on.
+     */
+    ENGINE_PROGRESS("engine_progress");

    private final String value;

@@ -23,8 +23,17 @@ public class AiWorkflowProgressEvent {
    /** Total number of plan steps, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
    private Integer stepCount;

+    /**
+     * Engine-emitted event payload, for {@link AiWorkflowPhase#ENGINE_PROGRESS} events. The payload
+     * is a typed subtype keyed on its {@code phase} string (e.g. {@code "whole_doc_slice_done"})
+     * carrying phase-specific fields (slice index, page range, durations, etc.) that the frontend
+     * can render as detailed progress.
+     */
+    private AiEngineProgressDetail engineDetail;
+
    public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
-        return new AiWorkflowProgressEvent(phase, System.currentTimeMillis(), null, null, null);
+        return new AiWorkflowProgressEvent(
+                phase, System.currentTimeMillis(), null, null, null, null);
    }

    public static AiWorkflowProgressEvent executingTool(String tool, int stepIndex, int stepCount) {
@@ -33,6 +42,17 @@ public class AiWorkflowProgressEvent {
                System.currentTimeMillis(),
                tool,
                stepIndex,
-                stepCount);
+                stepCount,
+                null);
+    }
+
+    public static AiWorkflowProgressEvent engineProgress(AiEngineProgressDetail detail) {
+        return new AiWorkflowProgressEvent(
+                AiWorkflowPhase.ENGINE_PROGRESS,
+                System.currentTimeMillis(),
+                null,
+                null,
+                null,
+                detail);
    }
 }
@@ -7,6 +7,8 @@ import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.net.http.HttpTimeoutException;
 import java.time.Duration;
+import java.util.function.Consumer;
+import java.util.stream.Stream;

 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpStatus;
@@ -83,6 +85,70 @@ public class AiEngineClient {
        return response.body();
    }

+    /**
+     * POST a JSON body and consume the response as a stream of NDJSON lines. Each line is passed to
+     * {@code lineConsumer} in arrival order; the call returns when the engine closes the stream.
+     *
+     * <p>This is the right shape for long-running orchestrator calls that emit incremental
+     * progress. The total HTTP timeout is the long-running timeout (typically 600s+), but in
+     * practice line arrival keeps the connection logically alive: as long as the engine emits
+     * events, the work is progressing. Genuine engine hangs still hit the total timeout.
+     */
+    public void streamPost(String path, String jsonBody, Consumer<String> lineConsumer)
+            throws IOException {
+        ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
+        if (!config.isEnabled()) {
+            throw new ResponseStatusException(
+                    HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
+        }
+
+        String url = config.getUrl().stripTrailing() + path;
+        Duration timeout = Duration.ofSeconds(config.getLongRunningTimeoutSeconds());
+        log.debug(
+                "Proxying AI engine streaming request to {} (timeout {}s)",
+                url,
+                timeout.toSeconds());
+
+        HttpRequest request =
+                HttpRequest.newBuilder()
+                        .uri(URI.create(url))
+                        .header("Content-Type", "application/json")
+                        .header("Accept", "application/x-ndjson")
+                        .timeout(timeout)
+                        .POST(HttpRequest.BodyPublishers.ofString(jsonBody))
+                        .build();
+
+        HttpResponse<Stream<String>> response;
+        try {
+            response = httpClient.send(request, HttpResponse.BodyHandlers.ofLines());
+        } catch (HttpTimeoutException e) {
+            throw new ResponseStatusException(HttpStatus.GATEWAY_TIMEOUT, "AI engine timed out", e);
+        } catch (IOException e) {
+            throw new ResponseStatusException(
+                    HttpStatus.SERVICE_UNAVAILABLE, "AI engine unreachable: " + e.getMessage(), e);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new ResponseStatusException(
+                    HttpStatus.SERVICE_UNAVAILABLE, "AI engine request was interrupted");
+        }
+
+        int status = response.statusCode();
+        if (status >= 400) {
+            throw new ResponseStatusException(
+                    HttpStatus.valueOf(status >= 500 ? 502 : status),
+                    "AI engine returned error: " + status);
+        }
+
+        try (Stream<String> lines = response.body()) {
+            lines.forEach(
+                    line -> {
+                        if (!line.isEmpty()) {
+                            lineConsumer.accept(line);
+                        }
+                    });
+        }
+    }
+
    public String get(String path) throws IOException {
        ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
        if (!config.isEnabled()) {
@@ -37,9 +37,10 @@ import stirling.software.common.util.TempFile;
 import stirling.software.common.util.TempFileManager;
 import stirling.software.common.util.ZipExtractionUtils;
 import stirling.software.proprietary.model.api.ai.AiConversationMessage;
+import stirling.software.proprietary.model.api.ai.AiDocumentIngestRequest;
+import stirling.software.proprietary.model.api.ai.AiEngineProgressDetail;
 import stirling.software.proprietary.model.api.ai.AiFile;
-import stirling.software.proprietary.model.api.ai.AiRagIngestRequest;
-import stirling.software.proprietary.model.api.ai.AiRagPageText;
+import stirling.software.proprietary.model.api.ai.AiPageText;
 import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
 import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
 import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
@@ -61,7 +62,7 @@ import tools.jackson.databind.ObjectMapper;
@RequiredArgsConstructor
 public class AiWorkflowService {

-    private static final String RAG_DOCUMENTS_ENDPOINT = "/api/v1/rag/documents";
+    private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents";

    private final CustomPDFDocumentFactory pdfDocumentFactory;
    private final AiEngineClient aiEngineClient;
@@ -77,6 +78,14 @@ public class AiWorkflowService {
    @FunctionalInterface
    public interface ProgressListener {
        void onProgress(AiWorkflowProgressEvent event);
+
+        /**
+         * Called when the engine emits a keep-alive heartbeat. Default is a no-op; consumers that
+         * forward to a downstream connection (e.g. an SSE emitter) override this to push a
+         * heartbeat through, so the next downstream-disconnect surfaces immediately rather than
+         * waiting for the next real progress event.
+         */
+        default void onHeartbeat() {}
    }

    private static final ProgressListener NOOP_LISTENER = event -> {};
@@ -144,7 +153,7 @@ public class AiWorkflowService {
            ProgressListener listener)
            throws IOException {
        listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
-        AiWorkflowResponse response = invokeOrchestrator(request);
+        AiWorkflowResponse response = invokeOrchestrator(request, listener);
        return switch (response.getOutcome()) {
            case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
            case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
@@ -278,22 +287,22 @@ public class AiWorkflowService {
    }

    private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException {
-        List<AiRagPageText> pages = new ArrayList<>();
+        List<AiPageText> pages = new ArrayList<>();
        try (PDDocument document = pdfDocumentFactory.load(multipartFile, true)) {
            int pageCount = document.getNumberOfPages();
            for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
                String pageText = pdfContentExtractor.extractPageTextRaw(document, pageNumber);
                if (pageText != null && !pageText.isBlank()) {
-                    pages.add(new AiRagPageText(pageNumber, pageText));
+                    pages.add(new AiPageText(pageNumber, pageText));
                }
            }
        }
-        AiRagIngestRequest ingestRequest =
-                new AiRagIngestRequest(file.getId(), file.getName(), pages);
+        AiDocumentIngestRequest ingestRequest =
+                new AiDocumentIngestRequest(file.getId(), file.getName(), pages);
        String body = objectMapper.writeValueAsString(ingestRequest);
-        aiEngineClient.postLongRunning(RAG_DOCUMENTS_ENDPOINT, body);
+        aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body);
        log.debug(
-                "Ingested file into RAG: id={}, name={}, pages={}",
+                "Ingested document: id={}, name={}, pages={}",
                file.getId(),
                file.getName(),
                pages.size());
@@ -651,10 +660,57 @@ public class AiWorkflowService {
        return response;
    }

-    private AiWorkflowResponse invokeOrchestrator(WorkflowTurnRequest request) throws IOException {
+    /**
+     * Drive the engine's streaming orchestrator endpoint. Progress events are forwarded to {@code
+     * listener} as they arrive (each one keeps the SSE connection to the frontend alive too). The
+     * final {@code result} event carries the full {@link AiWorkflowResponse}; an {@code error}
+     * event surfaces engine-side failures.
+     */
+    private AiWorkflowResponse invokeOrchestrator(
+            WorkflowTurnRequest request, ProgressListener listener) throws IOException {
        String requestBody = objectMapper.writeValueAsString(request);
-        String responseBody = aiEngineClient.post("/api/v1/orchestrator", requestBody);
-        return objectMapper.readValue(responseBody, AiWorkflowResponse.class);
+        AiWorkflowResponse[] resultHolder = new AiWorkflowResponse[1];
+        String[] errorHolder = new String[1];
+
+        aiEngineClient.streamPost(
+                "/api/v1/orchestrator",
+                requestBody,
+                line -> handleStreamLine(line, listener, resultHolder, errorHolder));
+
+        if (errorHolder[0] != null) {
+            throw new IOException("AI engine returned error: " + errorHolder[0]);
+        }
+        if (resultHolder[0] == null) {
+            throw new IOException("AI engine stream ended without a result");
+        }
+        return resultHolder[0];
+    }
+
+    private void handleStreamLine(
+            String line,
+            ProgressListener listener,
+            AiWorkflowResponse[] resultHolder,
+            String[] errorHolder) {
+        try {
+            JsonNode node = objectMapper.readTree(line);
+            String event = node.path("event").asText();
+            switch (event) {
+                case "progress" -> {
+                    AiEngineProgressDetail detail =
+                            objectMapper.treeToValue(node, AiEngineProgressDetail.class);
+                    listener.onProgress(AiWorkflowProgressEvent.engineProgress(detail));
+                }
+                case "result" -> {
+                    JsonNode response = node.path("response");
+                    resultHolder[0] = objectMapper.treeToValue(response, AiWorkflowResponse.class);
+                }
+                case "error" -> errorHolder[0] = node.path("message").asText("unknown error");
+                case "heartbeat" -> listener.onHeartbeat();
+                default -> log.warn("Ignoring unknown engine stream event: {}", event);
+            }
+        } catch (JacksonException e) {
+            log.warn("Failed to parse engine stream line: {}", line, e);
+        }
    }

    @Data
@@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.lenient;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -22,6 +23,7 @@ import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;

@@ -419,7 +421,7 @@ class AiWorkflowServiceTest {
    }

    @Test
-    void needIngestExtractsPageTextAndPostsToRagThenRetries() throws IOException {
+    void needIngestExtractsPageTextAndPostsThenRetries() throws IOException {
        MockMultipartFile input = pdf("report.pdf", "bytes");
        when(fileIdStrategy.idFor(any())).thenReturn("report-id");

@@ -431,37 +433,64 @@ class AiWorkflowServiceTest {
                .thenReturn("page content");

        int[] orchestratorCalls = {0};
-        when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString()))
-                .thenAnswer(
+        doAnswer(
                        inv -> {
                            orchestratorCalls[0]++;
+                            String responseJson;
                            if (orchestratorCalls[0] == 1) {
-                                return """
-                                       {
-                                         "outcome":"need_ingest",
-                                         "resumeWith":"pdf_question",
-                                         "reason":"ingest first",
-                                         "filesToIngest":[{"id":"report-id","name":"report.pdf"}],
-                                         "contentTypes":["page_text"]
-                                       }
-                                       """;
+                                responseJson =
+                                        """
+                                        {
+                                          "outcome":"need_ingest",
+                                          "resumeWith":"pdf_question",
+                                          "reason":"ingest first",
+                                          "filesToIngest":[{"id":"report-id","name":"report.pdf"}],
+                                          "contentTypes":["page_text"]
+                                        }
+                                        """;
+                            } else {
+                                responseJson =
+                                        """
+                                        {"outcome":"answer","answer":"done","evidence":[]}
+                                        """;
                            }
-                            return """
-                                   {"outcome":"answer","answer":"done","evidence":[]}
-                                   """;
-                        });
+                            Consumer<String> consumer = inv.getArgument(2);
+                            consumer.accept(wrapAsResultEvent(responseJson));
+                            return null;
+                        })
+                .when(aiEngineClient)
+                .streamPost(eq("/api/v1/orchestrator"), anyString(), any());

        AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));

        assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
-        verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/rag/documents"), anyString());
-        verify(aiEngineClient, times(2)).post(eq("/api/v1/orchestrator"), anyString());
+        verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString());
+        verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any());
    }

    // --- helpers ---

    private void stubOrchestrator(String responseJson) throws IOException {
-        when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString())).thenReturn(responseJson);
+        doAnswer(
+                        inv -> {
+                            Consumer<String> consumer = inv.getArgument(2);
+                            consumer.accept(wrapAsResultEvent(responseJson));
+                            return null;
+                        })
+                .when(aiEngineClient)
+                .streamPost(eq("/api/v1/orchestrator"), anyString(), any());
+    }
+
+    /**
+     * Wrap a unary orchestrator response JSON as the NDJSON "result" event the streaming endpoint
+     * emits, so existing tests can keep their per-outcome JSON literals.
+     */
+    private String wrapAsResultEvent(String responseJson) throws IOException {
+        return objectMapper
+                .createObjectNode()
+                .put("event", "result")
+                .set("response", objectMapper.readTree(responseJson))
+                .toString();
    }

    private void stubEndpoint(String endpoint, Resource body) {