diff --git a/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java b/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java index dba7deca2..e8a77bd08 100644 --- a/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java +++ b/app/common/src/main/java/stirling/software/common/model/ApplicationProperties.java @@ -237,6 +237,13 @@ public class ApplicationProperties { private boolean enabled = false; private String url = "http://localhost:5001"; private int timeoutSeconds = 120; + + /** + * Longer timeout for heavy operations like RAG ingestion, which embeds the whole document + * and can take multiple minutes for large books. Applied per-call when the caller + * explicitly requests it via {@code AiEngineClient.postWithTimeout}. + */ + private int longRunningTimeoutSeconds = 600; } @Data diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiFile.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiFile.java new file mode 100644 index 000000000..5b09fe8d5 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiFile.java @@ -0,0 +1,27 @@ +package stirling.software.proprietary.model.api.ai; + +import io.swagger.v3.oas.annotations.media.Schema; + +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +/** + * A file supplied to the AI engine, identified by a stable opaque id plus a display name. + * + *

Values MUST match {@code AiFile} in {@code engine/src/stirling/contracts/common.py}. + */ +@Data +@NoArgsConstructor +@AllArgsConstructor +@Schema(description = "File reference sent to the AI engine") +public class AiFile { + + @Schema( + description = + "Opaque, stable identifier. Owned by Java; used as the RAG collection key.") + private String id; + + @Schema(description = "Original filename, used by agents in user-facing prompts and responses.") + private String name; +} 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 new file mode 100644 index 000000000..ddb9081b6 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagIngestRequest.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/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/AiRagPageText.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagPageText.java new file mode 100644 index 000000000..417b04a38 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiRagPageText.java @@ -0,0 +1,16 @@ +package stirling.software.proprietary.model.api.ai; + +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +/** A single page of extracted text for RAG ingest requests. */ +@Data +@NoArgsConstructor +@AllArgsConstructor +public class AiRagPageText { + + private int pageNumber; + + private String text; +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowEditPlan.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowEditPlan.java deleted file mode 100644 index c41d6450a..000000000 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowEditPlan.java +++ /dev/null @@ -1,35 +0,0 @@ -package stirling.software.proprietary.model.api.ai; - -import java.util.ArrayList; -import java.util.List; -import java.util.Map; - -import io.swagger.v3.oas.annotations.media.Schema; - -import lombok.Data; - -/** - * Embedded plan optionally carried inside a question answer response. When present, the consumer - * (Java) runs the plan steps before delivering the answer; on the resume turn the engine returns - * the real answer using the captured tool reports. - * - *

Mirrors the engine's {@code EditPlanResponse} shape but is nested inside an answer rather than - * acting as the top-level outcome — matches the engine's {@code - * PdfQuestionAnswerResponse.edit_plan} field. - */ -@Data -@Schema(description = "Plan that must run before the answer is final") -public class AiWorkflowEditPlan { - - @Schema(description = "Optional human-readable summary of the plan") - private String summary; - - @Schema(description = "Optional rationale for the plan") - private String rationale; - - @Schema(description = "Tool steps to execute before resuming") - private List> steps = new ArrayList<>(); - - @Schema(description = "AI engine capability to resume with after running the steps") - private String resumeWith; -} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java index f23867028..cd53ba0ac 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java @@ -11,8 +11,8 @@ import lombok.Data; @Schema(description = "Per-file content extraction request from the AI engine") public class AiWorkflowFileRequest { - @Schema(description = "Original filename of the requested file", example = "contract.pdf") - private String fileName; + @Schema(description = "The file the engine wants content extracted for") + private AiFile file; @Schema(description = "Specific 1-based page numbers to extract from this file") private List pageNumbers = new ArrayList<>(); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java index 577f24e5f..ee2d5a584 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java @@ -12,6 +12,7 @@ public enum AiWorkflowOutcome { ANSWER("answer"), NOT_FOUND("not_found"), NEED_CONTENT("need_content"), + NEED_INGEST("need_ingest"), PLAN("plan"), NEED_CLARIFICATION("need_clarification"), CANNOT_DO("cannot_do"), diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java index 5cbb9c9b2..d1cb186f7 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java @@ -73,6 +73,12 @@ public class AiWorkflowResponse { @Schema(description = "Per-file text extraction requests from the AI engine") private List files = new ArrayList<>(); + @Schema( + description = + "Files the AI engine requires to be ingested into RAG before it can continue" + + " the workflow. Populated on need_ingest outcomes.") + private List filesToIngest = new ArrayList<>(); + @Schema(description = "Maximum number of pages the AI engine wants text extracted from") private Integer maxPages; @@ -89,11 +95,4 @@ public class AiWorkflowResponse { + " body or via the X-Stirling-Tool-Report header. May be null for tools" + " that produce only a file.") private JsonNode report; - - @Schema( - description = - "Optional plan attached to an answer outcome. When non-null on outcome=ANSWER," - + " run the plan steps before delivering the answer; the resumed call" - + " produces the real answer.") - private AiWorkflowEditPlan editPlan; } 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 102a7cdb2..fb9963ba8 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 @@ -43,20 +43,36 @@ public class AiEngineClient { public String post(String path, String jsonBody) throws IOException { ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); + return postWithTimeout(path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds())); + } + + /** + * POST with an explicit per-call timeout, for heavy operations (e.g. RAG ingestion of a large + * document) that legitimately take longer than the default timeout. + */ + public String postLongRunning(String path, String jsonBody) throws IOException { + ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); + return postWithTimeout( + path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds())); + } + + private String postWithTimeout(String path, String jsonBody, Duration timeout) + 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; - log.debug("Proxying AI engine request to {}", url); + log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds()); HttpRequest request = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Content-Type", "application/json") .header("Accept", "application/json") - .timeout(Duration.ofSeconds(config.getTimeoutSeconds())) + .timeout(timeout) .POST(HttpRequest.BodyPublishers.ofString(jsonBody)) .build(); 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 5eaf6ecf1..a8d923053 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 @@ -36,7 +36,9 @@ 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.AiWorkflowEditPlan; +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.AiWorkflowFileInput; import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest; import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome; @@ -58,6 +60,8 @@ import tools.jackson.databind.ObjectMapper; @RequiredArgsConstructor public class AiWorkflowService { + private static final String RAG_DOCUMENTS_ENDPOINT = "/api/v1/rag/documents"; + private final CustomPDFDocumentFactory pdfDocumentFactory; private final AiEngineClient aiEngineClient; private final PdfContentExtractor pdfContentExtractor; @@ -66,6 +70,7 @@ public class AiWorkflowService { private final FileStorage fileStorage; private final ToolMetadataService toolMetadataService; private final TempFileManager tempFileManager; + private final FileIdStrategy fileIdStrategy; @FunctionalInterface public interface ProgressListener { @@ -98,15 +103,24 @@ public class AiWorkflowService { throws IOException { validateRequest(request); - Map filesByName = new LinkedHashMap<>(); + // Key by opaque file id, not filename. Filenames aren't guaranteed unique across an + // upload (users can rotate the same 'scan.pdf' twice), and the engine identifies files + // by id in every response shape that asks Java to look a file up again. + Map filesById = new LinkedHashMap<>(); + List files = new ArrayList<>(); for (AiWorkflowFileInput fileInput : request.getFileInputs()) { - filesByName.put( - fileInput.getFileInput().getOriginalFilename(), fileInput.getFileInput()); + MultipartFile multipartFile = fileInput.getFileInput(); + AiFile aiFile = + new AiFile( + fileIdStrategy.idFor(multipartFile), + multipartFile.getOriginalFilename()); + filesById.put(aiFile.getId(), multipartFile); + files.add(aiFile); } WorkflowTurnRequest initialRequest = new WorkflowTurnRequest(); initialRequest.setUserMessage(request.getUserMessage().trim()); - initialRequest.setFileNames(new ArrayList<>(filesByName.keySet())); + initialRequest.setFiles(files); initialRequest.setConversationHistory( request.getConversationHistory() == null ? new ArrayList<>() @@ -116,23 +130,24 @@ public class AiWorkflowService { WorkflowState state = new WorkflowState.Pending(initialRequest); while (state instanceof WorkflowState.Pending pending) { - state = advance(pending.request(), filesByName, listener); + state = advance(pending.request(), filesById, listener); } return ((WorkflowState.Terminal) state).response(); } private WorkflowState advance( WorkflowTurnRequest request, - Map filesByName, + Map filesById, ProgressListener listener) throws IOException { listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE)); AiWorkflowResponse response = invokeOrchestrator(request); return switch (response.getOutcome()) { - case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener); - case TOOL_CALL -> onToolCall(response, filesByName, listener); - case PLAN -> onPlan(response, filesByName, request, listener); - case ANSWER -> onAnswer(response, filesByName, request, listener); + case NEED_CONTENT -> onNeedContent(response, filesById, request, listener); + case NEED_INGEST -> onNeedIngest(response, filesById, request, listener); + case TOOL_CALL -> onToolCall(response, filesById, listener); + case PLAN -> onPlan(response, filesById, request, listener); + case ANSWER -> onAnswer(response, filesById, request, listener); case NOT_FOUND, NEED_CLARIFICATION, CANNOT_DO, @@ -146,7 +161,7 @@ public class AiWorkflowService { private WorkflowState onNeedContent( AiWorkflowResponse response, - Map filesByName, + Map filesById, WorkflowTurnRequest request, ProgressListener listener) throws IOException { @@ -157,43 +172,42 @@ public class AiWorkflowService { List requestedFiles = response.getFiles(); - // Validate requested file names before loading anything + // Validate requested file ids before loading anything if (requestedFiles != null && !requestedFiles.isEmpty()) { for (AiWorkflowFileRequest fileReq : requestedFiles) { - if (!filesByName.containsKey(fileReq.getFileName())) { + AiFile file = fileReq.getFile(); + if (file == null || !filesById.containsKey(file.getId())) { + String display = file == null ? "" : file.getName(); return new WorkflowState.Terminal( - cannotContinue( - "AI engine requested unknown file: " + fileReq.getFileName())); + cannotContinue("AI engine requested unknown file: " + display)); } } } - List fileNamesToLoad = + List filesToLoad = (requestedFiles == null || requestedFiles.isEmpty()) - ? new ArrayList<>(filesByName.keySet()) - : requestedFiles.stream().map(AiWorkflowFileRequest::getFileName).toList(); + ? new ArrayList<>(request.getFiles()) + : requestedFiles.stream().map(AiWorkflowFileRequest::getFile).toList(); - Map requestedByName = + Map requestedById = requestedFiles == null || requestedFiles.isEmpty() ? Map.of() : requestedFiles.stream() - .collect( - Collectors.toMap( - AiWorkflowFileRequest::getFileName, r -> r)); + .collect(Collectors.toMap(r -> r.getFile().getId(), r -> r)); listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT)); List loadedFiles = new ArrayList<>(); try { - for (String fileName : fileNamesToLoad) { - PDDocument doc = pdfDocumentFactory.load(filesByName.get(fileName), true); - loadedFiles.add(new LoadedFile(fileName, doc)); + for (AiFile file : filesToLoad) { + PDDocument doc = pdfDocumentFactory.load(filesById.get(file.getId()), true); + loadedFiles.add(new LoadedFile(file.getId(), file.getName(), doc)); } List contentResults = pdfContentExtractor.extractContent( loadedFiles, - requestedByName, + requestedById, response.getMaxPages(), response.getMaxCharacters()); @@ -201,7 +215,7 @@ public class AiWorkflowService { WorkflowTurnRequest nextRequest = new WorkflowTurnRequest(); nextRequest.setUserMessage(request.getUserMessage()); - nextRequest.setFileNames(request.getFileNames()); + nextRequest.setFiles(request.getFiles()); nextRequest.setConversationHistory(request.getConversationHistory()); nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults)); nextRequest.setResumeWith(response.getResumeWith()); @@ -217,10 +231,74 @@ public class AiWorkflowService { } } + private WorkflowState onNeedIngest( + AiWorkflowResponse response, + Map filesById, + WorkflowTurnRequest request, + ProgressListener listener) + throws IOException { + List filesToIngest = response.getFilesToIngest(); + if (filesToIngest == null || filesToIngest.isEmpty()) { + return new WorkflowState.Terminal( + cannotContinue( + "AI engine returned need_ingest without listing any files to ingest.")); + } + // Guard against a retry loop: if we've already ingested this turn and the engine still + // asks for more, something is wrong on its side. + if (!request.getArtifacts().isEmpty() || request.getResumeWith() != null) { + return new WorkflowState.Terminal( + cannotContinue( + "AI engine requested ingest after the workflow had already been resumed.")); + } + + listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT)); + + for (AiFile file : filesToIngest) { + MultipartFile multipartFile = filesById.get(file.getId()); + if (multipartFile == null) { + return new WorkflowState.Terminal( + cannotContinue( + "AI engine requested ingest for unknown file: " + file.getName())); + } + ingestFile(file, multipartFile); + } + + listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.PROCESSING)); + + WorkflowTurnRequest nextRequest = new WorkflowTurnRequest(); + nextRequest.setUserMessage(request.getUserMessage()); + nextRequest.setFiles(request.getFiles()); + nextRequest.setConversationHistory(request.getConversationHistory()); + nextRequest.setResumeWith(response.getResumeWith()); + return new WorkflowState.Pending(nextRequest); + } + + private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException { + 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)); + } + } + } + AiRagIngestRequest ingestRequest = + new AiRagIngestRequest(file.getId(), file.getName(), pages); + String body = objectMapper.writeValueAsString(ingestRequest); + aiEngineClient.postLongRunning(RAG_DOCUMENTS_ENDPOINT, body); + log.debug( + "Ingested file into RAG: id={}, name={}, pages={}", + file.getId(), + file.getName(), + pages.size()); + } + @SuppressWarnings("unchecked") private WorkflowState onToolCall( AiWorkflowResponse response, - Map filesByName, + Map filesById, ProgressListener listener) { String endpointPath = response.getTool(); Map parameters = response.getParameters(); @@ -233,14 +311,14 @@ public class AiWorkflowService { } try { - List inputFiles = toResources(filesByName); + List inputFiles = toResources(filesById); listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1)); ToolResult result = executeStep(endpointPath, parameters, inputFiles); return new WorkflowState.Terminal( buildCompletedResponse( response.getRationale(), result.files(), - new ArrayList<>(filesByName.keySet()), + inputFileNames(filesById), result.report())); } catch (Exception e) { log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e); @@ -251,36 +329,23 @@ public class AiWorkflowService { private WorkflowState onPlan( AiWorkflowResponse response, - Map filesByName, + Map filesById, WorkflowTurnRequest previousRequest, ProgressListener listener) { return runPlan( response.getSteps(), response.getResumeWith(), response.getSummary(), - filesByName, + filesById, previousRequest, listener); } private WorkflowState onAnswer( AiWorkflowResponse response, - Map filesByName, + Map filesById, WorkflowTurnRequest previousRequest, ProgressListener listener) { - AiWorkflowEditPlan plan = response.getEditPlan(); - if (plan != null) { - // The engine wants us to run a side-quest before the answer is final. - // Run the embedded plan and resume the orchestrator with the captured - // report; the real answer arrives on the resume turn. - return runPlan( - plan.getSteps(), - plan.getResumeWith(), - plan.getSummary(), - filesByName, - previousRequest, - listener); - } return new WorkflowState.Terminal(response); } @@ -289,7 +354,7 @@ public class AiWorkflowService { List> steps, String resumeWith, String summary, - Map filesByName, + Map filesById, WorkflowTurnRequest previousRequest, ProgressListener listener) { if (steps == null || steps.isEmpty()) { @@ -298,7 +363,7 @@ public class AiWorkflowService { } try { - List currentFiles = toResources(filesByName); + List currentFiles = toResources(filesById); // Propagate the *last* non-null report — the terminal step defines the output. JsonNode lastReport = null; String lastReportTool = null; @@ -331,7 +396,7 @@ public class AiWorkflowService { if (resumeWith != null && !resumeWith.isBlank() && lastReport != null) { WorkflowTurnRequest resumeRequest = new WorkflowTurnRequest(); resumeRequest.setUserMessage(previousRequest.getUserMessage()); - resumeRequest.setFileNames(previousRequest.getFileNames()); + resumeRequest.setFiles(previousRequest.getFiles()); resumeRequest.setConversationHistory(previousRequest.getConversationHistory()); resumeRequest.setArtifacts(new ArrayList<>(previousRequest.getArtifacts())); resumeRequest @@ -345,10 +410,7 @@ public class AiWorkflowService { return new WorkflowState.Terminal( buildCompletedResponse( - summary, - currentFiles, - new ArrayList<>(filesByName.keySet()), - lastReport)); + summary, currentFiles, inputFileNames(filesById), lastReport)); } catch (Exception e) { log.error("Failed to execute plan: {}", e.getMessage(), e); return new WorkflowState.Terminal( @@ -356,6 +418,10 @@ public class AiWorkflowService { } } + private static List inputFileNames(Map filesById) { + return filesById.values().stream().map(MultipartFile::getOriginalFilename).toList(); + } + /** * Execute a single tool step. If the endpoint accepts multiple files, all files are sent in one * call. Otherwise, the endpoint is called once per file. ZIP responses are unpacked so each @@ -460,9 +526,9 @@ public class AiWorkflowService { } } - private List toResources(Map filesByName) throws IOException { + private List toResources(Map filesById) throws IOException { List resources = new ArrayList<>(); - for (MultipartFile file : filesByName.values()) { + for (MultipartFile file : filesById.values()) { TempFile tempFile = tempFileManager.createManagedTempFile("ai-workflow"); file.transferTo(tempFile.getPath()); final String originalName = Filenames.toSimpleFileName(file.getOriginalFilename()); @@ -554,7 +620,7 @@ public class AiWorkflowService { @Data private static class WorkflowTurnRequest { private String userMessage; - private List fileNames = new ArrayList<>(); + private List files = new ArrayList<>(); private List conversationHistory = new ArrayList<>(); private List artifacts = new ArrayList<>(); private String resumeWith; diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/ByteHashFileIdStrategy.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/ByteHashFileIdStrategy.java new file mode 100644 index 000000000..a30d7bc56 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/ByteHashFileIdStrategy.java @@ -0,0 +1,53 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; +import java.io.InputStream; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; + +import org.springframework.stereotype.Component; +import org.springframework.web.multipart.MultipartFile; + +/** + * Content-addressable id derived from the SHA-256 hash of the uploaded bytes. Same content always + * hashes to the same id, so re-uploads dedupe naturally in RAG. Suitable for session and SaaS + * deployments; a folder-watch deployment would use a different strategy keyed by path. + */ +@Component +public class ByteHashFileIdStrategy implements FileIdStrategy { + + /** + * Hex-char length of the returned id. 16 chars = 64 bits of collision space, which is plenty + * for per-user document sets and keeps RAG collection names short. + */ + private static final int ID_HEX_LENGTH = 16; + + private static final int BUFFER_SIZE = 64 * 1024; + + @Override + public String idFor(MultipartFile file) throws IOException { + MessageDigest digest = sha256(); + try (InputStream in = file.getInputStream()) { + byte[] buffer = new byte[BUFFER_SIZE]; + int read; + while ((read = in.read(buffer)) != -1) { + digest.update(buffer, 0, read); + } + } + byte[] hash = digest.digest(); + StringBuilder hex = new StringBuilder(ID_HEX_LENGTH); + for (int i = 0; hex.length() < ID_HEX_LENGTH; i++) { + hex.append(String.format("%02x", hash[i])); + } + return hex.toString(); + } + + private static MessageDigest sha256() { + try { + return MessageDigest.getInstance("SHA-256"); + } catch (NoSuchAlgorithmException e) { + // SHA-256 is mandated by the JDK; absent only if the platform is broken. + throw new IllegalStateException("SHA-256 not available", e); + } + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/FileIdStrategy.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/FileIdStrategy.java new file mode 100644 index 000000000..b5e616935 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/FileIdStrategy.java @@ -0,0 +1,16 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; + +import org.springframework.web.multipart.MultipartFile; + +/** + * Produces stable identifiers for uploaded files. The identifier is opaque to the AI engine and + * serves as the RAG collection key when content is ingested. Swapping implementations (content + * hash, filesystem path, tenant-scoped id, etc.) is how the system adapts to different deployment + * models without any engine-side change. + */ +public interface FileIdStrategy { + + String idFor(MultipartFile file) throws IOException; +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java b/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java index c73a64af5..8e89cada3 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java @@ -44,7 +44,11 @@ public class PdfContentExtractor { private static final int TEXT_PRESENCE_THRESHOLD = 20; - record LoadedFile(String fileName, PDDocument document) {} + /** + * A loaded PDF alongside the opaque file id used by the AI engine as its RAG collection key. + * Keyed by id (not name) because filenames aren't unique across an upload. + */ + record LoadedFile(String id, String fileName, PDDocument document) {} // ----------------------------------------------------------------------- // Low-level extraction methods (usable by any agent) @@ -126,7 +130,7 @@ public class PdfContentExtractor { */ List extractContent( List loadedFiles, - Map requestedByName, + Map requestedById, int maxPages, int maxCharacters) throws IOException { @@ -136,7 +140,7 @@ public class PdfContentExtractor { for (LoadedFile lf : loadedFiles) { if (remainingPages <= 0 || remainingCharacters <= 0) break; - AiWorkflowFileRequest fileReq = requestedByName.get(lf.fileName()); + AiWorkflowFileRequest fileReq = requestedById.get(lf.id()); List contentTypes = fileReq != null && !fileReq.getContentTypes().isEmpty() ? fileReq.getContentTypes() 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 2872e0db9..620c62e5d 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 @@ -3,8 +3,11 @@ package stirling.software.proprietary.service; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; +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.lenient; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -21,6 +24,8 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.zip.ZipEntry; import java.util.zip.ZipOutputStream; +import org.apache.pdfbox.pdmodel.PDDocument; +import org.apache.pdfbox.pdmodel.PDPage; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -32,6 +37,7 @@ import org.springframework.core.io.Resource; import org.springframework.http.ResponseEntity; import org.springframework.mock.web.MockMultipartFile; import org.springframework.util.MultiValueMap; +import org.springframework.web.multipart.MultipartFile; import stirling.software.common.model.ApplicationProperties; import stirling.software.common.service.CustomPDFDocumentFactory; @@ -72,6 +78,7 @@ class AiWorkflowServiceTest { @Mock private InternalApiClient internalApiClient; @Mock private FileStorage fileStorage; @Mock private ToolMetadataService toolMetadataService; + @Mock private FileIdStrategy fileIdStrategy; @TempDir Path tempDir; @@ -80,13 +87,19 @@ class AiWorkflowServiceTest { private AiWorkflowService service; @BeforeEach - void setUp() { + void setUp() throws IOException { ApplicationProperties props = new ApplicationProperties(); props.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString()); props.getSystem().getTempFileManagement().setPrefix("ai-test-"); tempFileManager = new TempFileManager(new TempFileRegistry(), props); objectMapper = JsonMapper.builder().build(); + // Mock strategy yields the filename as id so each MockMultipartFile in a test gets a + // distinct collection key. Real strategy (ByteHashFileIdStrategy) hashes bytes. + lenient() + .when(fileIdStrategy.idFor(any(MultipartFile.class))) + .thenAnswer(inv -> ((MultipartFile) inv.getArgument(0)).getOriginalFilename()); + service = new AiWorkflowService( pdfDocumentFactory, @@ -96,7 +109,8 @@ class AiWorkflowServiceTest { internalApiClient, fileStorage, toolMetadataService, - tempFileManager); + tempFileManager, + fileIdStrategy); } @Test @@ -246,6 +260,46 @@ class AiWorkflowServiceTest { verify(internalApiClient, never()).post(anyString(), any()); } + @Test + void needIngestExtractsPageTextAndPostsToRagThenRetries() throws IOException { + MockMultipartFile input = pdf("report.pdf", "bytes"); + when(fileIdStrategy.idFor(any())).thenReturn("report-id"); + + PDDocument document = new PDDocument(); + document.addPage(new PDPage()); + document.addPage(new PDPage()); + when(pdfDocumentFactory.load(any(MultipartFile.class), anyBoolean())).thenReturn(document); + when(pdfContentExtractor.extractPageTextRaw(eq(document), anyInt())) + .thenReturn("page content"); + + int[] orchestratorCalls = {0}; + when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString())) + .thenAnswer( + inv -> { + orchestratorCalls[0]++; + if (orchestratorCalls[0] == 1) { + return """ + { + "outcome":"need_ingest", + "resumeWith":"pdf_question", + "reason":"ingest first", + "filesToIngest":[{"id":"report-id","name":"report.pdf"}], + "contentTypes":["page_text"] + } + """; + } + return """ + {"outcome":"answer","answer":"done","evidence":[]} + """; + }); + + 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()); + } + // --- helpers --- private void stubOrchestrator(String responseJson) throws IOException { diff --git a/engine/.env b/engine/.env index 7c5d2b4a2..18dcfa99d 100644 --- a/engine/.env +++ b/engine/.env @@ -30,7 +30,12 @@ STIRLING_RAG_PGVECTOR_DSN= STIRLING_RAG_CHUNK_SIZE=512 STIRLING_RAG_CHUNK_OVERLAP=64 -STIRLING_RAG_TOP_K=5 +STIRLING_RAG_TOP_K=20 + +# Per-run cap on ``search_knowledge`` calls. After this many calls the tool is +# removed from the agent's toolset so it must answer from what it already retrieved +# rather than chain more searches. +STIRLING_RAG_MAX_SEARCHES=5 # Upper bounds on PDF page text the engine will request per extraction round. STIRLING_MAX_PAGES=200 diff --git a/engine/src/stirling/agents/orchestrator.py b/engine/src/stirling/agents/orchestrator.py index b6eb11ff3..6dedd5079 100644 --- a/engine/src/stirling/agents/orchestrator.py +++ b/engine/src/stirling/agents/orchestrator.py @@ -18,10 +18,11 @@ from stirling.contracts import ( OrchestratorRequest, OrchestratorResponse, PdfEditResponse, - PdfQuestionResponse, + PdfQuestionOrchestrateResponse, SupportedCapability, UnsupportedCapabilityResponse, format_conversation_history, + format_file_names, ) from stirling.contracts.pdf_edit import EditPlanResponse from stirling.services import AppRuntime @@ -78,12 +79,13 @@ class OrchestratorAgent: "You are the top-level orchestrator. " "Choose exactly one output function that best handles the request. " "Use delegate_pdf_edit for requested modifications of single or multiple PDFs. " - "Use delegate_pdf_question for questions about PDF contents. " + "Use delegate_pdf_question for questions about the contents of the attached PDFs. " "Use delegate_user_spec for requests to create or define an agent spec. " "Use delegate_pdf_review when the user wants the PDF returned with review" " comments attached — anything like 'review this', 'annotate with comments'," " 'leave feedback on the PDF'. " - "Use unsupported_capability only when none of the other outputs fit." + "Use unsupported_capability when the user asks about the assistant itself " + "or when none of the other outputs fit; supply a helpful message." ), model_settings=runtime.fast_model_settings, ) @@ -91,7 +93,7 @@ class OrchestratorAgent: async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse: logger.info( "[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r", - request.file_names, + [file.name for file in request.files], request.resume_with, [type(a).__name__ for a in request.artifacts], request.user_message, @@ -137,10 +139,10 @@ class OrchestratorAgent: async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse: return await PdfEditAgent(self.runtime).orchestrate(request) - async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse: + async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionOrchestrateResponse: return await self._run_pdf_question(ctx.deps.request) - async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse: + async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse: return await PdfQuestionAgent(self.runtime).orchestrate(request) async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse: @@ -165,12 +167,11 @@ class OrchestratorAgent: def _build_prompt(self, request: OrchestratorRequest) -> str: artifact_summary = self._describe_artifacts(request) - file_names = ", ".join(request.file_names) if request.file_names else "Unknown files" history = format_conversation_history(request.conversation_history) return ( f"Conversation history:\n{history}\n" f"User message: {request.user_message}\n" - f"Files: {file_names}\n" + f"Files: {format_file_names(request.files)}\n" f"Available artifacts:\n{artifact_summary}" ) diff --git a/engine/src/stirling/agents/pdf_edit.py b/engine/src/stirling/agents/pdf_edit.py index 54bc965ac..aecea50aa 100644 --- a/engine/src/stirling/agents/pdf_edit.py +++ b/engine/src/stirling/agents/pdf_edit.py @@ -22,6 +22,7 @@ from stirling.contracts import ( SupportedCapability, ToolOperationStep, format_conversation_history, + format_file_names, ) from stirling.logging import Pretty from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint @@ -116,7 +117,6 @@ class PdfEditParameterSelector: ) -> str: operation_id = operation_plan[operation_index] operation_list = ", ".join(operation.name for operation in operation_plan) - file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided." generated_steps_text = ( "\n".join( f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps) @@ -127,7 +127,7 @@ class PdfEditParameterSelector: return ( f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n" f"User request: {request.user_message}\n" - f"Files: {file_names}\n" + f"Files: {format_file_names(request.files)}\n" f"Operation plan: {operation_list}\n" f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n" f"Selected operation: {operation_id.name}\n" @@ -153,7 +153,7 @@ class PdfEditAgent: return await self.handle( PdfEditRequest( user_message=request.user_message, - file_names=request.file_names, + files=request.files, conversation_history=request.conversation_history, page_text=extracted_text.files if extracted_text is not None else [], ) @@ -166,7 +166,7 @@ class PdfEditAgent: async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse: logger.info( "[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r", - request.file_names, + [file.name for file in request.files], has_page_text(request.page_text), allow_need_content, request.user_message, @@ -225,11 +225,10 @@ class PdfEditAgent: ) def _build_selection_prompt(self, request: PdfEditRequest) -> str: - file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided." return ( f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n" f"User request: {request.user_message}\n" - f"Files: {file_names}\n" + f"Files: {format_file_names(request.files)}\n" f"Supported operations: {self._supported_operations_prompt()}\n" f"Extracted page text:\n{format_page_text(request.page_text)}" ) @@ -243,8 +242,7 @@ class PdfEditAgent: request: PdfEditRequest, ) -> NeedContentResponse: files = selection.files or [ - NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT]) - for file_name in request.file_names + NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in request.files ] return NeedContentResponse( resume_with=SupportedCapability.PDF_EDIT, diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index b4e520f48..c27b7e612 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -1,32 +1,67 @@ from __future__ import annotations +import logging + from pydantic_ai import Agent from pydantic_ai.output import NativeOutput -from stirling.agents._page_text import ( - format_page_text, - get_extracted_text_artifact, - has_page_text, -) from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict from stirling.contracts import ( + AiFile, EditPlanResponse, - NeedContentFileRequest, - NeedContentResponse, + NeedIngestResponse, OrchestratorRequest, PdfContentType, PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse, + PdfQuestionOrchestrateResponse, PdfQuestionRequest, PdfQuestionResponse, + PdfQuestionTerminalResponse, SupportedCapability, ToolOperationStep, Verdict, format_conversation_history, + format_file_names, ) 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" + "\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" + "\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" + "\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" + "- 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." +) + _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 " @@ -41,26 +76,6 @@ _MATH_SYNTH_SYSTEM_PROMPT = ( class PdfQuestionAgent: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime - rag = runtime.rag_capability - self.agent = Agent( - model=runtime.smart_model, - output_type=NativeOutput( - [ - PdfQuestionAnswerResponse, - PdfQuestionNotFoundResponse, - ] - ), - system_prompt=( - "Answer questions about PDFs using only the extracted page text provided in the prompt. " - "Do not guess or use outside knowledge. " - "If the answer is not supported by the provided text, return not_found. " - "When answering, include a short list of evidence snippets with their page numbers. " - "Reply in the SAME LANGUAGE as the question." - ), - instructions=rag.instructions, - toolsets=[rag.toolset], - model_settings=runtime.smart_model_settings, - ) self._math_synth_agent: Agent[None, str] = Agent( model=runtime.fast_model, output_type=str, @@ -70,30 +85,31 @@ class PdfQuestionAgent: self._math_intent_classifier = MathIntentClassifier(runtime) async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse: - if not has_page_text(request.page_text): - return NeedContentResponse( + logger.info( + "[pdf-question] handle: files=%s question=%r", + [file.name for file in request.files], + request.question, + ) + missing = await self._find_missing_files(request.files) + if missing: + logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing]) + return NeedIngestResponse( resume_with=SupportedCapability.PDF_QUESTION, - reason="No extracted PDF page text was provided, so the question cannot be answered yet.", - files=[ - NeedContentFileRequest( - file_name=file_name, - content_types=[PdfContentType.PAGE_TEXT], - ) - for file_name in request.file_names - ], - max_pages=self.runtime.settings.max_pages, - max_characters=self.runtime.settings.max_characters, + reason="Some files have not been ingested into RAG yet.", + files_to_ingest=missing, + content_types=[PdfContentType.PAGE_TEXT], ) return await self._run_answer_agent(request) - async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionResponse: + async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse: """Entry point for the orchestrator delegate. Decides math intent locally via a small classifier LLM (language-agnostic). - On a math first turn, embeds an :class:`EditPlanResponse` in the answer - response; on the resume turn, digests the captured :class:`Verdict` into - a localised prose answer. Non-math first turns fall through to the - text-grounded :meth:`handle` pipeline. + On a math first turn, returns an :class:`EditPlanResponse` (``outcome=PLAN``) + with ``resume_with=PDF_QUESTION`` so the caller runs the math specialist + and re-invokes the orchestrator. On the resume turn, the captured + :class:`Verdict` is digested into a localised prose answer. Non-math + first turns fall through to the text-grounded :meth:`handle` pipeline. """ verdict = extract_math_verdict(request) if verdict is not None: @@ -104,36 +120,53 @@ class PdfQuestionAgent: return PdfQuestionAnswerResponse(answer=answer, evidence=[]) if await self._math_intent_classifier.classify(request.user_message): - # First turn — ask the caller to run the math specialist and come back. - # The plan rides on the answer response as a nullable member; ``answer`` - # is empty on this turn and the caller resumes once the plan is run. - return PdfQuestionAnswerResponse( - answer="", - evidence=[], - edit_plan=EditPlanResponse( - summary="", - steps=[ - ToolOperationStep( - tool=AgentToolId.MATH_AUDITOR_AGENT, - parameters=MathAuditorAgentParams(), - ) - ], - resume_with=SupportedCapability.PDF_QUESTION, - ), + # First turn — emit a one-step plan calling the math specialist, + # with resume_with set so the caller comes back with the verdict + # in artifacts (handled by the resume branch above). + return EditPlanResponse( + summary="", + steps=[ + ToolOperationStep( + tool=AgentToolId.MATH_AUDITOR_AGENT, + parameters=MathAuditorAgentParams(), + ) + ], + resume_with=SupportedCapability.PDF_QUESTION, ) - extracted_text = get_extracted_text_artifact(request) return await self.handle( PdfQuestionRequest( question=request.user_message, - file_names=request.file_names, - page_text=extracted_text.files if extracted_text is not None else [], + files=request.files, conversation_history=request.conversation_history, ) ) - async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse: - result = await self.agent.run(self._build_prompt(request)) + 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): + missing.append(file) + return missing + + async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse: + rag = RagCapability( + rag_service=self.runtime.rag_service, + 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, + ) + agent = Agent( + model=self.runtime.smart_model, + output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]), + system_prompt=PDF_QUESTION_SYSTEM_PROMPT, + instructions=rag.instructions, + toolsets=[rag.toolset], + model_settings=self.runtime.smart_model_settings, + ) + prompt = self._build_prompt(request) + logger.debug("[pdf-question] prompt:\n%s", prompt) + result = await agent.run(prompt) return result.output async def _synthesise_math_answer(self, user_message: str, verdict: Verdict) -> str: @@ -146,12 +179,10 @@ class PdfQuestionAgent: return result.output def _build_prompt(self, request: PdfQuestionRequest) -> str: - file_names = ", ".join(request.file_names) if request.file_names else "Unknown files" - pages = format_page_text(request.page_text, empty="") history = format_conversation_history(request.conversation_history) return ( f"Conversation history:\n{history}\n" - f"Files: {file_names}\n" + f"Files: {format_file_names(request.files)}\n" f"Question: {request.question}\n" - f"Extracted page text:\n{pages}" + "Use search_knowledge to retrieve the relevant content, then answer." ) diff --git a/engine/src/stirling/api/app.py b/engine/src/stirling/api/app.py index c5ffac1b1..97cc700da 100644 --- a/engine/src/stirling/api/app.py +++ b/engine/src/stirling/api/app.py @@ -7,7 +7,13 @@ from fastapi import Depends, FastAPI from pydantic_ai import Agent from pydantic_ai.models.instrumented import InstrumentationSettings -from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent +from stirling.agents import ( + ExecutionPlanningAgent, + OrchestratorAgent, + PdfEditAgent, + PdfQuestionAgent, + UserSpecAgent, +) from stirling.agents.ledger import MathAuditorAgent from stirling.agents.pdf_comment import PdfCommentAgent from stirling.api.middleware import UserIdMiddleware @@ -51,6 +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() if tracer_provider: tracer_provider.shutdown() diff --git a/engine/src/stirling/api/dependencies.py b/engine/src/stirling/api/dependencies.py index 88b24a309..158f67cd6 100644 --- a/engine/src/stirling/api/dependencies.py +++ b/engine/src/stirling/api/dependencies.py @@ -2,7 +2,13 @@ from __future__ import annotations from fastapi import Request -from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent +from stirling.agents import ( + ExecutionPlanningAgent, + OrchestratorAgent, + PdfEditAgent, + PdfQuestionAgent, + UserSpecAgent, +) from stirling.agents.ledger import MathAuditorAgent from stirling.agents.pdf_comment import PdfCommentAgent from stirling.rag import RagService @@ -37,10 +43,6 @@ def get_rag_service(request: Request) -> RagService: return request.app.state.runtime.rag_service -def get_rag_embedding_model(request: Request) -> str: - return request.app.state.runtime.settings.rag_embedding_model - - def get_math_auditor_agent(request: Request) -> MathAuditorAgent: return request.app.state.math_auditor_agent diff --git a/engine/src/stirling/api/routes/rag.py b/engine/src/stirling/api/routes/rag.py index 2eb9acfcc..8e36ca8a3 100644 --- a/engine/src/stirling/api/routes/rag.py +++ b/engine/src/stirling/api/routes/rag.py @@ -4,75 +4,60 @@ from typing import Annotated from fastapi import APIRouter, Depends -from stirling.api.dependencies import get_rag_embedding_model, get_rag_service +from stirling.api.dependencies import get_rag_service from stirling.contracts import ( - RagCollectionsResponse, - RagDeleteCollectionResponse, - RagIndexRequest, - RagIndexResponse, - RagSearchRequest, - RagSearchResponse, - RagSearchResultItem, - RagStatusResponse, + DeleteDocumentResponse, + IngestDocumentRequest, + IngestDocumentResponse, + PdfContentType, ) -from stirling.rag import RagService +from stirling.models import FileId +from stirling.rag import Document, RagService router = APIRouter(prefix="/api/v1/rag", tags=["rag"]) -@router.get("/status", response_model=RagStatusResponse) -async def rag_status( +@router.post("/documents", response_model=IngestDocumentResponse) +async def ingest_document( + request: IngestDocumentRequest, rag: Annotated[RagService, Depends(get_rag_service)], - embedding_model: Annotated[str, Depends(get_rag_embedding_model)], -) -> RagStatusResponse: - collections = await rag.list_collections() - return RagStatusResponse(embedding_model=embedding_model, collections=collections) +) -> 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.post("/index", response_model=RagIndexResponse) -async def rag_index( - request: RagIndexRequest, +@router.delete("/documents/{document_id}", response_model=DeleteDocumentResponse) +async def delete_document( + document_id: FileId, rag: Annotated[RagService, Depends(get_rag_service)], -) -> RagIndexResponse: - count = await rag.index_text( - collection=request.collection, - text=request.text, - source=request.source, - metadata=request.metadata, - ) - return RagIndexResponse(collection=request.collection, chunks_indexed=count) - - -@router.post("/search", response_model=RagSearchResponse) -async def rag_search( - request: RagSearchRequest, - rag: Annotated[RagService, Depends(get_rag_service)], -) -> RagSearchResponse: - results = await rag.search(query=request.query, collection=request.collection, top_k=request.top_k) - items = [ - RagSearchResultItem( - text=r.document.text, - source=r.document.metadata.get("source", ""), - chunk_id=r.document.metadata.get("chunk_index", ""), - score=r.score, - ) - for r in results - ] - return RagSearchResponse(query=request.query, results=items) - - -@router.get("/collections", response_model=RagCollectionsResponse) -async def rag_collections( - rag: Annotated[RagService, Depends(get_rag_service)], -) -> RagCollectionsResponse: - collections = await rag.list_collections() - return RagCollectionsResponse(collections=collections) - - -@router.delete("/collections/{name}", response_model=RagDeleteCollectionResponse) -async def rag_delete_collection( - name: str, - rag: Annotated[RagService, Depends(get_rag_service)], -) -> RagDeleteCollectionResponse: - await rag.delete_collection(name) - return RagDeleteCollectionResponse(status="deleted", collection=name) +) -> 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 2aca8d3a0..9087307a6 100644 --- a/engine/src/stirling/config/settings.py +++ b/engine/src/stirling/config/settings.py @@ -36,6 +36,7 @@ class AppSettings(BaseSettings): rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE") rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") + rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES") max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES") max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS") diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 300c5c586..d6657ed3e 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -10,12 +10,14 @@ from .agent_drafts import ( from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep from .comments import CommentSpec from .common import ( + AiFile, ArtifactKind, ConversationMessage, ExtractedFileText, MathAuditorToolReportArtifact, NeedContentFileRequest, NeedContentResponse, + NeedIngestResponse, PdfContentType, PdfTextSelection, StepKind, @@ -24,6 +26,7 @@ from .common import ( ToolReportArtifact, WorkflowOutcome, format_conversation_history, + format_file_names, ) from .execution import ( AgentExecutionRequest, @@ -71,24 +74,20 @@ from .pdf_edit import ( from .pdf_questions import ( PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse, + PdfQuestionOrchestrateResponse, PdfQuestionRequest, PdfQuestionResponse, PdfQuestionTerminalResponse, ) from .rag import ( - MAX_INDEX_TEXT_LENGTH, - RagCollectionsResponse, - RagDeleteCollectionResponse, - RagIndexRequest, - RagIndexResponse, - RagSearchRequest, - RagSearchResponse, - RagSearchResultItem, - RagStatusResponse, + DeleteDocumentResponse, + IngestDocumentRequest, + IngestDocumentResponse, + IngestedPageText, ) __all__ = [ - "MAX_INDEX_TEXT_LENGTH", + "AiFile", "AgentDraft", "AgentDraftRequest", "AgentDraftResponse", @@ -105,6 +104,7 @@ __all__ = [ "CommentSpec", "CompletedExecutionAction", "ConversationMessage", + "DeleteDocumentResponse", "Discrepancy", "DiscrepancyKind", "EditCannotDoResponse", @@ -119,10 +119,15 @@ __all__ = [ "FolioManifest", "FolioType", "format_conversation_history", + "format_file_names", "HealthResponse", + "IngestDocumentRequest", + "IngestDocumentResponse", + "IngestedPageText", "MathAuditorToolReportArtifact", "NeedContentFileRequest", "NeedContentResponse", + "NeedIngestResponse", "NextExecutionAction", "OrchestratorRequest", "OrchestratorResponse", @@ -136,18 +141,11 @@ __all__ = [ "PdfEditTerminalResponse", "PdfQuestionAnswerResponse", "PdfQuestionNotFoundResponse", + "PdfQuestionOrchestrateResponse", "PdfQuestionRequest", "PdfQuestionResponse", "PdfQuestionTerminalResponse", "PdfTextSelection", - "RagCollectionsResponse", - "RagDeleteCollectionResponse", - "RagIndexRequest", - "RagIndexResponse", - "RagSearchRequest", - "RagSearchResponse", - "RagSearchResultItem", - "RagStatusResponse", "Requisition", "Severity", "StepKind", diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 69fa813ca..21e618ab1 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -6,7 +6,7 @@ from typing import Literal, assert_never from pydantic import Field, model_validator from stirling.contracts.ledger import Verdict -from stirling.models import OPERATIONS, ApiModel, ToolEndpoint +from stirling.models import OPERATIONS, ApiModel, FileId, ToolEndpoint from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId @@ -50,6 +50,7 @@ class WorkflowOutcome(StrEnum): ANSWER = "answer" NEED_CONTENT = "need_content" + NEED_INGEST = "need_ingest" NOT_FOUND = "not_found" PLAN = "plan" NEED_CLARIFICATION = "need_clarification" @@ -94,12 +95,30 @@ class ConversationMessage(ApiModel): content: str +class AiFile(ApiModel): + """A file the user has supplied, identified by both a stable id and a display name. + + The id is opaque to the engine: Java generates it (content hash, file path, UUID, etc.) + and the engine uses it as the RAG collection key for any agent that indexes content. + The name is used in user-facing prompts and responses. + """ + + id: FileId = Field(min_length=1) + name: str = Field(min_length=1) + + def format_conversation_history(conversation_history: list[ConversationMessage]) -> str: if not conversation_history: return "None" return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history) +def format_file_names(files: list[AiFile]) -> str: + if not files: + return "No file names were provided." + return ", ".join(file.name for file in files) + + class PdfTextSelection(ApiModel): page_number: int | None = None text: str @@ -111,7 +130,7 @@ class ExtractedFileText(ApiModel): class NeedContentFileRequest(ApiModel): - file_name: str + file: AiFile page_numbers: list[int] = Field(default_factory=list) content_types: list[PdfContentType] @@ -146,6 +165,20 @@ class MathAuditorToolReportArtifact(ApiModel): ToolReportArtifact = MathAuditorToolReportArtifact +class NeedIngestResponse(ApiModel): + """Signal that the listed files must be ingested into RAG 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. + """ + + outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST + resume_with: SupportedCapability + reason: str + files_to_ingest: list[AiFile] + content_types: list[PdfContentType] = Field(default_factory=list) + + class ToolOperationStep(ApiModel): kind: Literal[StepKind.TOOL] = StepKind.TOOL tool: AnyToolId diff --git a/engine/src/stirling/contracts/orchestrator.py b/engine/src/stirling/contracts/orchestrator.py index fac85a923..dc7c1f152 100644 --- a/engine/src/stirling/contracts/orchestrator.py +++ b/engine/src/stirling/contracts/orchestrator.py @@ -8,10 +8,12 @@ from stirling.models import ApiModel from .agent_drafts import AgentDraftResponse from .common import ( + AiFile, ArtifactKind, ConversationMessage, ExtractedFileText, NeedContentResponse, + NeedIngestResponse, SupportedCapability, ToolReportArtifact, WorkflowOutcome, @@ -31,7 +33,7 @@ WorkflowArtifact = Annotated[ExtractedTextArtifact | ToolReportArtifact, Field(d class OrchestratorRequest(ApiModel): user_message: str - file_names: list[str] + files: list[AiFile] = Field(default_factory=list) conversation_history: list[ConversationMessage] = Field(default_factory=list) artifacts: list[WorkflowArtifact] = Field(default_factory=list) resume_with: SupportedCapability | None = None @@ -47,6 +49,7 @@ type OrchestratorResponse = Annotated[ PdfEditTerminalResponse | PdfQuestionTerminalResponse | NeedContentResponse + | NeedIngestResponse | AgentDraftResponse | NextExecutionAction | UnsupportedCapabilityResponse, diff --git a/engine/src/stirling/contracts/pdf_edit.py b/engine/src/stirling/contracts/pdf_edit.py index 40d9f95c0..be79bf563 100644 --- a/engine/src/stirling/contracts/pdf_edit.py +++ b/engine/src/stirling/contracts/pdf_edit.py @@ -7,6 +7,7 @@ from pydantic import Field from stirling.models import ApiModel from .common import ( + AiFile, ConversationMessage, ExtractedFileText, NeedContentResponse, @@ -18,7 +19,7 @@ from .common import ( class PdfEditRequest(ApiModel): user_message: str - file_names: list[str] = Field(default_factory=list) + files: list[AiFile] = Field(default_factory=list) conversation_history: list[ConversationMessage] = Field(default_factory=list) page_text: list[ExtractedFileText] = Field(default_factory=list) diff --git a/engine/src/stirling/contracts/pdf_questions.py b/engine/src/stirling/contracts/pdf_questions.py index a987dc5d2..873164351 100644 --- a/engine/src/stirling/contracts/pdf_questions.py +++ b/engine/src/stirling/contracts/pdf_questions.py @@ -7,9 +7,10 @@ from pydantic import Field from stirling.models import ApiModel from .common import ( + AiFile, ConversationMessage, ExtractedFileText, - NeedContentResponse, + NeedIngestResponse, WorkflowOutcome, ) from .pdf_edit import EditPlanResponse @@ -17,8 +18,7 @@ from .pdf_edit import EditPlanResponse class PdfQuestionRequest(ApiModel): question: str - page_text: list[ExtractedFileText] = Field(default_factory=list) - file_names: list[str] + files: list[AiFile] = Field(default_factory=list) conversation_history: list[ConversationMessage] = Field(default_factory=list) @@ -26,15 +26,6 @@ class PdfQuestionAnswerResponse(ApiModel): outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER answer: str evidence: list[ExtractedFileText] = Field(default_factory=list) - edit_plan: EditPlanResponse | None = Field( - default=None, - description=( - "Optional plan the caller must run before the answer is final. When" - " populated, ``answer`` is empty on this turn — the caller executes" - " the plan and re-invokes the orchestrator with ``resume_with`` set" - " to PDF_QUESTION; the real answer arrives on the resume turn." - ), - ) class PdfQuestionNotFoundResponse(ApiModel): @@ -44,6 +35,14 @@ class PdfQuestionNotFoundResponse(ApiModel): type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse type PdfQuestionResponse = Annotated[ - PdfQuestionTerminalResponse | NeedContentResponse, + PdfQuestionTerminalResponse | NeedIngestResponse, Field(discriminator="outcome"), ] + + +# ``orchestrate`` may also emit an ``EditPlanResponse`` on the math-routing +# first turn (``outcome=PLAN`` with ``resume_with=PDF_QUESTION``). It's not in +# ``PdfQuestionTerminalResponse`` because that alias would otherwise duplicate +# the PLAN branch already provided by ``PdfEditTerminalResponse`` in the +# top-level :class:`OrchestratorResponse` discriminated union. +type PdfQuestionOrchestrateResponse = PdfQuestionResponse | EditPlanResponse diff --git a/engine/src/stirling/contracts/rag.py b/engine/src/stirling/contracts/rag.py index c4cea35af..9f655bbb4 100644 --- a/engine/src/stirling/contracts/rag.py +++ b/engine/src/stirling/contracts/rag.py @@ -4,48 +4,36 @@ from pydantic import Field from stirling.models import ApiModel -MAX_INDEX_TEXT_LENGTH = 1_000_000 # 1MB text limit per index request +from .common import FileId -class RagStatusResponse(ApiModel): - embedding_model: str - collections: list[str] +class IngestedPageText(ApiModel): + page_number: int = Field(ge=1) + text: str -class RagIndexRequest(ApiModel): - collection: str = Field(min_length=1) - text: str = Field(max_length=MAX_INDEX_TEXT_LENGTH) - source: str = "" - metadata: dict[str, str] = Field(default_factory=dict) +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 RagIndexResponse(ApiModel): - collection: str +class IngestDocumentResponse(ApiModel): + document_id: FileId chunks_indexed: int -class RagSearchRequest(ApiModel): - query: str - collection: str | None = Field(default=None, min_length=1) - top_k: int = 5 - - -class RagSearchResultItem(ApiModel): - text: str - source: str - chunk_id: str - score: float - - -class RagSearchResponse(ApiModel): - query: str - results: list[RagSearchResultItem] - - -class RagCollectionsResponse(ApiModel): - collections: list[str] - - -class RagDeleteCollectionResponse(ApiModel): - status: str - collection: str +class DeleteDocumentResponse(ApiModel): + document_id: FileId + deleted: bool diff --git a/engine/src/stirling/models/__init__.py b/engine/src/stirling/models/__init__.py index 0d4c774d0..9205f60a1 100644 --- a/engine/src/stirling/models/__init__.py +++ b/engine/src/stirling/models/__init__.py @@ -1,9 +1,10 @@ from . import tool_models -from .base import ApiModel +from .base import ApiModel, FileId from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint __all__ = [ "ApiModel", + "FileId", "OPERATIONS", "ParamToolModel", "ToolEndpoint", diff --git a/engine/src/stirling/models/base.py b/engine/src/stirling/models/base.py index 6b4d9bdb4..815c09854 100644 --- a/engine/src/stirling/models/base.py +++ b/engine/src/stirling/models/base.py @@ -1,8 +1,15 @@ from __future__ import annotations +from typing import NewType + from pydantic import BaseModel, ConfigDict from pydantic.alias_generators import to_camel +# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's +# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key +# throughout the engine. +FileId = NewType("FileId", str) + class ApiModel(BaseModel): model_config = ConfigDict( diff --git a/engine/src/stirling/rag/capability.py b/engine/src/stirling/rag/capability.py index 9950cf30a..494000f08 100644 --- a/engine/src/stirling/rag/capability.py +++ b/engine/src/stirling/rag/capability.py @@ -1,11 +1,16 @@ from __future__ import annotations +import logging from collections.abc import Awaitable, Callable -from pydantic_ai import FunctionToolset +from pydantic_ai import FunctionToolset, RunContext, ToolDefinition from pydantic_ai.toolsets import AbstractToolset +from stirling.models import FileId from stirling.rag.service import RagService +from stirling.rag.store import SearchResult + +logger = logging.getLogger(__name__) class RagCapability: @@ -22,19 +27,29 @@ class RagCapability: When no collections are pinned, the instructions are generated dynamically at run time so the agent sees the current list of collections in the store. + + Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a + single agent run. """ def __init__( self, rag_service: RagService, - collections: list[str] | None = None, + collections: list[FileId] | None = None, top_k: int = 5, + max_searches: int = 5, ) -> None: self._rag_service = rag_service self._collections = collections self._top_k = top_k + self._max_searches = max_searches + self._search_count = 0 toolset: FunctionToolset[None] = FunctionToolset() - toolset.add_function(self._search_knowledge, name="search_knowledge") + toolset.add_function( + self._search_knowledge, + name="search_knowledge", + prepare=self._prepare_search_knowledge, + ) self._toolset = toolset @property @@ -48,7 +63,7 @@ class RagCapability: return self._toolset @staticmethod - def _static_instructions_text(collections: list[str]) -> str: + def _static_instructions_text(collections: list[FileId]) -> str: collection_desc = f"collections: {', '.join(collections)}" return ( "You have access to a knowledge base search tool called 'search_knowledge'. " @@ -73,6 +88,18 @@ class RagCapability: "You do not have to use it if the answer is already clear from the provided text." ) + async def _prepare_search_knowledge( + self, + ctx: RunContext[None], + tool_def: ToolDefinition, + ) -> ToolDefinition | None: + """Remove the search tool from the agent's toolset once the per-run search + budget is exhausted. The agent then has no choice but to answer from what it + has already retrieved, which prevents runaway search loops.""" + if self._search_count >= self._max_searches: + return None + return tool_def + async def _search_knowledge(self, query: str, max_results: int | None = None) -> str: """Search the knowledge base for information relevant to the query. @@ -83,6 +110,7 @@ class RagCapability: Returns: Formatted text with the most relevant knowledge base excerpts. """ + self._search_count += 1 k = max_results if max_results is not None else self._top_k if self._collections: all_results = [] @@ -95,8 +123,21 @@ class RagCapability: results = await self._rag_service.search(query, top_k=k) if not results: + logger.info("[rag] search_knowledge query=%r -> 0 results", query) return "No relevant results found in the knowledge base." + formatted = self._format_results(results) + logger.info( + "[rag] search_knowledge query=%r -> %d results, %d chars", + query, + len(results), + len(formatted), + ) + logger.debug("[rag] search_knowledge query=%r returned:\n%s", query, formatted) + return formatted + + @staticmethod + def _format_results(results: list[SearchResult]) -> str: sections = [] for i, result in enumerate(results, 1): source = result.document.metadata.get("source", "unknown") diff --git a/engine/src/stirling/rag/embedder.py b/engine/src/stirling/rag/embedder.py index 5df71bc80..c97c28759 100644 --- a/engine/src/stirling/rag/embedder.py +++ b/engine/src/stirling/rag/embedder.py @@ -5,14 +5,27 @@ from pydantic_ai import Embedder from stirling.rag.chunker import chunk_text from stirling.rag.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 +# of round trips. VoyageAI caps at 1000, OpenAI at 2048, Cohere at 96; 256 is a good +# default for Voyage/OpenAI. Cohere users should pass a lower value via construction. +DEFAULT_EMBED_BATCH_SIZE = 256 + class EmbeddingService: """Wraps Pydantic AI's Embedder to provide document chunking and embedding.""" - def __init__(self, model_name: str, chunk_size: int = 512, chunk_overlap: int = 64) -> None: + def __init__( + self, + model_name: str, + chunk_size: int = 512, + chunk_overlap: int = 64, + embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE, + ) -> None: self._embedder = Embedder(model_name) self._chunk_size = chunk_size self._chunk_overlap = chunk_overlap + self._embed_batch_size = embed_batch_size async def embed_query(self, text: str) -> list[float]: """Embed a search query, optimised for retrieval.""" @@ -20,11 +33,19 @@ class EmbeddingService: return list(result.embeddings[0]) async def embed_documents(self, texts: list[str]) -> list[list[float]]: - """Embed multiple document texts for indexing.""" + """Embed multiple document texts for indexing. + + Splits the input into batches of ``embed_batch_size`` so callers can hand us + any number of chunks without hitting provider per-request limits. + """ if not texts: return [] - result = await self._embedder.embed_documents(texts) - return [list(emb) for emb in result.embeddings] + all_embeddings: list[list[float]] = [] + for start in range(0, len(texts), self._embed_batch_size): + batch = texts[start : start + self._embed_batch_size] + result = await self._embedder.embed_documents(batch) + all_embeddings.extend(list(emb) for emb in result.embeddings) + return all_embeddings def chunk_and_prepare( self, diff --git a/engine/src/stirling/rag/pgvector_store.py b/engine/src/stirling/rag/pgvector_store.py index 9eedffb47..b73596e49 100644 --- a/engine/src/stirling/rag/pgvector_store.py +++ b/engine/src/stirling/rag/pgvector_store.py @@ -131,3 +131,7 @@ class PgVectorStore(VectorStore): ) row = await cur.fetchone() return row is not None + + async def close(self) -> None: + # Connections are opened and closed per call, so nothing persistent to release. + return None diff --git a/engine/src/stirling/rag/service.py b/engine/src/stirling/rag/service.py index b8c9c4ab1..e818cdccd 100644 --- a/engine/src/stirling/rag/service.py +++ b/engine/src/stirling/rag/service.py @@ -2,6 +2,7 @@ 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 @@ -18,7 +19,7 @@ class RagService: async def index_text( self, - collection: str, + collection: FileId, text: str, source: str = "", metadata: dict[str, str] | None = None, @@ -31,7 +32,7 @@ class RagService: await self._store.add_documents(collection, documents, embeddings) return len(documents) - async def index_documents(self, collection: str, documents: list[Document]) -> int: + 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 @@ -39,10 +40,23 @@ class RagService: 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: str | None = None, + collection: FileId | None = None, top_k: int | None = None, ) -> list[SearchResult]: """Embed query and search across one or all collections. @@ -71,10 +85,18 @@ class RagService: all_results.sort(key=lambda r: r.score, reverse=True) return all_results[:k] - async def delete_collection(self, collection: str) -> None: + async def delete_collection(self, collection: FileId) -> None: """Remove a collection and all its documents.""" await self._store.delete_collection(collection) - async def list_collections(self) -> list[str]: + 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 await self._store.list_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/sqlite_vec_store.py b/engine/src/stirling/rag/sqlite_vec_store.py index b3008dcf2..25bbffcd3 100644 --- a/engine/src/stirling/rag/sqlite_vec_store.py +++ b/engine/src/stirling/rag/sqlite_vec_store.py @@ -225,3 +225,19 @@ class SqliteVecStore(VectorStore): def _sync_has_collection(self, collection: str) -> bool: row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone() return row is not None + + async def close(self) -> None: + async with self._lock: + await asyncio.to_thread(self._sync_close) + + def _sync_close(self) -> None: + """Checkpoint the WAL into the main database file and close the connection so + the .db-shm and .db-wal files are cleaned up on graceful shutdown.""" + if self._db_path is not None: + try: + self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + self._conn.commit() + except sqlite3.Error: + # Best effort: if checkpointing fails we still want to close the connection. + pass + self._conn.close() diff --git a/engine/src/stirling/rag/store.py b/engine/src/stirling/rag/store.py index 1ad0dffbf..1a73300d7 100644 --- a/engine/src/stirling/rag/store.py +++ b/engine/src/stirling/rag/store.py @@ -57,3 +57,7 @@ class VectorStore(ABC): @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/tests/agents/test_math_presentation.py b/engine/tests/agents/test_math_presentation.py index d4af9d671..6d50afc5f 100644 --- a/engine/tests/agents/test_math_presentation.py +++ b/engine/tests/agents/test_math_presentation.py @@ -14,6 +14,7 @@ from pydantic import ValidationError from stirling.agents.math_presentation import extract_math_verdict from stirling.contracts import ( + AiFile, ExtractedFileText, ExtractedTextArtifact, MathAuditorToolReportArtifact, @@ -21,6 +22,7 @@ from stirling.contracts import ( WorkflowArtifact, ) from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict +from stirling.models import FileId def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict: @@ -42,7 +44,7 @@ def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict: def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest: return OrchestratorRequest( user_message="review the math", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], artifacts=artifacts, ) diff --git a/engine/tests/agents/test_orchestrator_pdf_comment.py b/engine/tests/agents/test_orchestrator_pdf_comment.py index e17cf373e..a9aabaa56 100644 --- a/engine/tests/agents/test_orchestrator_pdf_comment.py +++ b/engine/tests/agents/test_orchestrator_pdf_comment.py @@ -23,8 +23,9 @@ from unittest.mock import AsyncMock, patch import pytest from stirling.agents import OrchestratorAgent -from stirling.contracts import OrchestratorRequest +from stirling.contracts import AiFile, OrchestratorRequest from stirling.contracts.pdf_edit import EditPlanResponse +from stirling.models import FileId from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams from stirling.services.runtime import AppRuntime @@ -39,7 +40,7 @@ async def test_delegate_pdf_review_wires_prompt_to_tool_step(runtime: AppRuntime orchestrator = OrchestratorAgent(runtime) request = OrchestratorRequest( user_message="please add review comments flagging ambiguous dates", - file_names=["contract.pdf"], + files=[AiFile(id=FileId("contract-id"), name="contract.pdf")], ) ctx = SimpleNamespace(deps=_FakeDeps(request=request)) diff --git a/engine/tests/agents/test_pdf_questions_orchestrate.py b/engine/tests/agents/test_pdf_questions_orchestrate.py index a34ba11c0..224f20914 100644 --- a/engine/tests/agents/test_pdf_questions_orchestrate.py +++ b/engine/tests/agents/test_pdf_questions_orchestrate.py @@ -12,12 +12,15 @@ import pytest from stirling.agents.pdf_questions import _MATH_SYNTH_SYSTEM_PROMPT, PdfQuestionAgent from stirling.contracts import ( + AiFile, + EditPlanResponse, MathAuditorToolReportArtifact, OrchestratorRequest, PdfQuestionAnswerResponse, SupportedCapability, ) from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict +from stirling.models import FileId from stirling.models.agent_tool_models import AgentToolId from stirling.services.runtime import AppRuntime @@ -49,25 +52,23 @@ def _make_verdict() -> Verdict: @pytest.mark.anyio -async def test_orchestrate_classifier_true_embeds_plan_in_answer(runtime: AppRuntime) -> None: - """First turn — classifier says math; the response is a PdfQuestionAnswerResponse - with the math-auditor plan attached as a nullable ``edit_plan`` field. The - answer is empty on this turn; the caller runs the embedded plan and resumes.""" +async def test_orchestrate_classifier_true_returns_math_audit_plan(runtime: AppRuntime) -> None: + """First turn — classifier says math; the response is an EditPlanResponse + (``outcome=PLAN``) with ``resume_with=PDF_QUESTION``. The caller runs the + plan and re-invokes the orchestrator with the verdict in artifacts.""" agent = PdfQuestionAgent(runtime) request = OrchestratorRequest( user_message="ist die mathematik korrekt?", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], ) with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)): response = await agent.orchestrate(request) - assert isinstance(response, PdfQuestionAnswerResponse) - assert response.answer == "" - assert response.edit_plan is not None - assert response.edit_plan.resume_with == SupportedCapability.PDF_QUESTION - assert len(response.edit_plan.steps) == 1 - assert response.edit_plan.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT + assert isinstance(response, EditPlanResponse) + assert response.resume_with == SupportedCapability.PDF_QUESTION + assert len(response.steps) == 1 + assert response.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT @pytest.mark.anyio @@ -81,7 +82,7 @@ async def test_orchestrate_resume_synthesises_answer_without_calling_classifier( verdict = _make_verdict() request = OrchestratorRequest( user_message="ist die mathematik korrekt?", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], artifacts=[MathAuditorToolReportArtifact(report=verdict)], ) canned_answer = "Die Summe stimmt nicht: angegeben $215,000, erwartet $215,500." diff --git a/engine/tests/agents/test_pdf_review.py b/engine/tests/agents/test_pdf_review.py index f902d4789..112623c74 100644 --- a/engine/tests/agents/test_pdf_review.py +++ b/engine/tests/agents/test_pdf_review.py @@ -20,9 +20,9 @@ from stirling.agents.pdf_review import ( _LocalisedComment, _LocalisedVerdict, ) -from stirling.contracts import EditPlanResponse, OrchestratorRequest, SupportedCapability +from stirling.contracts import AiFile, EditPlanResponse, OrchestratorRequest, SupportedCapability from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict -from stirling.models import ToolEndpoint +from stirling.models import FileId, ToolEndpoint from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams from stirling.services.runtime import AppRuntime @@ -144,7 +144,10 @@ async def test_orchestrate_classifier_true_emits_math_audit_plan(runtime: AppRun """First turn — when the math-intent classifier says yes, emit a one-step plan calling the math auditor with resume_with=PDF_REVIEW.""" agent = PdfReviewAgent(runtime) - request = OrchestratorRequest(user_message="vérifie les totaux", file_names=["report.pdf"]) + request = OrchestratorRequest( + user_message="vérifie les totaux", + files=[AiFile(id=FileId("report-id"), name="report.pdf")], + ) with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)): response = await agent.orchestrate(request) @@ -161,7 +164,7 @@ async def test_orchestrate_classifier_false_routes_to_pdf_comment_agent(runtime: agent = PdfReviewAgent(runtime) request = OrchestratorRequest( user_message="review the invoices for ambiguous wording", - file_names=["contract.pdf"], + files=[AiFile(id=FileId("contract-id"), name="contract.pdf")], ) with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=False)): @@ -187,7 +190,7 @@ async def test_orchestrate_resume_uses_verdict_without_calling_classifier( verdict = _make_verdict([_discrepancy(page=0, stated="$100")]) request = OrchestratorRequest( user_message="flag math errors", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], artifacts=[MathAuditorToolReportArtifact(report=verdict)], ) canned = _LocalisedVerdict( diff --git a/engine/tests/conftest.py b/engine/tests/conftest.py index 3bc409f61..36194b37f 100644 --- a/engine/tests/conftest.py +++ b/engine/tests/conftest.py @@ -30,6 +30,7 @@ def build_app_settings() -> AppSettings: rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, + rag_max_searches=5, 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 55df5e689..9d35a615e 100644 --- a/engine/tests/pdf_comment/test_routes.py +++ b/engine/tests/pdf_comment/test_routes.py @@ -42,6 +42,7 @@ class StubSettingsProvider: rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, + rag_max_searches=5, max_pages=100, max_characters=100_000, posthog_enabled=False, diff --git a/engine/tests/test_pdf_edit_agent.py b/engine/tests/test_pdf_edit_agent.py index 22716a50b..dd9bf7f51 100644 --- a/engine/tests/test_pdf_edit_agent.py +++ b/engine/tests/test_pdf_edit_agent.py @@ -7,6 +7,7 @@ import pytest from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection from stirling.agents.pdf_edit import PdfEditPlanOutput from stirling.contracts import ( + AiFile, EditCannotDoResponse, EditClarificationRequest, EditPlanResponse, @@ -19,6 +20,7 @@ from stirling.contracts import ( SupportedCapability, ToolOperationStep, ) +from stirling.models import FileId from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint from stirling.services.runtime import AppRuntime @@ -91,7 +93,7 @@ async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> Non response = await agent.handle( PdfEditRequest( user_message="Rotate the PDF clockwise and then compress it.", - file_names=["scan.pdf"], + files=[AiFile(id=FileId("scan-id"), name="scan.pdf")], ) ) @@ -117,7 +119,7 @@ async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtim request = PdfEditRequest( user_message="Rotate the PDF clockwise and then compress it.", - file_names=["scan.pdf"], + files=[AiFile(id=FileId("scan-id"), name="scan.pdf")], ) response = await agent.handle(request) @@ -181,13 +183,18 @@ async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime response = await agent.handle( PdfEditRequest( user_message="Split after every page that says 'NEW PAGE'.", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], ) ) assert isinstance(response, NeedContentResponse) assert response.resume_with == SupportedCapability.PDF_EDIT - assert response.files == [NeedContentFileRequest(file_name="report.pdf", content_types=[PdfContentType.PAGE_TEXT])] + assert response.files == [ + NeedContentFileRequest( + file=AiFile(id=FileId("report-id"), name="report.pdf"), + content_types=[PdfContentType.PAGE_TEXT], + ) + ] assert response.max_pages == runtime.settings.max_pages assert response.max_characters == runtime.settings.max_characters assert parameter_selector.calls == [] @@ -269,7 +276,7 @@ async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: Ap await agent.handle( PdfEditRequest( user_message="Rotate clockwise.", - file_names=["report.pdf"], + files=[AiFile(id=FileId("report-id"), name="report.pdf")], page_text=page_text, ) ) diff --git a/engine/tests/test_pdf_question_agent.py b/engine/tests/test_pdf_question_agent.py index b638c6bc4..714406339 100644 --- a/engine/tests/test_pdf_question_agent.py +++ b/engine/tests/test_pdf_question_agent.py @@ -1,64 +1,133 @@ from __future__ import annotations +from dataclasses import replace + import pytest from stirling.agents import PdfQuestionAgent from stirling.contracts import ( + AiFile, ExtractedFileText, - NeedContentResponse, + NeedIngestResponse, + PdfContentType, PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, + PdfQuestionTerminalResponse, PdfTextSelection, + SupportedCapability, ) +from stirling.models import FileId +from stirling.rag import Document, RagService, SqliteVecStore from stirling.services.runtime import AppRuntime -class StubPdfQuestionAgent(PdfQuestionAgent): - def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None: - super().__init__(runtime) - self.response = response +class StubEmbedder: + """Deterministic embeddings so RAG lookups work in tests without network.""" - async def _run_answer_agent( + 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, - request: PdfQuestionRequest, - ) -> PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse: - return self.response + text: str, + 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: 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 -def invoice_page() -> ExtractedFileText: - return ExtractedFileText( - file_name="invoice.pdf", - pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")], +class StubPdfQuestionAgent(PdfQuestionAgent): + def __init__(self, runtime: AppRuntime, response: PdfQuestionTerminalResponse) -> None: + super().__init__(runtime) + self._response = response + + async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse: + return self._response + + +@pytest.fixture +def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime: + """A runtime whose RAG service uses a stub embedder + ephemeral store.""" + stub = RagService( + embedder=StubEmbedder(), # type: ignore[arg-type] + store=SqliteVecStore.ephemeral(), + default_top_k=runtime.settings.rag_default_top_k, ) + return replace(runtime, rag_service=stub) @pytest.mark.anyio -async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None: - agent = PdfQuestionAgent(runtime) +async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag: AppRuntime) -> None: + agent = PdfQuestionAgent(runtime_with_stub_rag) - response = await agent.handle( - PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"]) - ) + missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf") + response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[missing_file])) - assert isinstance(response, NeedContentResponse) + assert isinstance(response, NeedIngestResponse) + assert response.resume_with == SupportedCapability.PDF_QUESTION + assert response.files_to_ingest == [missing_file] + assert PdfContentType.PAGE_TEXT in response.content_types @pytest.mark.anyio -async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None: +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.", + source="present.pdf", + ) + agent = PdfQuestionAgent(runtime_with_stub_rag) + + present_file = AiFile(id=FileId("present-id"), name="present.pdf") + missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf") + response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[present_file, missing_file])) + + assert isinstance(response, NeedIngestResponse) + assert response.files_to_ingest == [missing_file] + + +@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.", + source="invoice.pdf", + ) agent = StubPdfQuestionAgent( - runtime, + runtime_with_stub_rag, PdfQuestionAnswerResponse( answer="The invoice total is 120.00.", - evidence=[invoice_page()], + evidence=[ + ExtractedFileText( + file_name="invoice.pdf", + pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")], + ) + ], ), ) response = await agent.handle( PdfQuestionRequest( question="What is the total?", - page_text=[invoice_page()], - file_names=["invoice.pdf"], + files=[AiFile(id=FileId("invoice-id"), name="invoice.pdf")], ) ) @@ -67,19 +136,21 @@ async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) - @pytest.mark.anyio -async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None: - agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text.")) +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.", + source="shipping.pdf", + ) + agent = StubPdfQuestionAgent( + runtime_with_stub_rag, + PdfQuestionNotFoundResponse(reason="The answer is not present in the text."), + ) response = await agent.handle( PdfQuestionRequest( question="What is the total?", - page_text=[ - ExtractedFileText( - file_name="invoice.pdf", - pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")], - ) - ], - file_names=["invoice.pdf"], + files=[AiFile(id=FileId("shipping-id"), name="shipping.pdf")], ) ) diff --git a/engine/tests/test_rag.py b/engine/tests/test_rag.py index da1d52465..9616ae834 100644 --- a/engine/tests/test_rag.py +++ b/engine/tests/test_rag.py @@ -2,6 +2,7 @@ from __future__ import annotations import pytest +from stirling.models import FileId from stirling.rag.capability import RagCapability from stirling.rag.chunker import chunk_text from stirling.rag.service import RagService @@ -163,38 +164,38 @@ class TestRagService: @pytest.mark.anyio async def test_index_and_search(self, rag_service: RagService) -> None: text = "Python is great for data science. It has many libraries like pandas and numpy." - count = await rag_service.index_text("docs", text, source="guide.pdf") + count = await rag_service.index_text(FileId("docs"), text, source="guide.pdf") assert count > 0 - results = await rag_service.search("Python libraries", collection="docs") + results = await rag_service.search("Python libraries", collection=FileId("docs")) assert len(results) > 0 assert results[0].document.text # non-empty text @pytest.mark.anyio async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None: - count = await rag_service.index_text("docs", "", source="empty.pdf") + count = await rag_service.index_text(FileId("docs"), "", 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="nonexistent") + results = await rag_service.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("col-a", "Machine learning overview.", source="ml.pdf") - await rag_service.index_text("col-b", "Deep learning with neural networks.", source="dl.pdf") + 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") results = await rag_service.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("temp", "Temporary data.", source="tmp.pdf") + await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf") collections = await rag_service.list_collections() assert "temp" in collections - await rag_service.delete_collection("temp") + await rag_service.delete_collection(FileId("temp")) collections = await rag_service.list_collections() assert "temp" not in collections @@ -214,7 +215,7 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re class TestRagCapability: def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None: - cap = RagCapability(rag_service, collections=["docs", "manuals"]) + cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")]) instructions = cap.instructions assert isinstance(instructions, str) assert "docs, manuals" in instructions @@ -227,8 +228,8 @@ class TestRagCapability: @pytest.mark.anyio async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None: - await rag_service.index_text("col-a", "Alpha content.", source="a.pdf") - await rag_service.index_text("col-b", "Beta content.", source="b.pdf") + 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) instructions_fn = cap.instructions assert callable(instructions_fn) @@ -252,7 +253,7 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None: - await rag_service.index_text("docs", "Python is a programming language.", source="guide.pdf") + await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf") cap = RagCapability(rag_service) output = await _invoke_search_knowledge(cap, "Python") assert "[Result 1" in output @@ -262,10 +263,10 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None: - await rag_service.index_text("pinned", "Pinned collection content.", source="pinned.pdf") - await rag_service.index_text("other", "Content in another collection.", source="other.pdf") + 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") - cap = RagCapability(rag_service, collections=["pinned"]) + cap = RagCapability(rag_service, collections=[FileId("pinned")]) output = await _invoke_search_knowledge(cap, "content") assert "pinned.pdf" in output assert "other.pdf" not in output @@ -273,7 +274,7 @@ class TestRagCapability: @pytest.mark.anyio async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None: paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10)) - await rag_service.index_text("bulk", paragraphs, source="bulk.pdf") + await rag_service.index_text(FileId("bulk"), paragraphs, source="bulk.pdf") cap = RagCapability(rag_service) output = await _invoke_search_knowledge(cap, "topic", max_results=2) @@ -281,3 +282,27 @@ class TestRagCapability: 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: + """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) + 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] + + +def _dummy_tool_def() -> object: + """Sentinel passed to ``_prepare_search_knowledge``. The callback only inspects + ``_search_count``; it doesn't read anything off the tool_def or context.""" + return object() diff --git a/engine/tests/test_rag_routes.py b/engine/tests/test_rag_routes.py index 622715927..f9f61c996 100644 --- a/engine/tests/test_rag_routes.py +++ b/engine/tests/test_rag_routes.py @@ -6,14 +6,13 @@ import pytest from fastapi.testclient import TestClient from stirling.api import app -from stirling.api.dependencies import get_rag_embedding_model, get_rag_service +from stirling.api.dependencies import get_rag_service +from stirling.models import FileId from stirling.rag import Document, RagService, SqliteVecStore -TEST_EMBEDDING_MODEL = "test-embedder" - class StubEmbedder: - """Deterministic embeddings for route tests — no network, no provider needed.""" + """Deterministic embeddings for route tests: no network, no provider needed.""" def __init__(self, dim: int = 8) -> None: self._dim = dim @@ -53,153 +52,154 @@ def _build_service() -> RagService: @pytest.fixture -def client() -> Iterator[TestClient]: - service = _build_service() +def service() -> RagService: + return _build_service() + + +@pytest.fixture +def client(service: RagService) -> Iterator[TestClient]: app.dependency_overrides[get_rag_service] = lambda: service - app.dependency_overrides[get_rag_embedding_model] = lambda: TEST_EMBEDDING_MODEL try: yield TestClient(app) finally: app.dependency_overrides.pop(get_rag_service, None) - app.dependency_overrides.pop(get_rag_embedding_model, None) -# ── /status ───────────────────────────────────────────────────────────── +# ── POST /documents ───────────────────────────────────────────────────── -def test_status_reports_embedding_model_and_collections(client: TestClient) -> None: +def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None: + response = client.post( + "/api/v1/rag/documents", + json={ + "documentId": "doc-123", + "source": "report.pdf", + "pageText": [ + {"pageNumber": 1, "text": "The introduction covers the main topic."}, + {"pageNumber": 2, "text": "The conclusion summarises the findings."}, + ], + }, + ) + assert response.status_code == 200 + body = response.json() + assert body["documentId"] == "doc-123" + assert body["chunksIndexed"] >= 2 + + +@pytest.mark.anyio +async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None: client.post( - "/api/v1/rag/index", - json={"collection": "my-docs", "text": "Hello world.", "source": "a.pdf"}, + "/api/v1/rag/documents", + json={ + "documentId": "replace-me", + "source": "replace-me.pdf", + "pageText": [{"pageNumber": 1, "text": "Original content that existed before."}], + }, ) - response = client.get("/api/v1/rag/status") - assert response.status_code == 200 - body = response.json() - assert body["embeddingModel"] == TEST_EMBEDDING_MODEL - assert "my-docs" in body["collections"] - - -def test_status_when_empty(client: TestClient) -> None: - response = client.get("/api/v1/rag/status") - assert response.status_code == 200 - body = response.json() - assert body == {"embeddingModel": TEST_EMBEDDING_MODEL, "collections": []} - - -# ── /index ────────────────────────────────────────────────────────────── - - -def test_index_returns_chunk_count(client: TestClient) -> None: + # Second ingest with different content should replace the first entirely response = client.post( - "/api/v1/rag/index", - json={"collection": "indexed", "text": "Short text.", "source": "doc.pdf"}, + "/api/v1/rag/documents", + json={ + "documentId": "replace-me", + "source": "replace-me.pdf", + "pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}], + }, ) assert response.status_code == 200 - body = response.json() - assert body["collection"] == "indexed" - assert body["chunksIndexed"] >= 1 + + results = await service.search("New content", collection=FileId("replace-me"), top_k=5) + texts = [r.document.text for r in results] + assert any("New content" in t for t in texts) + assert not any("Original content" in t for t in texts) -def test_index_rejects_empty_collection_name(client: TestClient) -> None: +def test_ingest_document_skips_empty_pages(client: TestClient) -> None: response = client.post( - "/api/v1/rag/index", - json={"collection": "", "text": "Text.", "source": "x.pdf"}, + "/api/v1/rag/documents", + json={ + "documentId": "mixed", + "source": "mixed.pdf", + "pageText": [ + {"pageNumber": 1, "text": " "}, + {"pageNumber": 2, "text": "Real content on page 2."}, + ], + }, + ) + assert response.status_code == 200 + assert response.json()["chunksIndexed"] >= 1 + + +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"}) + 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", + json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]}, ) assert response.status_code == 422 -def test_index_rejects_oversized_text(client: TestClient) -> None: - huge = "x" * 1_000_001 # Just over the 1MB cap +def test_ingest_document_rejects_missing_source(client: TestClient) -> None: response = client.post( - "/api/v1/rag/index", - json={"collection": "toobig", "text": huge}, + "/api/v1/rag/documents", + json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]}, ) assert response.status_code == 422 -# ── /search ───────────────────────────────────────────────────────────── - - -def test_search_returns_results(client: TestClient) -> None: - client.post( - "/api/v1/rag/index", - json={"collection": "search-test", "text": "Python is fun.", "source": "guide.pdf"}, - ) +def test_ingest_document_rejects_empty_source(client: TestClient) -> None: response = client.post( - "/api/v1/rag/search", - json={"query": "Python", "collection": "search-test", "topK": 3}, - ) - assert response.status_code == 200 - body = response.json() - assert body["query"] == "Python" - assert len(body["results"]) >= 1 - first = body["results"][0] - assert first["source"] == "guide.pdf" - assert "score" in first - - -def test_search_rejects_empty_collection_name(client: TestClient) -> None: - response = client.post( - "/api/v1/rag/search", - json={"query": "anything", "collection": ""}, + "/api/v1/rag/documents", + json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]}, ) assert response.status_code == 422 -def test_search_without_collection_searches_all(client: TestClient) -> None: - client.post( - "/api/v1/rag/index", - json={"collection": "col-one", "text": "Alpha content.", "source": "one.pdf"}, - ) - client.post( - "/api/v1/rag/index", - json={"collection": "col-two", "text": "Beta content.", "source": "two.pdf"}, - ) +def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None: response = client.post( - "/api/v1/rag/search", - json={"query": "content"}, + "/api/v1/rag/documents", + json={ + "documentId": "bad-page", + "source": "bad-page.pdf", + "pageText": [{"pageNumber": 0, "text": "something"}], + }, ) - assert response.status_code == 200 - body = response.json() - assert len(body["results"]) >= 1 + assert response.status_code == 422 -# ── /collections ──────────────────────────────────────────────────────── +# ── DELETE /documents/{id} ────────────────────────────────────────────── -def test_collections_empty_when_no_data(client: TestClient) -> None: - response = client.get("/api/v1/rag/collections") - assert response.status_code == 200 - assert response.json() == {"collections": []} - - -def test_collections_lists_indexed(client: TestClient) -> None: +def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None: client.post( - "/api/v1/rag/index", - json={"collection": "list-me", "text": "Text.", "source": "x.pdf"}, + "/api/v1/rag/documents", + json={ + "documentId": "to-delete", + "source": "to-delete.pdf", + "pageText": [{"pageNumber": 1, "text": "Text."}], + }, ) - response = client.get("/api/v1/rag/collections") + response = client.delete("/api/v1/rag/documents/to-delete") assert response.status_code == 200 - assert "list-me" in response.json()["collections"] + assert response.json() == {"documentId": "to-delete", "deleted": True} -# ── DELETE /collections/{name} ────────────────────────────────────────── +def test_delete_document_is_idempotent(client: TestClient) -> None: + response = client.delete("/api/v1/rag/documents/never-existed") + assert response.status_code == 200 + assert response.json() == {"documentId": "never-existed", "deleted": False} -def test_delete_collection_removes_it(client: TestClient) -> None: +@pytest.mark.anyio +async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None: client.post( - "/api/v1/rag/index", - json={"collection": "to-delete", "text": "Text.", "source": "x.pdf"}, + "/api/v1/rag/documents", + json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]}, ) - response = client.delete("/api/v1/rag/collections/to-delete") - assert response.status_code == 200 - assert response.json() == {"status": "deleted", "collection": "to-delete"} - - listing = client.get("/api/v1/rag/collections").json() - assert "to-delete" not in listing["collections"] - - -def test_delete_nonexistent_collection_is_idempotent(client: TestClient) -> None: - response = client.delete("/api/v1/rag/collections/never-existed") - assert response.status_code == 200 - assert response.json() == {"status": "deleted", "collection": "never-existed"} + assert await service.has_collection(FileId("gone")) + client.delete("/api/v1/rag/documents/gone") + assert not await service.has_collection(FileId("gone")) diff --git a/engine/tests/test_stirling_api.py b/engine/tests/test_stirling_api.py index 96fd98268..547b1e2a5 100644 --- a/engine/tests/test_stirling_api.py +++ b/engine/tests/test_stirling_api.py @@ -88,7 +88,10 @@ def test_health_route() -> None: def test_orchestrator_route() -> None: - response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]}) + response = client.post( + "/api/v1/orchestrator", + json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]}, + ) assert response.status_code == 200 assert response.json()["outcome"] == "need_content" @@ -106,8 +109,7 @@ def test_pdf_questions_route() -> None: "/api/v1/pdf/questions", json={ "question": "what is this?", - "fileNames": ["test.pdf"], - "pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}], + "files": [{"id": "test-id", "name": "test.pdf"}], }, ) diff --git a/engine/tests/test_stirling_contracts.py b/engine/tests/test_stirling_contracts.py index 7fc1d9c64..7e4a2416d 100644 --- a/engine/tests/test_stirling_contracts.py +++ b/engine/tests/test_stirling_contracts.py @@ -3,6 +3,7 @@ from stirling.contracts import ( AgentExecutionRequest, AgentSpec, AgentSpecStep, + AiFile, EditPlanResponse, ExecutionContext, ExtractedFileText, @@ -12,13 +13,14 @@ from stirling.contracts import ( PdfTextSelection, ToolOperationStep, ) +from stirling.models import FileId from stirling.models.tool_models import Angle, RotatePdfParams, ToolEndpoint def test_orchestrator_request_accepts_user_message() -> None: request = OrchestratorRequest( user_message="Rotate the PDF", - file_names=["test.pdf"], + files=[AiFile(id=FileId("test-id"), name="test.pdf")], artifacts=[ ExtractedTextArtifact( files=[ @@ -89,6 +91,7 @@ def test_app_settings_accepts_model_configuration() -> None: rag_chunk_size=512, rag_chunk_overlap=64, rag_default_top_k=5, + rag_max_searches=5, max_pages=200, max_characters=200_000, posthog_enabled=False,