mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 19:33:11 +02:00
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.
This commit is contained in:
James Brunton
+83
@@ -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<String> 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<AiWorkflowResponse> 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<String> 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");
|
||||
}
|
||||
}
|
||||
}
|
||||
+57
@@ -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.
|
||||
*
|
||||
* <p>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);
|
||||
}
|
||||
}
|
||||
+22
@@ -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;
|
||||
}
|
||||
+22
@@ -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<Integer> pageNumbers = new ArrayList<>();
|
||||
|
||||
@Schema(description = "Content types to extract from this file")
|
||||
private List<AiPdfContentType> contentTypes = new ArrayList<>();
|
||||
}
|
||||
+43
@@ -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.
|
||||
*
|
||||
* <p>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);
|
||||
}
|
||||
}
|
||||
+23
@@ -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<AiWorkflowFileInput> fileInputs;
|
||||
|
||||
@NotBlank
|
||||
@Schema(description = "The user message to orchestrate", example = "Summarise these documents")
|
||||
private String userMessage;
|
||||
}
|
||||
+58
@@ -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<AiWorkflowTextSelection> evidence = new ArrayList<>();
|
||||
|
||||
@Schema(description = "Structured tool steps when the workflow returns a plan")
|
||||
private List<Map<String, Object>> steps = new ArrayList<>();
|
||||
|
||||
@Schema(description = "Per-file text extraction requests from the AI engine")
|
||||
private List<AiWorkflowFileRequest> 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;
|
||||
}
|
||||
+16
@@ -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;
|
||||
}
|
||||
+108
@@ -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<String> 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<String> response = sendRequest(request);
|
||||
|
||||
log.debug("AI engine responded with status {}", response.statusCode());
|
||||
checkResponseStatus(response);
|
||||
return response.body();
|
||||
}
|
||||
|
||||
private HttpResponse<String> 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<String> 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
+181
@@ -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<String, MultipartFile> 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<String, MultipartFile> 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<String, MultipartFile> filesByName,
|
||||
WorkflowTurnRequest request)
|
||||
throws IOException {
|
||||
if (!request.getArtifacts().isEmpty()) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue("AI engine requested content extraction more than once."));
|
||||
}
|
||||
|
||||
List<AiWorkflowFileRequest> 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<String> fileNamesToLoad =
|
||||
(requestedFiles == null || requestedFiles.isEmpty())
|
||||
? new ArrayList<>(filesByName.keySet())
|
||||
: requestedFiles.stream().map(AiWorkflowFileRequest::getFileName).toList();
|
||||
|
||||
Map<String, AiWorkflowFileRequest> requestedByName =
|
||||
requestedFiles == null || requestedFiles.isEmpty()
|
||||
? Map.of()
|
||||
: requestedFiles.stream()
|
||||
.collect(
|
||||
Collectors.toMap(
|
||||
AiWorkflowFileRequest::getFileName, r -> r));
|
||||
|
||||
List<LoadedFile> loadedFiles = new ArrayList<>();
|
||||
try {
|
||||
for (String fileName : fileNamesToLoad) {
|
||||
PDDocument doc = pdfDocumentFactory.load(filesByName.get(fileName), true);
|
||||
loadedFiles.add(new LoadedFile(fileName, doc));
|
||||
}
|
||||
|
||||
List<PdfContentResult> 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<String> fileNames = new ArrayList<>();
|
||||
private List<WorkflowArtifact> artifacts = new ArrayList<>();
|
||||
private String resumeWith;
|
||||
}
|
||||
}
|
||||
+279
@@ -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<PdfContentResult> extractContent(
|
||||
List<LoadedFile> loadedFiles,
|
||||
Map<String, AiWorkflowFileRequest> requestedByName,
|
||||
int maxPages,
|
||||
int maxCharacters)
|
||||
throws IOException {
|
||||
List<PdfContentResult> 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<AiPdfContentType> contentTypes =
|
||||
fileReq != null && !fileReq.getContentTypes().isEmpty()
|
||||
? fileReq.getContentTypes()
|
||||
: List.of(AiPdfContentType.PAGE_TEXT);
|
||||
|
||||
for (AiPdfContentType contentType : contentTypes) {
|
||||
Optional<PdfContentResult> 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<WorkflowArtifact> buildArtifacts(List<PdfContentResult> results) {
|
||||
List<WorkflowArtifact> artifacts = new ArrayList<>();
|
||||
Map<ArtifactKind, List<PdfContentResult>> byKind =
|
||||
results.stream().collect(Collectors.groupingBy(PdfContentResult::getArtifactKind));
|
||||
for (var entry : byKind.entrySet()) {
|
||||
artifacts.add(buildArtifact(entry.getKey(), entry.getValue()));
|
||||
}
|
||||
return artifacts;
|
||||
}
|
||||
|
||||
private Optional<PdfContentResult> dispatchContentType(
|
||||
AiPdfContentType contentType,
|
||||
LoadedFile lf,
|
||||
AiWorkflowFileRequest fileReq,
|
||||
int remainingPages,
|
||||
int remainingCharacters)
|
||||
throws IOException {
|
||||
return switch (contentType) {
|
||||
case PAGE_TEXT, FULL_TEXT ->
|
||||
Optional.<PdfContentResult>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<Integer> requestedPages = fileReq != null ? fileReq.getPageNumbers() : null;
|
||||
List<Integer> pages =
|
||||
selectPages(lf.document().getNumberOfPages(), requestedPages, remainingPages);
|
||||
List<AiWorkflowTextSelection> extracted =
|
||||
extractPageText(lf.document(), pages, remainingCharacters);
|
||||
return extracted.isEmpty() ? null : buildExtractedFileText(lf.fileName(), extracted);
|
||||
}
|
||||
|
||||
private WorkflowArtifact buildArtifact(ArtifactKind kind, List<PdfContentResult> results) {
|
||||
return switch (kind) {
|
||||
case EXTRACTED_TEXT -> {
|
||||
ExtractedTextArtifact artifact = new ExtractedTextArtifact();
|
||||
artifact.setFiles(results.stream().map(ExtractedFileText.class::cast).toList());
|
||||
yield artifact;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private List<Integer> selectPages(
|
||||
int totalPages, List<Integer> requestedPageNumbers, int maxPages) {
|
||||
if (totalPages <= 0) {
|
||||
throw ExceptionUtils.createPdfNoPages();
|
||||
}
|
||||
|
||||
List<Integer> pages = new ArrayList<>();
|
||||
|
||||
if (requestedPageNumbers == null || requestedPageNumbers.isEmpty()) {
|
||||
for (int p = 1; p <= totalPages && pages.size() < maxPages; p++) {
|
||||
pages.add(p);
|
||||
}
|
||||
return pages;
|
||||
}
|
||||
|
||||
Set<Integer> 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<AiWorkflowTextSelection> extractPageText(
|
||||
PDDocument document, List<Integer> selectedPages, int maxCharacters)
|
||||
throws IOException {
|
||||
PDFTextStripper textStripper = new PDFTextStripper();
|
||||
List<AiWorkflowTextSelection> 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<AiWorkflowTextSelection> 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<AiWorkflowTextSelection> 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<ExtractedFileText> files = new ArrayList<>();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user