From b130242688cb08fd5a8e3b20e1b74a80dc1dc936 Mon Sep 17 00:00:00 2001 From: James Brunton Date: Thu, 9 Apr 2026 09:04:38 +0100 Subject: [PATCH] Add Java orchestrator to connect to the AI engine (#6003) # Description of Changes Add Java orchestration layer which can connect and go back and forth with the AI engine to get results for the user. It's expected that the AI engine will not be publicly available and this Java layer will always be in front of it, to manage sessions and auth etc. --- .gitignore | 1 + .../common/model/ApplicationProperties.java | 8 + .../src/main/resources/settings.yml.template | 5 + .../controller/api/AiEngineController.java | 83 ++++++ .../model/api/ai/AiPdfContentType.java | 57 ++++ .../model/api/ai/AiWorkflowFileInput.java | 22 ++ .../model/api/ai/AiWorkflowFileRequest.java | 22 ++ .../model/api/ai/AiWorkflowOutcome.java | 43 +++ .../model/api/ai/AiWorkflowRequest.java | 23 ++ .../model/api/ai/AiWorkflowResponse.java | 58 ++++ .../model/api/ai/AiWorkflowTextSelection.java | 16 + .../proprietary/service/AiEngineClient.java | 108 +++++++ .../service/AiWorkflowService.java | 181 ++++++++++++ .../service/PdfContentExtractor.java | 279 ++++++++++++++++++ engine/pyproject.toml | 6 +- engine/src/stirling/agents/orchestrator.py | 79 ++++- engine/src/stirling/agents/pdf_questions.py | 40 ++- engine/src/stirling/contracts/__init__.py | 33 ++- engine/src/stirling/contracts/agent_drafts.py | 8 +- engine/src/stirling/contracts/agent_specs.py | 4 +- engine/src/stirling/contracts/common.py | 86 +++++- engine/src/stirling/contracts/execution.py | 7 +- engine/src/stirling/contracts/orchestrator.py | 21 +- engine/src/stirling/contracts/pdf_edit.py | 9 +- .../src/stirling/contracts/pdf_questions.py | 29 +- engine/tests/test_pdf_question_agent.py | 32 +- engine/tests/test_stirling_api.py | 19 +- engine/tests/test_stirling_contracts.py | 19 +- 28 files changed, 1222 insertions(+), 76 deletions(-) create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPdfContentType.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileInput.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowTextSelection.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java create mode 100644 app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java diff --git a/.gitignore b/.gitignore index b5025f1b4..48d38e3e5 100644 --- a/.gitignore +++ b/.gitignore @@ -181,6 +181,7 @@ venv.bak/ .idea/ *.iml out/ +.junie/ # Ignore Mac DS_Store files .DS_Store 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 fa36998e7..dba7deca2 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 @@ -75,6 +75,7 @@ public class ApplicationProperties { private AutoPipeline autoPipeline = new AutoPipeline(); private ProcessExecutor processExecutor = new ProcessExecutor(); private PdfEditor pdfEditor = new PdfEditor(); + private AiEngine aiEngine = new AiEngine(); @Bean public PropertySource dynamicYamlPropertySource(ConfigurableEnvironment environment) @@ -231,6 +232,13 @@ public class ApplicationProperties { } } + @Data + public static class AiEngine { + private boolean enabled = false; + private String url = "http://localhost:5001"; + private int timeoutSeconds = 120; + } + @Data public static class Legal { private String termsAndConditions; diff --git a/app/core/src/main/resources/settings.yml.template b/app/core/src/main/resources/settings.yml.template index c78a515bd..2540c2f51 100644 --- a/app/core/src/main/resources/settings.yml.template +++ b/app/core/src/main/resources/settings.yml.template @@ -325,6 +325,11 @@ processExecutor: ghostscriptTimeoutMinutes: 30 ocrMyPdfTimeoutMinutes: 30 +aiEngine: + enabled: false # Set to 'true' to enable the AI engine integration + url: http://localhost:5001 # URL of the Python AI engine + timeoutSeconds: 120 # Timeout in seconds for AI engine requests + pdfEditor: fallback-font: classpath:/static/fonts/NotoSans-Regular.ttf # Override to point at a custom fallback font cache: diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java new file mode 100644 index 000000000..6262c2c4e --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java @@ -0,0 +1,83 @@ +package stirling.software.proprietary.controller.api; + +import java.io.IOException; + +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.ModelAttribute; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.server.ResponseStatusException; + +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.tags.Tag; + +import jakarta.validation.Valid; + +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import stirling.software.proprietary.model.api.ai.AiWorkflowRequest; +import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; +import stirling.software.proprietary.service.AiEngineClient; +import stirling.software.proprietary.service.AiWorkflowService; + +import tools.jackson.core.JacksonException; +import tools.jackson.databind.JsonNode; +import tools.jackson.databind.ObjectMapper; + +@Slf4j +@RestController +@RequestMapping("/api/v1/ai") +@RequiredArgsConstructor +@Tag(name = "AI Engine", description = "Endpoints for AI-powered PDF workflows") +public class AiEngineController { + + private final AiEngineClient aiEngineClient; + private final AiWorkflowService aiWorkflowService; + private final ObjectMapper objectMapper; + + @GetMapping("/health") + @Operation( + summary = "AI engine health check", + description = "Returns the health status of the AI engine including configured models") + public ResponseEntity health() throws IOException { + String response = aiEngineClient.get("/health"); + return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response); + } + + @PostMapping(value = "/orchestrate", consumes = MediaType.MULTIPART_FORM_DATA_VALUE) + @Operation( + summary = "Run an AI workflow against a PDF", + description = + "Accepts a PDF upload and a user message and returns an AI workflow result") + public ResponseEntity orchestrate( + @Valid @ModelAttribute AiWorkflowRequest request) throws IOException { + return ResponseEntity.ok(aiWorkflowService.orchestrate(request)); + } + + @PostMapping(value = "/pdf/edit", consumes = MediaType.APPLICATION_JSON_VALUE) + @Operation( + summary = "Generate a PDF edit plan", + description = + "Sends a user message to the PDF edit agent which returns a structured plan" + + " of tool operations to perform") + public ResponseEntity pdfEdit(@RequestBody String requestBody) throws IOException { + validateJson(requestBody); + String response = aiEngineClient.post("/api/v1/pdf/edit", requestBody); + return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response); + } + + private void validateJson(String body) { + try { + objectMapper.readValue(body, JsonNode.class); + } catch (JacksonException e) { + throw new ResponseStatusException( + HttpStatus.BAD_REQUEST, "Request body is not valid JSON"); + } + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPdfContentType.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPdfContentType.java new file mode 100644 index 000000000..d7c38a723 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiPdfContentType.java @@ -0,0 +1,57 @@ +package stirling.software.proprietary.model.api.ai; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; + +/** + * Types of content that can be extracted from a PDF and sent to the AI. + * + *

Values MUST match {@code PdfContentType} in {@code engine/src/stirling/contracts/common.py}. + */ +public enum AiPdfContentType { + // Document-level structured data + PAGE_LAYOUT("page_layout"), + DOCUMENT_METADATA("document_metadata"), + ENCRYPTION_INFO("encryption_info"), + BOOKMARKS("bookmarks"), + LAYERS("layers"), + EMBEDDED_FILES("embedded_files"), + JAVASCRIPT("javascript"), + LINKS("links"), + IMAGE_INFO("image_info"), + FONTS("fonts"), + + // Text and content + PAGE_TEXT("page_text"), + FULL_TEXT("full_text"), + FORM_FIELDS("form_fields"), + ANNOTATIONS("annotations"), + SIGNATURES("signatures"), + STRUCTURE_TREE("structure_tree"), + XMP_METADATA("xmp_metadata"), + + // Heavy content + COMPLIANCE("compliance"), + IMAGES("images"); + + private final String value; + + AiPdfContentType(String value) { + this.value = value; + } + + @JsonValue + public String getValue() { + return value; + } + + @JsonCreator + public static AiPdfContentType fromValue(String value) { + for (AiPdfContentType type : values()) { + if (type.value.equals(value)) { + return type; + } + } + throw new IllegalArgumentException("Unknown PDF content type: " + value); + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileInput.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileInput.java new file mode 100644 index 000000000..c83fa5569 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileInput.java @@ -0,0 +1,22 @@ +package stirling.software.proprietary.model.api.ai; + +import org.springframework.http.MediaType; +import org.springframework.web.multipart.MultipartFile; + +import io.swagger.v3.oas.annotations.media.Schema; + +import jakarta.validation.constraints.NotNull; + +import lombok.Data; + +@Data +@Schema(description = "A single PDF file input") +public class AiWorkflowFileInput { + + @NotNull + @Schema( + description = "The input PDF file", + contentMediaType = MediaType.APPLICATION_PDF_VALUE, + format = "binary") + private MultipartFile fileInput; +} 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 new file mode 100644 index 000000000..f23867028 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowFileRequest.java @@ -0,0 +1,22 @@ +package stirling.software.proprietary.model.api.ai; + +import java.util.ArrayList; +import java.util.List; + +import io.swagger.v3.oas.annotations.media.Schema; + +import lombok.Data; + +@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 = "Specific 1-based page numbers to extract from this file") + private List pageNumbers = new ArrayList<>(); + + @Schema(description = "Content types to extract from this file") + private List contentTypes = 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 new file mode 100644 index 000000000..78ce09b7f --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowOutcome.java @@ -0,0 +1,43 @@ +package stirling.software.proprietary.model.api.ai; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; + +/** + * Discriminator values for AI workflow responses. + * + *

Values MUST match {@code WorkflowOutcome} in {@code engine/src/stirling/contracts/common.py}. + */ +public enum AiWorkflowOutcome { + ANSWER("answer"), + NOT_FOUND("not_found"), + NEED_CONTENT("need_content"), + PLAN("plan"), + NEED_CLARIFICATION("need_clarification"), + CANNOT_DO("cannot_do"), + TOOL_CALL("tool_call"), + COMPLETED("completed"), + UNSUPPORTED_CAPABILITY("unsupported_capability"), + CANNOT_CONTINUE("cannot_continue"); + + private final String value; + + AiWorkflowOutcome(String value) { + this.value = value; + } + + @JsonValue + public String getValue() { + return value; + } + + @JsonCreator + public static AiWorkflowOutcome fromValue(String value) { + for (AiWorkflowOutcome outcome : values()) { + if (outcome.value.equals(value)) { + return outcome; + } + } + throw new IllegalArgumentException("Unknown AI workflow outcome: " + value); + } +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java new file mode 100644 index 000000000..22228d2aa --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowRequest.java @@ -0,0 +1,23 @@ +package stirling.software.proprietary.model.api.ai; + +import java.util.List; + +import io.swagger.v3.oas.annotations.media.Schema; + +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; + +import lombok.Data; + +@Data +@Schema(description = "Run an AI workflow against one or more PDF files") +public class AiWorkflowRequest { + + @NotNull + @Schema(description = "The input PDF files") + private List fileInputs; + + @NotBlank + @Schema(description = "The user message to orchestrate", example = "Summarise these documents") + private String userMessage; +} 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 new file mode 100644 index 000000000..1e04bece8 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowResponse.java @@ -0,0 +1,58 @@ +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; + +@Data +@Schema(description = "Structured AI workflow result") +public class AiWorkflowResponse { + + @Schema(description = "Workflow outcome") + private AiWorkflowOutcome outcome; + + @Schema(description = "Answer returned by the AI workflow when applicable") + private String answer; + + @Schema(description = "Summary returned by the AI workflow when applicable") + private String summary; + + @Schema(description = "Rationale returned by the AI workflow when applicable") + private String rationale; + + @Schema(description = "Reason when the AI workflow cannot proceed") + private String reason; + + @Schema(description = "Clarification question for the user when more input is required") + private String question; + + @Schema( + description = + "Unsupported capability identifier when the workflow cannot route the request") + private String capability; + + @Schema(description = "Message returned for unsupported capability outcomes") + private String message; + + @Schema(description = "Supporting evidence snippets from extracted PDF text") + private List evidence = new ArrayList<>(); + + @Schema(description = "Structured tool steps when the workflow returns a plan") + private List> steps = new ArrayList<>(); + + @Schema(description = "Per-file text extraction requests from the AI engine") + private List files = new ArrayList<>(); + + @Schema(description = "Maximum number of pages the AI engine wants text extracted from") + private Integer maxPages; + + @Schema(description = "Maximum number of characters the AI engine wants extracted") + private Integer maxCharacters; + + @Schema(description = "AI engine capability to resume with on the next turn") + private String resumeWith; +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowTextSelection.java b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowTextSelection.java new file mode 100644 index 000000000..265d989da --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiWorkflowTextSelection.java @@ -0,0 +1,16 @@ +package stirling.software.proprietary.model.api.ai; + +import io.swagger.v3.oas.annotations.media.Schema; + +import lombok.Data; + +@Data +@Schema(description = "Page-scoped extracted text selection") +public class AiWorkflowTextSelection { + + @Schema(description = "1-based page number", example = "2") + private Integer pageNumber; + + @Schema(description = "Extracted text or evidence snippet") + private String text; +} 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 new file mode 100644 index 000000000..753331b12 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java @@ -0,0 +1,108 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.time.Duration; + +import org.springframework.http.HttpStatus; +import org.springframework.stereotype.Service; +import org.springframework.web.server.ResponseStatusException; + +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.model.ApplicationProperties; + +@Slf4j +@Service +public class AiEngineClient { + + private final ApplicationProperties applicationProperties; + private final HttpClient httpClient; + + public AiEngineClient(ApplicationProperties applicationProperties) { + this.applicationProperties = applicationProperties; + this.httpClient = + HttpClient.newBuilder() + .connectTimeout( + Duration.ofSeconds( + applicationProperties.getAiEngine().getTimeoutSeconds())) + .build(); + } + + public String post(String path, String jsonBody) 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); + + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url)) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .timeout(Duration.ofSeconds(config.getTimeoutSeconds())) + .POST(HttpRequest.BodyPublishers.ofString(jsonBody)) + .build(); + + HttpResponse response = sendRequest(request); + + log.debug("AI engine responded with status {}", response.statusCode()); + checkResponseStatus(response); + return response.body(); + } + + public String get(String path) 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 GET request to {}", url); + + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url)) + .header("Accept", "application/json") + .timeout(Duration.ofSeconds(config.getTimeoutSeconds())) + .GET() + .build(); + + HttpResponse response = sendRequest(request); + + log.debug("AI engine responded with status {}", response.statusCode()); + checkResponseStatus(response); + return response.body(); + } + + private HttpResponse sendRequest(HttpRequest request) throws IOException { + try { + return httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new ResponseStatusException( + HttpStatus.SERVICE_UNAVAILABLE, "AI engine request was interrupted"); + } + } + + private void checkResponseStatus(HttpResponse response) { + int status = response.statusCode(); + if (status >= 500) { + throw new ResponseStatusException( + HttpStatus.BAD_GATEWAY, "AI engine returned error: " + status); + } + if (status >= 400) { + throw new ResponseStatusException( + HttpStatus.valueOf(status), + "AI engine returned client error: " + response.body()); + } + } +} 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 new file mode 100644 index 000000000..817d2e4a2 --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java @@ -0,0 +1,181 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import org.apache.pdfbox.pdmodel.PDDocument; +import org.springframework.stereotype.Service; +import org.springframework.web.multipart.MultipartFile; + +import lombok.Data; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.service.CustomPDFDocumentFactory; +import stirling.software.common.util.ExceptionUtils; +import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput; +import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest; +import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome; +import stirling.software.proprietary.model.api.ai.AiWorkflowRequest; +import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; +import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile; +import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult; +import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact; + +import tools.jackson.databind.ObjectMapper; + +@Slf4j +@Service +@RequiredArgsConstructor +public class AiWorkflowService { + + private final CustomPDFDocumentFactory pdfDocumentFactory; + private final AiEngineClient aiEngineClient; + private final PdfContentExtractor pdfContentExtractor; + private final ObjectMapper objectMapper; + + private sealed interface WorkflowState { + record Pending(WorkflowTurnRequest request) implements WorkflowState {} + + record Terminal(AiWorkflowResponse response) implements WorkflowState {} + } + + public AiWorkflowResponse orchestrate(AiWorkflowRequest request) throws IOException { + validateRequest(request); + + Map filesByName = new LinkedHashMap<>(); + for (AiWorkflowFileInput fileInput : request.getFileInputs()) { + filesByName.put( + fileInput.getFileInput().getOriginalFilename(), fileInput.getFileInput()); + } + + WorkflowTurnRequest initialRequest = new WorkflowTurnRequest(); + initialRequest.setUserMessage(request.getUserMessage().trim()); + initialRequest.setFileNames(new ArrayList<>(filesByName.keySet())); + + WorkflowState state = new WorkflowState.Pending(initialRequest); + while (state instanceof WorkflowState.Pending pending) { + state = advance(pending.request(), filesByName); + } + return ((WorkflowState.Terminal) state).response(); + } + + private WorkflowState advance( + WorkflowTurnRequest request, Map filesByName) + throws IOException { + AiWorkflowResponse response = invokeOrchestrator(request); + return switch (response.getOutcome()) { + case NEED_CONTENT -> onNeedContent(response, filesByName, request); + case ANSWER, + NOT_FOUND, + PLAN, + NEED_CLARIFICATION, + CANNOT_DO, + TOOL_CALL, + COMPLETED, + UNSUPPORTED_CAPABILITY, + CANNOT_CONTINUE -> + new WorkflowState.Terminal(response); + }; + } + + private WorkflowState onNeedContent( + AiWorkflowResponse response, + Map filesByName, + WorkflowTurnRequest request) + throws IOException { + if (!request.getArtifacts().isEmpty()) { + return new WorkflowState.Terminal( + cannotContinue("AI engine requested content extraction more than once.")); + } + + List requestedFiles = response.getFiles(); + + // Validate requested file names before loading anything + if (requestedFiles != null && !requestedFiles.isEmpty()) { + for (AiWorkflowFileRequest fileReq : requestedFiles) { + if (!filesByName.containsKey(fileReq.getFileName())) { + return new WorkflowState.Terminal( + cannotContinue( + "AI engine requested unknown file: " + fileReq.getFileName())); + } + } + } + + List fileNamesToLoad = + (requestedFiles == null || requestedFiles.isEmpty()) + ? new ArrayList<>(filesByName.keySet()) + : requestedFiles.stream().map(AiWorkflowFileRequest::getFileName).toList(); + + Map requestedByName = + requestedFiles == null || requestedFiles.isEmpty() + ? Map.of() + : requestedFiles.stream() + .collect( + Collectors.toMap( + AiWorkflowFileRequest::getFileName, r -> r)); + + List loadedFiles = new ArrayList<>(); + try { + for (String fileName : fileNamesToLoad) { + PDDocument doc = pdfDocumentFactory.load(filesByName.get(fileName), true); + loadedFiles.add(new LoadedFile(fileName, doc)); + } + + List contentResults = + pdfContentExtractor.extractContent( + loadedFiles, + requestedByName, + response.getMaxPages(), + response.getMaxCharacters()); + + WorkflowTurnRequest nextRequest = new WorkflowTurnRequest(); + nextRequest.setUserMessage(request.getUserMessage()); + nextRequest.setFileNames(request.getFileNames()); + nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults)); + nextRequest.setResumeWith(response.getResumeWith()); + return new WorkflowState.Pending(nextRequest); + } finally { + for (LoadedFile lf : loadedFiles) { + try { + lf.document().close(); + } catch (IOException e) { + log.warn("Failed to close PDF document: {}", lf.fileName(), e); + } + } + } + } + + private void validateRequest(AiWorkflowRequest request) { + for (AiWorkflowFileInput fileInput : request.getFileInputs()) { + if (fileInput.getFileInput().isEmpty()) { + throw ExceptionUtils.createFileNullOrEmptyException(); + } + } + } + + private AiWorkflowResponse cannotContinue(String reason) { + AiWorkflowResponse response = new AiWorkflowResponse(); + response.setOutcome(AiWorkflowOutcome.CANNOT_CONTINUE); + response.setReason(reason); + return response; + } + + private AiWorkflowResponse invokeOrchestrator(WorkflowTurnRequest request) throws IOException { + String requestBody = objectMapper.writeValueAsString(request); + String responseBody = aiEngineClient.post("/api/v1/orchestrator", requestBody); + return objectMapper.readValue(responseBody, AiWorkflowResponse.class); + } + + @Data + private static class WorkflowTurnRequest { + private String userMessage; + private List fileNames = new ArrayList<>(); + private List artifacts = new ArrayList<>(); + private String resumeWith; + } +} 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 new file mode 100644 index 000000000..be3f86a5a --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/service/PdfContentExtractor.java @@ -0,0 +1,279 @@ +package stirling.software.proprietary.service; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +import org.apache.pdfbox.pdmodel.PDDocument; +import org.apache.pdfbox.text.PDFTextStripper; +import org.springframework.stereotype.Service; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonValue; + +import lombok.Data; +import lombok.extern.slf4j.Slf4j; + +import stirling.software.common.util.ExceptionUtils; +import stirling.software.proprietary.model.api.ai.AiPdfContentType; +import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest; +import stirling.software.proprietary.model.api.ai.AiWorkflowTextSelection; + +@Slf4j +@Service +public class PdfContentExtractor { + + private static final int MAX_CHARACTERS_PER_PAGE = 4_000; + + record LoadedFile(String fileName, PDDocument document) {} + + /** + * Extracts content from the loaded files according to the requested content types and budget + * constraints. + */ + List extractContent( + List loadedFiles, + Map requestedByName, + int maxPages, + int maxCharacters) + throws IOException { + List contentResults = new ArrayList<>(); + int remainingPages = maxPages; + int remainingCharacters = maxCharacters; + + for (LoadedFile lf : loadedFiles) { + if (remainingPages <= 0 || remainingCharacters <= 0) break; + AiWorkflowFileRequest fileReq = requestedByName.get(lf.fileName()); + List contentTypes = + fileReq != null && !fileReq.getContentTypes().isEmpty() + ? fileReq.getContentTypes() + : List.of(AiPdfContentType.PAGE_TEXT); + + for (AiPdfContentType contentType : contentTypes) { + Optional result = + dispatchContentType( + contentType, lf, fileReq, remainingPages, remainingCharacters); + if (result.isPresent()) { + PdfContentResult content = result.get(); + contentResults.add(content); + remainingPages -= content.pagesConsumed(); + remainingCharacters -= content.charactersConsumed(); + } + } + } + return contentResults; + } + + /** Groups content results by artifact kind and builds the corresponding workflow artifacts. */ + List buildArtifacts(List results) { + List artifacts = new ArrayList<>(); + Map> byKind = + results.stream().collect(Collectors.groupingBy(PdfContentResult::getArtifactKind)); + for (var entry : byKind.entrySet()) { + artifacts.add(buildArtifact(entry.getKey(), entry.getValue())); + } + return artifacts; + } + + private Optional dispatchContentType( + AiPdfContentType contentType, + LoadedFile lf, + AiWorkflowFileRequest fileReq, + int remainingPages, + int remainingCharacters) + throws IOException { + return switch (contentType) { + case PAGE_TEXT, FULL_TEXT -> + Optional.ofNullable( + extractText(lf, fileReq, remainingPages, remainingCharacters)); + default -> { + log.warn( + "Content type {} not yet implemented, skipping for {}", + contentType, + lf.fileName()); + yield Optional.empty(); + } + }; + } + + private ExtractedFileText extractText( + LoadedFile lf, + AiWorkflowFileRequest fileReq, + int remainingPages, + int remainingCharacters) + throws IOException { + List requestedPages = fileReq != null ? fileReq.getPageNumbers() : null; + List pages = + selectPages(lf.document().getNumberOfPages(), requestedPages, remainingPages); + List extracted = + extractPageText(lf.document(), pages, remainingCharacters); + return extracted.isEmpty() ? null : buildExtractedFileText(lf.fileName(), extracted); + } + + private WorkflowArtifact buildArtifact(ArtifactKind kind, List results) { + return switch (kind) { + case EXTRACTED_TEXT -> { + ExtractedTextArtifact artifact = new ExtractedTextArtifact(); + artifact.setFiles(results.stream().map(ExtractedFileText.class::cast).toList()); + yield artifact; + } + }; + } + + private List selectPages( + int totalPages, List requestedPageNumbers, int maxPages) { + if (totalPages <= 0) { + throw ExceptionUtils.createPdfNoPages(); + } + + List pages = new ArrayList<>(); + + if (requestedPageNumbers == null || requestedPageNumbers.isEmpty()) { + for (int p = 1; p <= totalPages && pages.size() < maxPages; p++) { + pages.add(p); + } + return pages; + } + + Set deduplicatedPages = new LinkedHashSet<>(requestedPageNumbers); + for (Integer pageNumber : deduplicatedPages) { + if (pageNumber == null || pageNumber < 1 || pageNumber > totalPages) { + throw ExceptionUtils.createIllegalArgumentException( + "error.invalidPageNumber", + "Requested page number %s is outside the PDF page range.", + pageNumber); + } + pages.add(pageNumber); + if (pages.size() >= maxPages) { + break; + } + } + return pages; + } + + private List extractPageText( + PDDocument document, List selectedPages, int maxCharacters) + throws IOException { + PDFTextStripper textStripper = new PDFTextStripper(); + List pages = new ArrayList<>(); + int remainingCharacters = maxCharacters; + + for (Integer pageNumber : selectedPages) { + if (remainingCharacters <= 0) { + break; + } + + textStripper.setStartPage(pageNumber); + textStripper.setEndPage(pageNumber); + + String pageText = textStripper.getText(document).trim(); + if (pageText.isBlank()) { + continue; + } + + int allowedCharacters = Math.min(remainingCharacters, MAX_CHARACTERS_PER_PAGE); + String clippedText = clip(pageText, allowedCharacters); + if (clippedText.isBlank()) { + continue; + } + + AiWorkflowTextSelection selection = new AiWorkflowTextSelection(); + selection.setPageNumber(pageNumber); + selection.setText(clippedText); + pages.add(selection); + remainingCharacters -= clippedText.length(); + } + return pages; + } + + private ExtractedFileText buildExtractedFileText( + String fileName, List pages) { + ExtractedFileText fileText = new ExtractedFileText(); + fileText.setFileName(fileName); + fileText.setPages(pages); + return fileText; + } + + private String clip(String text, int maxLength) { + if (text.length() <= maxLength) { + return text; + } + // Avoid splitting a surrogate pair at the boundary + int end = maxLength; + if (Character.isHighSurrogate(text.charAt(end - 1))) { + end--; + } + return text.substring(0, end); + } + + // --- Types shared with AiWorkflowService (package-private) --- + + interface PdfContentResult { + @JsonIgnore + ArtifactKind getArtifactKind(); + + @JsonIgnore + default int pagesConsumed() { + return 0; + } + + @JsonIgnore + default int charactersConsumed() { + return 0; + } + } + + /** + * Values MUST match {@code ArtifactKind} in {@code engine/src/stirling/contracts/common.py}. + */ + enum ArtifactKind { + EXTRACTED_TEXT("extracted_text"); + + private final String value; + + ArtifactKind(String value) { + this.value = value; + } + + @JsonValue + public String getValue() { + return value; + } + } + + interface WorkflowArtifact { + ArtifactKind getKind(); + } + + @Data + static class ExtractedFileText implements PdfContentResult { + private String fileName; + private List pages = new ArrayList<>(); + + @Override + public ArtifactKind getArtifactKind() { + return ArtifactKind.EXTRACTED_TEXT; + } + + @Override + public int pagesConsumed() { + return pages.size(); + } + + @Override + public int charactersConsumed() { + return pages.stream().mapToInt(p -> p.getText().length()).sum(); + } + } + + @Data + static final class ExtractedTextArtifact implements WorkflowArtifact { + private final ArtifactKind kind = ArtifactKind.EXTRACTED_TEXT; + private List files = new ArrayList<>(); + } +} diff --git a/engine/pyproject.toml b/engine/pyproject.toml index 693a28153..d8c665f37 100644 --- a/engine/pyproject.toml +++ b/engine/pyproject.toml @@ -42,9 +42,8 @@ select = [ "W", "RUF100", "UP", -] -ignore = [ - "E501", # Temporarily disable line length limit until codebase conformat + "PYI", # flake8-pyi: flags deprecated typing constructs + "FA", # flake8-future-annotations: flags missing future annotations imports ] [tool.pyright] @@ -55,6 +54,7 @@ reportUnnecessaryCast = "warning" reportUnnecessaryTypeIgnoreComment = "warning" reportUnusedImport = "warning" reportUnknownParameterType = "warning" +reportDeprecated = "warning" [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/engine/src/stirling/agents/orchestrator.py b/engine/src/stirling/agents/orchestrator.py index 1ca37c23f..9b58b4540 100644 --- a/engine/src/stirling/agents/orchestrator.py +++ b/engine/src/stirling/agents/orchestrator.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import assert_never from pydantic_ai import Agent from pydantic_ai.output import ToolOutput @@ -12,12 +13,14 @@ from stirling.agents.user_spec import UserSpecAgent from stirling.contracts import ( AgentDraftRequest, AgentDraftWorkflowResponse, + ExtractedTextArtifact, OrchestratorRequest, OrchestratorResponse, PdfEditRequest, PdfEditResponse, PdfQuestionRequest, PdfQuestionResponse, + SupportedCapability, UnsupportedCapabilityResponse, ) from stirling.services import AppRuntime @@ -61,7 +64,7 @@ class OrchestratorAgent: "You are the top-level orchestrator. " "Choose exactly one output function that best handles the request. " "Use delegate_pdf_edit for requested PDF modifications. " - "Use delegate_pdf_question for questions about the contents of a PDF. " + "Use delegate_pdf_question for questions about PDF contents. " "Use delegate_user_spec for requests to create or define an agent spec. " "Use unsupported_capability only when none of the other outputs fit." ), @@ -69,27 +72,56 @@ class OrchestratorAgent: ) async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse: + if request.resume_with is not None: + return await self._resume(request, request.resume_with) result = await self.agent.run( - request.user_message, + self._build_prompt(request), deps=OrchestratorDeps(runtime=self.runtime, request=request), ) return result.output + async def _resume(self, request: OrchestratorRequest, capability: SupportedCapability) -> OrchestratorResponse: + """Fast-path to get back to the correct endpoint without having to call AI.""" + match capability: + case SupportedCapability.PDF_QUESTION: + return await self._run_pdf_question(request) + case SupportedCapability.PDF_EDIT: + return await self._run_pdf_edit(request) + case SupportedCapability.AGENT_DRAFT: + return await self._run_agent_draft(request) + case ( + SupportedCapability.ORCHESTRATE + | SupportedCapability.AGENT_REVISE + | SupportedCapability.AGENT_NEXT_ACTION + ): + raise ValueError(f"Cannot resume orchestrator with capability: {capability}") + case _ as unreachable: + assert_never(unreachable) + async def delegate_pdf_edit(self, ctx: RunContext[OrchestratorDeps]) -> PdfEditResponse: - request = ctx.deps.request - return await PdfEditAgent(ctx.deps.runtime).handle( - PdfEditRequest(user_message=request.user_message, conversation_id=request.conversation_id) - ) + return await self._run_pdf_edit(ctx.deps.request) + + async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse: + return await PdfEditAgent(self.runtime).handle(PdfEditRequest(user_message=request.user_message)) async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse: - request = ctx.deps.request - return await PdfQuestionAgent(ctx.deps.runtime).handle( - PdfQuestionRequest(question=request.user_message, conversation_id=request.conversation_id) + return await self._run_pdf_question(ctx.deps.request) + + async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse: + extracted_text = self._get_extracted_text_artifact(request) + return await PdfQuestionAgent(self.runtime).handle( + PdfQuestionRequest( + question=request.user_message, + file_names=request.file_names, + page_text=extracted_text.files if extracted_text is not None else [], + ) ) async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse: - request = ctx.deps.request - return await UserSpecAgent(ctx.deps.runtime).draft(AgentDraftRequest(user_message=request.user_message)) + return await self._run_agent_draft(ctx.deps.request) + + async def _run_agent_draft(self, request: OrchestratorRequest) -> AgentDraftWorkflowResponse: + return await UserSpecAgent(self.runtime).draft(AgentDraftRequest(user_message=request.user_message)) async def unsupported_capability( self, @@ -98,3 +130,28 @@ class OrchestratorAgent: message: str, ) -> UnsupportedCapabilityResponse: return UnsupportedCapabilityResponse(capability=capability, message=message) + + def _get_extracted_text_artifact(self, request: OrchestratorRequest) -> ExtractedTextArtifact | None: + for artifact in request.artifacts: + if isinstance(artifact, ExtractedTextArtifact): + return artifact + return None + + 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" + return f"User message: {request.user_message}\nFiles: {file_names}\nAvailable artifacts:\n{artifact_summary}" + + def _describe_artifacts(self, request: OrchestratorRequest) -> str: + if not request.artifacts: + return "- none" + + descriptions: list[str] = [] + for artifact in request.artifacts: + if isinstance(artifact, ExtractedTextArtifact): + total_pages = sum(len(f.pages) for f in artifact.files) + file_names = [f.file_name for f in artifact.files] + descriptions.append(f"- extracted_text: {total_pages} pages from {file_names}") + continue + descriptions.append("- unknown artifact") + return "\n".join(descriptions) diff --git a/engine/src/stirling/agents/pdf_questions.py b/engine/src/stirling/agents/pdf_questions.py index b7ca33ac9..c8a63bd70 100644 --- a/engine/src/stirling/agents/pdf_questions.py +++ b/engine/src/stirling/agents/pdf_questions.py @@ -4,8 +4,11 @@ from pydantic_ai import Agent from pydantic_ai.output import NativeOutput from stirling.contracts import ( + ExtractedFileText, + NeedContentFileRequest, + PdfContentType, PdfQuestionAnswerResponse, - PdfQuestionNeedTextResponse, + PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfQuestionResponse, @@ -14,6 +17,9 @@ from stirling.services import AppRuntime class PdfQuestionAgent: + DEFAULT_MAX_PAGES = 12 + DEFAULT_MAX_CHARACTERS = 24_000 + def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime self.agent = Agent( @@ -25,18 +31,27 @@ class PdfQuestionAgent: ] ), system_prompt=( - "Answer questions about a PDF using only the extracted text provided in the 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 copied from the provided text." + "When answering, include a short list of evidence snippets with their page numbers." ), model_settings=runtime.smart_model_settings, ) async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse: - if not request.extracted_text.strip(): - return PdfQuestionNeedTextResponse( - reason="No extracted PDF text was provided, so the question cannot be answered yet." + if not self._has_page_text(request.page_text): + return PdfQuestionNeedContentResponse( + 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.DEFAULT_MAX_PAGES, + max_characters=self.DEFAULT_MAX_CHARACTERS, ) return await self._run_answer_agent(request) @@ -45,5 +60,14 @@ class PdfQuestionAgent: return result.output def _build_prompt(self, request: PdfQuestionRequest) -> str: - file_name = request.file_name or "Unknown file" - return f"File: {file_name}\nQuestion: {request.question}\nExtracted text:\n{request.extracted_text}" + file_names = ", ".join(request.file_names) if request.file_names else "Unknown files" + sections = [ + f"[File: {file_text.file_name}, Page {selection.page_number or '?'}]\n{selection.text}" + for file_text in request.page_text + for selection in file_text.pages + ] + pages = "\n\n".join(sections) + return f"Files: {file_names}\nQuestion: {request.question}\nExtracted page text:\n{pages}" + + def _has_page_text(self, page_text: list[ExtractedFileText]) -> bool: + return any(selection.text.strip() for file_text in page_text for selection in file_text.pages) diff --git a/engine/src/stirling/contracts/__init__.py b/engine/src/stirling/contracts/__init__.py index 6988233c5..0593f4dd2 100644 --- a/engine/src/stirling/contracts/__init__.py +++ b/engine/src/stirling/contracts/__init__.py @@ -8,7 +8,17 @@ from .agent_drafts import ( AgentRevisionWorkflowResponse, ) from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep -from .common import ConversationMessage, PdfTextSelection, ToolOperationStep +from .common import ( + ArtifactKind, + ConversationMessage, + ExtractedFileText, + PdfContentType, + PdfTextSelection, + StepKind, + SupportedCapability, + ToolOperationStep, + WorkflowOutcome, +) from .execution import ( AgentExecutionRequest, CannotContinueExecutionAction, @@ -19,7 +29,13 @@ from .execution import ( ToolCallExecutionAction, ) from .health import HealthResponse -from .orchestrator import OrchestratorRequest, OrchestratorResponse, SupportedCapability, UnsupportedCapabilityResponse +from .orchestrator import ( + ExtractedTextArtifact, + OrchestratorRequest, + OrchestratorResponse, + UnsupportedCapabilityResponse, + WorkflowArtifact, +) from .pdf_edit import ( EditCannotDoResponse, EditClarificationRequest, @@ -28,14 +44,16 @@ from .pdf_edit import ( PdfEditResponse, ) from .pdf_questions import ( + NeedContentFileRequest, PdfQuestionAnswerResponse, - PdfQuestionNeedTextResponse, + PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, PdfQuestionResponse, ) __all__ = [ + "ArtifactKind", "AgentDraft", "AgentDraftRequest", "AgentDraftResponse", @@ -49,6 +67,7 @@ __all__ = [ "AiToolAgentStep", "CannotContinueExecutionAction", "ConversationMessage", + "ExtractedFileText", "CompletedExecutionAction", "EditCannotDoResponse", "EditClarificationRequest", @@ -56,19 +75,25 @@ __all__ = [ "ExecutionContext", "ExecutionStepResult", "HealthResponse", + "NeedContentFileRequest", "NextExecutionAction", + "ExtractedTextArtifact", "OrchestratorRequest", "OrchestratorResponse", "PdfEditRequest", "PdfEditResponse", "PdfQuestionAnswerResponse", "PdfQuestionNotFoundResponse", - "PdfQuestionNeedTextResponse", + "PdfContentType", + "PdfQuestionNeedContentResponse", "PdfQuestionRequest", "PdfQuestionResponse", "PdfTextSelection", + "StepKind", "SupportedCapability", "ToolOperationStep", "ToolCallExecutionAction", + "WorkflowOutcome", "UnsupportedCapabilityResponse", + "WorkflowArtifact", ] diff --git a/engine/src/stirling/contracts/agent_drafts.py b/engine/src/stirling/contracts/agent_drafts.py index a59d279f2..752019d34 100644 --- a/engine/src/stirling/contracts/agent_drafts.py +++ b/engine/src/stirling/contracts/agent_drafts.py @@ -7,12 +7,12 @@ from pydantic import Field from stirling.models import ApiModel from .agent_specs import AgentSpecStep -from .common import ConversationMessage +from .common import ConversationMessage, StepKind, WorkflowOutcome from .pdf_edit import EditCannotDoResponse, EditClarificationRequest class AgentDraftStep(ApiModel): - kind: Literal["tool", "ai_tool"] + kind: Literal[StepKind.TOOL, StepKind.AI_TOOL] title: str description: str @@ -30,7 +30,7 @@ class AgentDraftRequest(ApiModel): class AgentDraftResponse(ApiModel): - outcome: Literal["draft"] = "draft" + outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT draft: AgentDraft @@ -41,7 +41,7 @@ class AgentRevisionRequest(ApiModel): class AgentRevisionResponse(ApiModel): - outcome: Literal["draft"] = "draft" + outcome: Literal[WorkflowOutcome.DRAFT] = WorkflowOutcome.DRAFT draft: AgentDraft diff --git a/engine/src/stirling/contracts/agent_specs.py b/engine/src/stirling/contracts/agent_specs.py index 403af1780..684872ea8 100644 --- a/engine/src/stirling/contracts/agent_specs.py +++ b/engine/src/stirling/contracts/agent_specs.py @@ -6,11 +6,11 @@ from pydantic import Field from stirling.models import ApiModel, OperationId -from .common import ToolOperationStep +from .common import StepKind, ToolOperationStep class AiToolAgentStep(ApiModel): - kind: Literal["ai_tool"] = "ai_tool" + kind: Literal[StepKind.AI_TOOL] = StepKind.AI_TOOL title: str description: str tool: OperationId diff --git a/engine/src/stirling/contracts/common.py b/engine/src/stirling/contracts/common.py index 65bb22400..7fd2af9aa 100644 --- a/engine/src/stirling/contracts/common.py +++ b/engine/src/stirling/contracts/common.py @@ -1,12 +1,89 @@ from __future__ import annotations +from enum import StrEnum from typing import Literal -from pydantic import model_validator +from pydantic import Field, model_validator from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel +class PdfContentType(StrEnum): + """Types of content that can be extracted from a PDF and sent to the AI. + + Java counterpart: AiPdfContentType.java - values must stay in sync. + """ + + # Document-level structured data + PAGE_LAYOUT = "page_layout" + DOCUMENT_METADATA = "document_metadata" + ENCRYPTION_INFO = "encryption_info" + BOOKMARKS = "bookmarks" + LAYERS = "layers" + EMBEDDED_FILES = "embedded_files" + JAVASCRIPT = "javascript" + LINKS = "links" + IMAGE_INFO = "image_info" + FONTS = "fonts" + + # Text and content + PAGE_TEXT = "page_text" + FULL_TEXT = "full_text" + FORM_FIELDS = "form_fields" + ANNOTATIONS = "annotations" + SIGNATURES = "signatures" + STRUCTURE_TREE = "structure_tree" + XMP_METADATA = "xmp_metadata" + + # Heavy content + COMPLIANCE = "compliance" + IMAGES = "images" + + +class WorkflowOutcome(StrEnum): + """Discriminator values for all workflow response unions (outcome field). + + Java counterpart: AiWorkflowOutcome.java - values must stay in sync. + """ + + ANSWER = "answer" + NEED_CONTENT = "need_content" + NOT_FOUND = "not_found" + PLAN = "plan" + NEED_CLARIFICATION = "need_clarification" + CANNOT_DO = "cannot_do" + DRAFT = "draft" + TOOL_CALL = "tool_call" + COMPLETED = "completed" + CANNOT_CONTINUE = "cannot_continue" + UNSUPPORTED_CAPABILITY = "unsupported_capability" + + +class ArtifactKind(StrEnum): + """Discriminator values for WorkflowArtifact unions (kind field). + + Java counterpart: PdfContentExtractor.ArtifactKind - values must stay in sync. + """ + + EXTRACTED_TEXT = "extracted_text" + + +class StepKind(StrEnum): + """Discriminator values for AgentSpecStep unions (kind field).""" + + TOOL = "tool" + AI_TOOL = "ai_tool" + + +class SupportedCapability(StrEnum): + ORCHESTRATE = "orchestrate" + PDF_EDIT = "pdf_edit" + PDF_QUESTION = "pdf_question" + AGENT_DRAFT = "agent_draft" + AGENT_REVISE = "agent_revise" + AGENT_NEXT_ACTION = "agent_next_action" + + class ConversationMessage(ApiModel): role: str content: str @@ -17,8 +94,13 @@ class PdfTextSelection(ApiModel): text: str +class ExtractedFileText(ApiModel): + file_name: str + pages: list[PdfTextSelection] = Field(default_factory=list) + + class ToolOperationStep(ApiModel): - kind: Literal["tool"] = "tool" + kind: Literal[StepKind.TOOL] = StepKind.TOOL tool: OperationId parameters: ParamToolModel diff --git a/engine/src/stirling/contracts/execution.py b/engine/src/stirling/contracts/execution.py index 64e70d682..6e1f9ad34 100644 --- a/engine/src/stirling/contracts/execution.py +++ b/engine/src/stirling/contracts/execution.py @@ -7,6 +7,7 @@ from pydantic import Field from stirling.models import ApiModel, OperationId, ParamToolModel from .agent_specs import AgentSpec +from .common import WorkflowOutcome class ExecutionStepResult(ApiModel): @@ -31,19 +32,19 @@ class AgentExecutionRequest(ApiModel): class ToolCallExecutionAction(ApiModel): - outcome: Literal["tool_call"] = "tool_call" + outcome: Literal[WorkflowOutcome.TOOL_CALL] = WorkflowOutcome.TOOL_CALL tool: OperationId parameters: ParamToolModel rationale: str | None = None class CompletedExecutionAction(ApiModel): - outcome: Literal["completed"] = "completed" + outcome: Literal[WorkflowOutcome.COMPLETED] = WorkflowOutcome.COMPLETED summary: str class CannotContinueExecutionAction(ApiModel): - outcome: Literal["cannot_continue"] = "cannot_continue" + outcome: Literal[WorkflowOutcome.CANNOT_CONTINUE] = WorkflowOutcome.CANNOT_CONTINUE reason: str diff --git a/engine/src/stirling/contracts/orchestrator.py b/engine/src/stirling/contracts/orchestrator.py index 563341880..3fa788fe6 100644 --- a/engine/src/stirling/contracts/orchestrator.py +++ b/engine/src/stirling/contracts/orchestrator.py @@ -1,6 +1,5 @@ from __future__ import annotations -from enum import StrEnum from typing import Annotated, Literal from pydantic import Field @@ -8,27 +7,29 @@ from pydantic import Field from stirling.models import ApiModel from .agent_drafts import AgentDraftResponse +from .common import ArtifactKind, ExtractedFileText, SupportedCapability, WorkflowOutcome from .execution import NextExecutionAction from .pdf_edit import PdfEditResponse from .pdf_questions import PdfQuestionResponse -class SupportedCapability(StrEnum): - ORCHESTRATE = "orchestrate" - PDF_EDIT = "pdf_edit" - PDF_QUESTION = "pdf_question" - AGENT_DRAFT = "agent_draft" - AGENT_REVISE = "agent_revise" - AGENT_NEXT_ACTION = "agent_next_action" +class ExtractedTextArtifact(ApiModel): + kind: Literal[ArtifactKind.EXTRACTED_TEXT] = ArtifactKind.EXTRACTED_TEXT + files: list[ExtractedFileText] = Field(default_factory=list) + + +WorkflowArtifact = Annotated[ExtractedTextArtifact, Field(discriminator="kind")] class OrchestratorRequest(ApiModel): user_message: str - conversation_id: str | None = None + file_names: list[str] + artifacts: list[WorkflowArtifact] = Field(default_factory=list) + resume_with: SupportedCapability | None = None class UnsupportedCapabilityResponse(ApiModel): - outcome: Literal["unsupported_capability"] = "unsupported_capability" + outcome: Literal[WorkflowOutcome.UNSUPPORTED_CAPABILITY] = WorkflowOutcome.UNSUPPORTED_CAPABILITY capability: str message: str diff --git a/engine/src/stirling/contracts/pdf_edit.py b/engine/src/stirling/contracts/pdf_edit.py index e3e1e6057..2bcfe7ac6 100644 --- a/engine/src/stirling/contracts/pdf_edit.py +++ b/engine/src/stirling/contracts/pdf_edit.py @@ -6,30 +6,29 @@ from pydantic import Field from stirling.models import ApiModel -from .common import ToolOperationStep +from .common import ToolOperationStep, WorkflowOutcome class PdfEditRequest(ApiModel): user_message: str - conversation_id: str | None = None file_names: list[str] = Field(default_factory=list) class EditPlanResponse(ApiModel): - outcome: Literal["plan"] = "plan" + outcome: Literal[WorkflowOutcome.PLAN] = WorkflowOutcome.PLAN summary: str rationale: str | None = None steps: list[ToolOperationStep] class EditClarificationRequest(ApiModel): - outcome: Literal["need_clarification"] = "need_clarification" + outcome: Literal[WorkflowOutcome.NEED_CLARIFICATION] = WorkflowOutcome.NEED_CLARIFICATION question: str reason: str class EditCannotDoResponse(ApiModel): - outcome: Literal["cannot_do"] = "cannot_do" + outcome: Literal[WorkflowOutcome.CANNOT_DO] = WorkflowOutcome.CANNOT_DO reason: str diff --git a/engine/src/stirling/contracts/pdf_questions.py b/engine/src/stirling/contracts/pdf_questions.py index 3ac12bdfe..987dc9b59 100644 --- a/engine/src/stirling/contracts/pdf_questions.py +++ b/engine/src/stirling/contracts/pdf_questions.py @@ -6,31 +6,42 @@ from pydantic import Field from stirling.models import ApiModel +from .common import ExtractedFileText, PdfContentType, SupportedCapability, WorkflowOutcome + class PdfQuestionRequest(ApiModel): question: str - conversation_id: str | None = None - extracted_text: str = "" - file_name: str | None = None + page_text: list[ExtractedFileText] = Field(default_factory=list) + file_names: list[str] class PdfQuestionAnswerResponse(ApiModel): - outcome: Literal["answer"] = "answer" + outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER answer: str - evidence: list[str] = Field(default_factory=list) + evidence: list[ExtractedFileText] = Field(default_factory=list) -class PdfQuestionNeedTextResponse(ApiModel): - outcome: Literal["need_text"] = "need_text" +class NeedContentFileRequest(ApiModel): + file_name: str + page_numbers: list[int] = Field(default_factory=list) + content_types: list[PdfContentType] + + +class PdfQuestionNeedContentResponse(ApiModel): + outcome: Literal[WorkflowOutcome.NEED_CONTENT] = WorkflowOutcome.NEED_CONTENT + resume_with: SupportedCapability = SupportedCapability.PDF_QUESTION reason: str + files: list[NeedContentFileRequest] = Field(default_factory=list) + max_pages: int + max_characters: int class PdfQuestionNotFoundResponse(ApiModel): - outcome: Literal["not_found"] = "not_found" + outcome: Literal[WorkflowOutcome.NOT_FOUND] = WorkflowOutcome.NOT_FOUND reason: str PdfQuestionResponse = Annotated[ - PdfQuestionAnswerResponse | PdfQuestionNeedTextResponse | PdfQuestionNotFoundResponse, + PdfQuestionAnswerResponse | PdfQuestionNeedContentResponse | PdfQuestionNotFoundResponse, Field(discriminator="outcome"), ] diff --git a/engine/tests/test_pdf_question_agent.py b/engine/tests/test_pdf_question_agent.py index 52df9474c..a4a1ce7bf 100644 --- a/engine/tests/test_pdf_question_agent.py +++ b/engine/tests/test_pdf_question_agent.py @@ -5,10 +5,12 @@ import pytest from stirling.agents import PdfQuestionAgent from stirling.config import AppSettings from stirling.contracts import ( + ExtractedFileText, PdfQuestionAnswerResponse, - PdfQuestionNeedTextResponse, + PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, + PdfTextSelection, ) from stirling.services import build_runtime @@ -34,13 +36,22 @@ def build_test_settings() -> AppSettings: ) +def invoice_page() -> ExtractedFileText: + return ExtractedFileText( + file_name="invoice.pdf", + pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")], + ) + + @pytest.mark.anyio async def test_pdf_question_agent_requires_extracted_text() -> None: agent = PdfQuestionAgent(build_runtime(build_test_settings())) - response = await agent.handle(PdfQuestionRequest(question="What is the total?", extracted_text="")) + response = await agent.handle( + PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"]) + ) - assert isinstance(response, PdfQuestionNeedTextResponse) + assert isinstance(response, PdfQuestionNeedContentResponse) @pytest.mark.anyio @@ -48,15 +59,15 @@ async def test_pdf_question_agent_returns_grounded_answer() -> None: agent = StubPdfQuestionAgent( PdfQuestionAnswerResponse( answer="The invoice total is 120.00.", - evidence=["Invoice total: 120.00"], + evidence=[invoice_page()], ) ) response = await agent.handle( PdfQuestionRequest( question="What is the total?", - extracted_text="Invoice total: 120.00", - file_name="invoice.pdf", + page_text=[invoice_page()], + file_names=["invoice.pdf"], ) ) @@ -71,8 +82,13 @@ async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient() response = await agent.handle( PdfQuestionRequest( question="What is the total?", - extracted_text="This page contains only a shipping address.", - file_name="invoice.pdf", + page_text=[ + ExtractedFileText( + file_name="invoice.pdf", + pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")], + ) + ], + file_names=["invoice.pdf"], ) ) diff --git a/engine/tests/test_stirling_api.py b/engine/tests/test_stirling_api.py index 9191eeb44..774fb28d9 100644 --- a/engine/tests/test_stirling_api.py +++ b/engine/tests/test_stirling_api.py @@ -20,9 +20,9 @@ from stirling.contracts import ( EditCannotDoResponse, OrchestratorRequest, PdfEditRequest, + PdfQuestionNeedContentResponse, PdfQuestionNotFoundResponse, PdfQuestionRequest, - UnsupportedCapabilityResponse, ) from stirling.models.tool_models import RotateParams @@ -38,8 +38,8 @@ class StubSettingsProvider: class StubOrchestratorAgent: - async def handle(self, request: OrchestratorRequest) -> UnsupportedCapabilityResponse: - return UnsupportedCapabilityResponse(capability="pdf_edit", message=request.user_message) + async def handle(self, request: OrchestratorRequest) -> PdfQuestionNeedContentResponse: + return PdfQuestionNeedContentResponse(reason=request.user_message, files=[], max_pages=1, max_characters=1000) class StubPdfEditAgent: @@ -115,10 +115,10 @@ def test_health_route() -> None: def test_orchestrator_route() -> None: - response = client.post("/api/v1/orchestrator", json={"userMessage": "route this"}) + response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]}) assert response.status_code == 200 - assert response.json()["outcome"] == "unsupported_capability" + assert response.json()["outcome"] == "need_content" def test_pdf_edit_route() -> None: @@ -129,7 +129,14 @@ def test_pdf_edit_route() -> None: def test_pdf_questions_route() -> None: - response = client.post("/api/v1/pdf/questions", json={"question": "what is this?"}) + response = client.post( + "/api/v1/pdf/questions", + json={ + "question": "what is this?", + "fileNames": ["test.pdf"], + "pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}], + }, + ) assert response.status_code == 200 assert response.json()["outcome"] == "not_found" diff --git a/engine/tests/test_stirling_contracts.py b/engine/tests/test_stirling_contracts.py index 7cf7df974..63e283a64 100644 --- a/engine/tests/test_stirling_contracts.py +++ b/engine/tests/test_stirling_contracts.py @@ -9,17 +9,34 @@ from stirling.contracts import ( AgentSpecStep, EditPlanResponse, ExecutionContext, + ExtractedFileText, + ExtractedTextArtifact, OrchestratorRequest, PdfQuestionAnswerResponse, + PdfTextSelection, ToolOperationStep, ) from stirling.models.tool_models import OperationId, RotateParams def test_orchestrator_request_accepts_user_message() -> None: - request = OrchestratorRequest(user_message="Rotate the PDF") + request = OrchestratorRequest( + user_message="Rotate the PDF", + file_names=["test.pdf"], + artifacts=[ + ExtractedTextArtifact( + files=[ + ExtractedFileText( + file_name="test.pdf", + pages=[PdfTextSelection(page_number=1, text="Hello")], + ) + ] + ) + ], + ) assert request.user_message == "Rotate the PDF" + assert len(request.artifacts) == 1 def test_agent_execution_request_uses_typed_agent_spec() -> None: