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: