Pdf comment agent (#6196)

Co-authored-by: James Brunton <[email protected]>
This commit is contained in:
ConnorYoh
2026-05-01 10:19:38 +01:00
committed by GitHub
co-authored by James Brunton
parent 2dc5276e8b
commit 86774d556e
78 changed files with 5091 additions and 112 deletions
@@ -20,21 +20,30 @@ import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.proprietary.model.api.ai.Verdict;
import stirling.software.proprietary.service.AiToolInputValidator;
import stirling.software.proprietary.service.MathAuditorOrchestrator;
/**
* Public entry point for the Math Auditor Agent (mathAuditorAgent).
*
* <p>Accepts a PDF from the client, hands it to the {@link MathAuditorOrchestrator} which runs the
* multi-round Java-Python negotiation, and returns the Auditor's {@link Verdict}.
* multi-round Java-Python negotiation, and returns the Auditor's {@link Verdict} as JSON.
*
* <p>This endpoint is a pure specialist — it produces the structured finding and nothing more.
* Presentation (rendering as a chat answer, projecting to PDF comments, etc.) is the responsibility
* of the caller (e.g. the orchestrator's {@code delegate_pdf_question} or {@code
* delegate_pdf_review} meta-agents).
*
* <p>Lives under {@code /api/v1/ai/tools/} so it is dispatchable by the AI orchestrator via the
* standard {@code InternalApiClient} allowlist — no special-case plumbing needed.
*
* <p>The raw PDF never leaves Java. Python receives only structured text and CSV data.
*/
@Slf4j
@RestController
@RequestMapping("/api/v1/ai")
@RequestMapping("/api/v1/ai/tools")
@RequiredArgsConstructor
@Tag(name = "AI Engine", description = "AI-powered document analysis endpoints.")
@Tag(name = "AI Tools", description = "Dispatchable AI-backed tools.")
public class MathAuditorAgentController {
private final MathAuditorOrchestrator orchestrator;
@@ -49,11 +58,12 @@ public class MathAuditorAgentController {
The auditor checks:
- Table row and column totals (tally errors)
- Inline arithmetic expressions (e.g. "100 + 200 = 300")
- Cross-page figure consistency (same figure cited differently on different pages)
- Cross-page figure consistency
- Prose claims about percentages, growth rates, and comparisons
The PDF is processed entirely on the Java side; only extracted text and table data
are sent to the AI engine.
Returns a JSON Verdict describing every discrepancy found. How the Verdict is
presented to the end user (chat answer, PDF annotations, etc.) is up to the
caller.
Input: PDF Output: JSON Type: SISO
""")
@@ -68,11 +78,7 @@ public class MathAuditorAgentController {
@RequestParam(value = "tolerance", defaultValue = "0.01")
BigDecimal tolerance) {
String contentType = fileInput.getContentType();
if (contentType == null || !contentType.equals("application/pdf")) {
return ResponseEntity.badRequest().build();
}
AiToolInputValidator.validatePdfUpload(fileInput);
if (tolerance.compareTo(BigDecimal.ZERO) < 0) {
return ResponseEntity.badRequest().build();
}
@@ -89,9 +95,6 @@ public class MathAuditorAgentController {
} catch (IOException e) {
log.error("[math-auditor-agent] IO error during audit", e);
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).build();
} catch (Exception e) {
log.error("[math-auditor-agent] unexpected error during audit", e);
return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).build();
}
}
}
@@ -0,0 +1,120 @@
package stirling.software.proprietary.controller.api;
import java.io.IOException;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.proprietary.service.AiToolResponseHeaders;
import stirling.software.proprietary.service.PdfCommentAgentOrchestrator;
import stirling.software.proprietary.service.PdfCommentAgentOrchestrator.AnnotatedPdf;
import tools.jackson.databind.ObjectMapper;
import tools.jackson.databind.node.ObjectNode;
/**
* Public entry point for the PDF Comment Agent (pdfCommentAgent).
*
* <p>Accepts a PDF and a natural-language prompt, delegates to {@link PdfCommentAgentOrchestrator}
* which consults the Python engine and applies {@code PDAnnotationText} sticky-note annotations,
* then streams the annotated PDF back in the response body. This shape matches the rest of the
* Stirling tool endpoints ({@code /api/v1/misc/*}, {@code /api/v1/general/*}) and is what the AI
* workflow orchestrator expects when dispatching this tool as a plan step.
*
* <p>The raw PDF never leaves Java. Python only receives positioned text chunks.
*/
@Slf4j
@RestController
@RequestMapping("/api/v1/ai/tools")
@RequiredArgsConstructor
@Tag(name = "AI Tools", description = "Dispatchable AI-backed tools.")
public class PdfCommentAgentController {
private final PdfCommentAgentOrchestrator orchestrator;
private final ObjectMapper objectMapper;
@PostMapping(
value = "/pdf-comment-agent",
consumes = MediaType.MULTIPART_FORM_DATA_VALUE,
produces = MediaType.APPLICATION_PDF_VALUE)
@Operation(
summary = "Annotate a PDF with AI-generated sticky-note comments",
description =
"""
Runs the PDF Comment Agent against the supplied PDF. Java extracts positioned
text chunks from the document, ships them (with the user's prompt) to the
AI engine, then applies the returned comments as standard PDF Text
annotations (sticky notes) anchored to the relevant chunks.
The annotated PDF is streamed back in the response body with
Content-Type: application/pdf.
Input: PDF + prompt Output: PDF Type: SISO
""")
public ResponseEntity<Resource> pdfCommentAgent(
@Parameter(description = "The PDF document to annotate", required = true)
@RequestParam("fileInput")
MultipartFile fileInput,
@Parameter(
description =
"Natural-language instructions for the AI — what to comment on",
required = true)
@RequestParam("prompt")
String prompt)
throws IOException {
String safeName =
fileInput.getOriginalFilename() != null
? fileInput.getOriginalFilename().replaceAll("[\\r\\n]", "_")
: "<unnamed>";
log.info(
"[pdf-comment-agent] request file={} promptLen={}",
safeName,
prompt == null ? 0 : prompt.length());
// ResponseStatusException (validation errors) propagates to Spring's default handler;
// IOException is re-thrown to produce a 500. Other RuntimeExceptions likewise propagate.
AnnotatedPdf annotated = orchestrator.applyComments(fileInput, prompt);
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_PDF);
headers.setContentDispositionFormData("attachment", annotated.fileName());
headers.setContentLength(annotated.bytes().length);
headers.set(AiToolResponseHeaders.TOOL_REPORT, buildReportHeader(annotated));
return ResponseEntity.ok().headers(headers).body(new ByteArrayResource(annotated.bytes()));
}
/**
* Build the metadata JSON surfaced in {@link AiToolResponseHeaders#TOOL_REPORT} alongside the
* annotated PDF. Kept small (fits comfortably in a header): counts and the agent's rationale so
* a chat UI can show "Added 3 comments: <rationale>" alongside the downloaded file.
*/
private String buildReportHeader(AnnotatedPdf annotated) {
ObjectNode node = objectMapper.createObjectNode();
node.put("annotationsApplied", annotated.annotationsApplied());
node.put("instructionsReceived", annotated.instructionsReceived());
if (annotated.rationale() != null) {
node.put("rationale", annotated.rationale());
}
try {
return objectMapper.writeValueAsString(node);
} catch (Exception e) {
log.warn("Failed to serialise pdf-comment-agent report header: {}", e.getMessage());
return "{\"annotationsApplied\":" + annotated.annotationsApplied() + "}";
}
}
}
@@ -0,0 +1,35 @@
package stirling.software.proprietary.model.api.ai;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.Data;
/**
* Embedded plan optionally carried inside a question answer response. When present, the consumer
* (Java) runs the plan steps before delivering the answer; on the resume turn the engine returns
* the real answer using the captured tool reports.
*
* <p>Mirrors the engine's {@code EditPlanResponse} shape but is nested inside an answer rather than
* acting as the top-level outcome — matches the engine's {@code
* PdfQuestionAnswerResponse.edit_plan} field.
*/
@Data
@Schema(description = "Plan that must run before the answer is final")
public class AiWorkflowEditPlan {
@Schema(description = "Optional human-readable summary of the plan")
private String summary;
@Schema(description = "Optional rationale for the plan")
private String rationale;
@Schema(description = "Tool steps to execute before resuming")
private List<Map<String, Object>> steps = new ArrayList<>();
@Schema(description = "AI engine capability to resume with after running the steps")
private String resumeWith;
}
@@ -8,6 +8,8 @@ import io.swagger.v3.oas.annotations.media.Schema;
import lombok.Data;
import tools.jackson.databind.JsonNode;
@Data
@Schema(description = "Structured AI workflow result")
public class AiWorkflowResponse {
@@ -79,4 +81,19 @@ public class AiWorkflowResponse {
@Schema(description = "AI engine capability to resume with on the next turn")
private String resumeWith;
@Schema(
description =
"Optional structured report from the tool (e.g. math-auditor Verdict, PDF"
+ " comment-agent summary). Tools surface this either via a JSON response"
+ " body or via the X-Stirling-Tool-Report header. May be null for tools"
+ " that produce only a file.")
private JsonNode report;
@Schema(
description =
"Optional plan attached to an answer outcome. When non-null on outcome=ANSWER,"
+ " run the plan steps before delivering the answer; the resumed call"
+ " produces the real answer.")
private AiWorkflowEditPlan editPlan;
}
@@ -0,0 +1,14 @@
package stirling.software.proprietary.model.api.ai.comments;
import java.util.List;
/**
* Request body sent from Java to the Python PDF Comment Agent at {@code POST
* /api/v1/ai/pdf-comment-agent/generate}.
*
* @param sessionId Random UUID that uniquely identifies this generate call.
* @param userMessage The user's natural-language prompt (e.g. "flag any ambiguous dates").
* @param chunks Positioned text chunks extracted from the PDF that the model may comment on.
*/
public record PdfCommentEngineRequest(
String sessionId, String userMessage, List<TextChunk> chunks) {}
@@ -0,0 +1,14 @@
package stirling.software.proprietary.model.api.ai.comments;
import java.util.List;
/**
* Response body returned by the Python PDF Comment Agent at {@code POST
* /api/v1/ai/pdf-comment-agent/generate}.
*
* @param sessionId Echoes the session id from the request.
* @param comments The comments the agent wants to place on the document.
* @param rationale Short free-text explanation of the agent's choices.
*/
public record PdfCommentEngineResponse(
String sessionId, List<PdfCommentInstruction> comments, String rationale) {}
@@ -0,0 +1,12 @@
package stirling.software.proprietary.model.api.ai.comments;
/**
* A single comment instruction returned by the Python PDF Comment Agent.
*
* @param chunkId The {@link TextChunk#id()} the comment should anchor to.
* @param commentText The comment body (required, non-null).
* @param author Optional author/title for the popup. May be {@code null}.
* @param subject Optional subject line for the popup. May be {@code null}.
*/
public record PdfCommentInstruction(
String chunkId, String commentText, String author, String subject) {}
@@ -0,0 +1,18 @@
package stirling.software.proprietary.model.api.ai.comments;
/**
* One positioned text chunk extracted from a PDF, sent to the Python PDF Comment Agent so it can
* pick which chunks to annotate.
*
* <p>The bounding box is in PDF user-space coordinates (origin at the page's bottom-left).
*
* @param id Stable chunk id in the form {@code "p{pageIdx}-c{chunkIdx}"} (both 0-indexed).
* @param page 0-indexed page number.
* @param x Bottom-left x coordinate of the chunk bbox (PDF user-space).
* @param y Bottom-left y coordinate of the chunk bbox (PDF user-space).
* @param width Width of the chunk bbox.
* @param height Height of the chunk bbox.
* @param text The plain-text content of the chunk (truncated to a sane length).
*/
public record TextChunk(
String id, int page, float x, float y, float width, float height, String text) {}
@@ -5,8 +5,10 @@ import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpTimeoutException;
import java.time.Duration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import org.springframework.web.server.ResponseStatusException;
@@ -22,14 +24,21 @@ public class AiEngineClient {
private final ApplicationProperties applicationProperties;
private final HttpClient httpClient;
@Autowired
public AiEngineClient(ApplicationProperties applicationProperties) {
this.applicationProperties = applicationProperties;
this.httpClient =
this(
applicationProperties,
HttpClient.newBuilder()
.connectTimeout(
Duration.ofSeconds(
applicationProperties.getAiEngine().getTimeoutSeconds()))
.build();
.build());
}
/** Package-private constructor that accepts an HttpClient directly; intended for tests. */
AiEngineClient(ApplicationProperties applicationProperties, HttpClient httpClient) {
this.applicationProperties = applicationProperties;
this.httpClient = httpClient;
}
public String post(String path, String jsonBody) throws IOException {
@@ -86,6 +95,14 @@ public class AiEngineClient {
private HttpResponse<String> sendRequest(HttpRequest request) throws IOException {
try {
return httpClient.send(request, HttpResponse.BodyHandlers.ofString());
} catch (HttpTimeoutException e) {
throw new ResponseStatusException(HttpStatus.GATEWAY_TIMEOUT, "AI engine timed out", e);
} catch (IOException e) {
// Connection refused, DNS failure, socket reset, etc. — surface as
// SERVICE_UNAVAILABLE so every caller of this client sees a structured
// status rather than a raw 500 from an unhandled IOException.
throw new ResponseStatusException(
HttpStatus.SERVICE_UNAVAILABLE, "AI engine unreachable: " + e.getMessage(), e);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new ResponseStatusException(
@@ -0,0 +1,48 @@
package stirling.software.proprietary.service;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException;
/**
* Shared input-validation for AI-backed tool endpoints.
*
* <p>Spring's {@code spring.servlet.multipart.max-file-size} is tuned for the regular PDF tools (2
* GB) — far too permissive for AI tools where upload size translates directly into token budget,
* memory, and engine cost. Every AI tool should call {@link #validatePdfUpload} on its input before
* doing any work.
*/
public final class AiToolInputValidator {
/**
* Upper bound on PDF size accepted by any AI tool. Chosen so that a realistic document fits
* (contracts, research papers, books) while capping pathological uploads that would blow the
* engine's token budget or memory.
*/
public static final long MAX_INPUT_FILE_BYTES = 50L * 1024 * 1024;
private AiToolInputValidator() {}
/**
* Validate a PDF uploaded to an AI tool endpoint. Throws {@link ResponseStatusException} with
* an appropriate HTTP status on any failure.
*/
public static void validatePdfUpload(MultipartFile file) {
if (file == null || file.isEmpty()) {
throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "fileInput is required");
}
String contentType = file.getContentType();
if (contentType == null || !contentType.equals(MediaType.APPLICATION_PDF_VALUE)) {
throw new ResponseStatusException(
HttpStatus.BAD_REQUEST, "Only application/pdf uploads are supported");
}
if (file.getSize() > MAX_INPUT_FILE_BYTES) {
throw new ResponseStatusException(
HttpStatus.PAYLOAD_TOO_LARGE,
"PDF exceeds maximum size of "
+ (MAX_INPUT_FILE_BYTES / (1024 * 1024))
+ " MB for AI tools");
}
}
}
@@ -0,0 +1,21 @@
package stirling.software.proprietary.service;
/**
* Custom response headers the AI tools use when returning a file body with structured metadata.
*
* <p>Kept in one place because the value is referenced both server-side (tools that produce the
* header, orchestrator code that consumes it) and client-side (HTTP response handling in the
* frontend). Changing the string requires updating every reader, so centralising avoids the "must
* stay in sync" coupling.
*/
public final class AiToolResponseHeaders {
/**
* Header tools set to surface a structured metadata report alongside a file body. Value is a
* JSON object whose shape depends on the tool (e.g. {@code annotationsApplied}, {@code
* rationale} for pdf-comment-agent). Absent when the tool has no metadata to report.
*/
public static final String TOOL_REPORT = "X-Stirling-Tool-Report";
private AiToolResponseHeaders() {}
}
@@ -11,6 +11,7 @@ import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.MediaTypeFactory;
@@ -35,6 +36,7 @@ import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
import stirling.software.common.util.ZipExtractionUtils;
import stirling.software.proprietary.model.api.ai.AiConversationMessage;
import stirling.software.proprietary.model.api.ai.AiWorkflowEditPlan;
import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
@@ -47,6 +49,8 @@ import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile;
import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult;
import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact;
import tools.jackson.core.JacksonException;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.ObjectMapper;
@Slf4j
@@ -76,6 +80,16 @@ public class AiWorkflowService {
record Terminal(AiWorkflowResponse response) implements WorkflowState {}
}
/**
* Internal value-class for tool responses. {@code files} holds any result files (typically one;
* multiple for ZIP-response tools). {@code report} holds an optional structured metadata
* payload the tool chose to surface alongside (or instead of) a file.
*
* <p>Tools populate the report either by returning a JSON body (whole body → report) or by
* adding the {@link AiToolResponseHeaders#TOOL_REPORT} header alongside a file body.
*/
private record ToolResult(List<Resource> files, JsonNode report) {}
public AiWorkflowResponse orchestrate(AiWorkflowRequest request) throws IOException {
return orchestrate(request, NOOP_LISTENER);
}
@@ -117,9 +131,9 @@ public class AiWorkflowService {
return switch (response.getOutcome()) {
case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener);
case TOOL_CALL -> onToolCall(response, filesByName, listener);
case PLAN -> onPlan(response, filesByName, listener);
case ANSWER,
NOT_FOUND,
case PLAN -> onPlan(response, filesByName, request, listener);
case ANSWER -> onAnswer(response, filesByName, request, listener);
case NOT_FOUND,
NEED_CLARIFICATION,
CANNOT_DO,
DRAFT,
@@ -221,12 +235,13 @@ public class AiWorkflowService {
try {
List<Resource> inputFiles = toResources(filesByName);
listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1));
List<Resource> results = executeStep(endpointPath, parameters, inputFiles);
ToolResult result = executeStep(endpointPath, parameters, inputFiles);
return new WorkflowState.Terminal(
buildCompletedResponse(
response.getRationale(),
results,
new ArrayList<>(filesByName.keySet())));
result.files(),
new ArrayList<>(filesByName.keySet()),
result.report()));
} catch (Exception e) {
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
return new WorkflowState.Terminal(
@@ -234,12 +249,49 @@ public class AiWorkflowService {
}
}
@SuppressWarnings("unchecked")
private WorkflowState onPlan(
AiWorkflowResponse response,
Map<String, MultipartFile> filesByName,
WorkflowTurnRequest previousRequest,
ProgressListener listener) {
return runPlan(
response.getSteps(),
response.getResumeWith(),
response.getSummary(),
filesByName,
previousRequest,
listener);
}
private WorkflowState onAnswer(
AiWorkflowResponse response,
Map<String, MultipartFile> filesByName,
WorkflowTurnRequest previousRequest,
ProgressListener listener) {
AiWorkflowEditPlan plan = response.getEditPlan();
if (plan != null) {
// The engine wants us to run a side-quest before the answer is final.
// Run the embedded plan and resume the orchestrator with the captured
// report; the real answer arrives on the resume turn.
return runPlan(
plan.getSteps(),
plan.getResumeWith(),
plan.getSummary(),
filesByName,
previousRequest,
listener);
}
return new WorkflowState.Terminal(response);
}
@SuppressWarnings("unchecked")
private WorkflowState runPlan(
List<Map<String, Object>> steps,
String resumeWith,
String summary,
Map<String, MultipartFile> filesByName,
WorkflowTurnRequest previousRequest,
ProgressListener listener) {
List<Map<String, Object>> steps = response.getSteps();
if (steps == null || steps.isEmpty()) {
return new WorkflowState.Terminal(
cannotContinue("AI engine returned a plan with no steps."));
@@ -247,6 +299,9 @@ public class AiWorkflowService {
try {
List<Resource> currentFiles = toResources(filesByName);
// Propagate the *last* non-null report — the terminal step defines the output.
JsonNode lastReport = null;
String lastReportTool = null;
for (int i = 0; i < steps.size(); i++) {
Map<String, Object> step = steps.get(i);
@@ -263,14 +318,37 @@ public class AiWorkflowService {
listener.onProgress(
AiWorkflowProgressEvent.executingTool(endpointPath, i + 1, steps.size()));
currentFiles = executeStep(endpointPath, parameters, currentFiles);
ToolResult stepResult = executeStep(endpointPath, parameters, currentFiles);
currentFiles = stepResult.files();
if (stepResult.report() != null) {
lastReport = stepResult.report();
lastReportTool = endpointPath;
}
}
// Multi-turn: if the plan was emitted with resume_with set, the delegate wants
// Java to re-invoke the orchestrator with any captured report as an artifact.
if (resumeWith != null && !resumeWith.isBlank() && lastReport != null) {
WorkflowTurnRequest resumeRequest = new WorkflowTurnRequest();
resumeRequest.setUserMessage(previousRequest.getUserMessage());
resumeRequest.setFileNames(previousRequest.getFileNames());
resumeRequest.setConversationHistory(previousRequest.getConversationHistory());
resumeRequest.setArtifacts(new ArrayList<>(previousRequest.getArtifacts()));
resumeRequest
.getArtifacts()
.add(
new PdfContentExtractor.ToolReportArtifact(
lastReportTool, lastReport));
resumeRequest.setResumeWith(resumeWith);
return new WorkflowState.Pending(resumeRequest);
}
return new WorkflowState.Terminal(
buildCompletedResponse(
response.getSummary(),
summary,
currentFiles,
new ArrayList<>(filesByName.keySet())));
new ArrayList<>(filesByName.keySet()),
lastReport));
} catch (Exception e) {
log.error("Failed to execute plan: {}", e.getMessage(), e);
return new WorkflowState.Terminal(
@@ -282,27 +360,46 @@ public class AiWorkflowService {
* Execute a single tool step. If the endpoint accepts multiple files, all files are sent in one
* call. Otherwise, the endpoint is called once per file. ZIP responses are unpacked so each
* inner file is treated as its own result (e.g. split outputs a ZIP of pages).
*
* <p>A structured {@code report} may be returned alongside (or instead of) files — see {@link
* ToolResult}. For per-file dispatch (single-input endpoints called once per input), the first
* non-null report wins.
*/
private List<Resource> executeStep(
private ToolResult executeStep(
String endpointPath, Map<String, Object> parameters, List<Resource> inputFiles)
throws IOException {
List<Resource> results = new ArrayList<>();
List<Resource> files = new ArrayList<>();
JsonNode report = null;
if (toolMetadataService.isMultiInput(endpointPath)) {
results.addAll(callEndpoint(endpointPath, parameters, inputFiles));
ToolResult r = callEndpoint(endpointPath, parameters, inputFiles);
files.addAll(r.files());
report = r.report();
} else {
for (Resource file : inputFiles) {
results.addAll(callEndpoint(endpointPath, parameters, List.of(file)));
ToolResult r = callEndpoint(endpointPath, parameters, List.of(file));
files.addAll(r.files());
if (report == null) {
report = r.report();
}
}
}
return results;
return new ToolResult(files, report);
}
/**
* Call an endpoint and return the response body. Endpoints that are declared as ZIP-returning
* in the API spec (multi-output, or {@code Output:ZIP-*}) are unpacked into their individual
* entries so callers always see a flat list of result files.
* Call an endpoint and return its result files and optional report.
*
* <ul>
* <li>JSON body (Content-Type: application/json) → the entire body is the report, no files
* are returned.
* <li>File body (PDF etc.) → the file is returned; if an {@link
* AiToolResponseHeaders#TOOL_REPORT} header is present, its (minified JSON) value is
* parsed as the report.
* <li>ZIP responses declared by the tool metadata service are unpacked so callers always see
* a flat list of result files.
* </ul>
*/
private List<Resource> callEndpoint(
private ToolResult callEndpoint(
String endpointPath, Map<String, Object> parameters, List<Resource> files)
throws IOException {
MultiValueMap<String, Object> body = new LinkedMultiValueMap<>();
@@ -324,10 +421,43 @@ public class AiWorkflowService {
"Tool returned HTTP " + response.getStatusCode() + " for " + endpointPath);
}
Resource resource = response.getBody();
HttpHeaders headers = response.getHeaders();
MediaType contentType = headers.getContentType();
// JSON-only response — the whole body is the structured report, no result file.
if (contentType != null && MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
try (java.io.InputStream is = resource.getInputStream()) {
JsonNode report = objectMapper.readTree(is);
return new ToolResult(List.of(), report);
}
}
JsonNode report = parseReportHeader(headers, endpointPath);
if (toolMetadataService.shouldUnpackZipResponse(endpointPath)) {
return ZipExtractionUtils.extractZip(resource, tempFileManager);
return new ToolResult(ZipExtractionUtils.extractZip(resource, tempFileManager), report);
}
return new ToolResult(List.of(resource), report);
}
/**
* Parse the optional {@link AiToolResponseHeaders#TOOL_REPORT} header into a {@link JsonNode},
* or return null.
*/
private JsonNode parseReportHeader(HttpHeaders headers, String endpointPath) {
String raw = headers.getFirst(AiToolResponseHeaders.TOOL_REPORT);
if (raw == null || raw.isBlank()) {
return null;
}
try {
return objectMapper.readTree(raw);
} catch (JacksonException e) {
log.warn(
"Ignoring malformed {} header from {}: {}",
AiToolResponseHeaders.TOOL_REPORT,
endpointPath,
e.getMessage());
return null;
}
return List.of(resource);
}
private List<Resource> toResources(Map<String, MultipartFile> filesByName) throws IOException {
@@ -348,7 +478,10 @@ public class AiWorkflowService {
}
private AiWorkflowResponse buildCompletedResponse(
String summary, List<Resource> resultFiles, List<String> inputFileNames)
String summary,
List<Resource> resultFiles,
List<String> inputFileNames,
JsonNode report)
throws IOException {
// Store every output file individually so each gets its own Stirling file ID and the
// frontend can add them as independent variants without going through a zip.
@@ -386,6 +519,7 @@ public class AiWorkflowService {
completed.setOutcome(AiWorkflowOutcome.COMPLETED);
completed.setSummary(summary);
completed.setResultFiles(descriptors);
completed.setReport(report);
// Mirror the first file into the legacy single-file fields so existing clients still work.
if (!descriptors.isEmpty()) {
AiWorkflowResultFile first = descriptors.getFirst();
@@ -0,0 +1,249 @@
package stirling.software.proprietary.service;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.api.comments.AnnotationLocation;
import stirling.software.common.model.api.comments.StickyNoteSpec;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.PdfAnnotationService;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineRequest;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineResponse;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentInstruction;
import stirling.software.proprietary.model.api.ai.comments.TextChunk;
import tools.jackson.databind.ObjectMapper;
/**
* Composed AI tool for PDF comment generation.
*
* <p>Runs the full flow:
*
* <ol>
* <li>Validate inputs (PDF, non-empty prompt within length limit).
* <li>Extract positioned text chunks from the PDF.
* <li>POST the chunks + prompt to the Python agent at {@code
* /api/v1/ai/pdf-comment-agent/generate}.
* <li>Resolve each returned chunk-id reference to an absolute {@link StickyNoteSpec}.
* <li>Hand the specs to {@link PdfAnnotationService} for deterministic placement.
* <li>Save the annotated PDF, return the bytes + filename.
* </ol>
*
* <p>Annotation primitives live in {@link PdfAnnotationService} (shared with {@code
* /api/v1/misc/add-comments}). This class owns only the AI-specific bits: chunk extraction and
* engine round-trip.
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class PdfCommentAgentOrchestrator {
private static final String GENERATE_PATH = "/api/v1/ai/pdf-comment-agent/generate";
private static final int MAX_PROMPT_LEN = 4000;
/** Width/height of the sticky-note icon placed on the page, in PDF user-space units. */
private static final float ANNOTATION_SIZE = 20f;
/** Filename used when the uploaded PDF has no usable original filename. */
private static final String FALLBACK_OUTPUT_NAME = "document-commented.pdf";
/**
* Small value record returned to the controller: the annotated PDF bytes, the suggested
* download filename (used in the {@code Content-Disposition} header), and metadata the
* controller emits in the {@code X-Stirling-Tool-Report} header so callers (frontend,
* orchestrator) can surface a chat-style summary alongside the file.
*/
public record AnnotatedPdf(
byte[] bytes,
String fileName,
int annotationsApplied,
int instructionsReceived,
String rationale) {}
private final AiEngineClient aiEngineClient;
private final PdfTextChunkExtractor pdfTextChunkExtractor;
private final CustomPDFDocumentFactory pdfDocumentFactory;
private final ObjectMapper objectMapper;
private final PdfAnnotationService pdfAnnotationService;
/**
* Run the full PDF comment generation flow.
*
* @param pdfFile the uploaded PDF
* @param prompt the user's natural-language instructions
* @return the annotated PDF bytes and suggested filename
*/
public AnnotatedPdf applyComments(MultipartFile pdfFile, String prompt) throws IOException {
AiToolInputValidator.validatePdfUpload(pdfFile);
String trimmedPrompt = prompt == null ? "" : prompt.trim();
if (trimmedPrompt.isEmpty()) {
throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Prompt is required");
}
if (trimmedPrompt.length() > MAX_PROMPT_LEN) {
throw new ResponseStatusException(
HttpStatus.BAD_REQUEST,
"Prompt exceeds maximum length of " + MAX_PROMPT_LEN + " characters");
}
String sessionId = UUID.randomUUID().toString();
log.info(
"[pdf-comment-agent] session={} file={} promptLen={}",
sessionId,
safeName(pdfFile.getOriginalFilename()),
trimmedPrompt.length());
try (PDDocument document = pdfDocumentFactory.load(pdfFile)) {
List<TextChunk> chunks = pdfTextChunkExtractor.extract(document);
if (chunks.isEmpty()) {
throw new ResponseStatusException(
HttpStatus.BAD_REQUEST, "PDF has no extractable text");
}
log.info(
"[pdf-comment-agent] session={} extracted {} chunks across {} pages",
sessionId,
chunks.size(),
document.getNumberOfPages());
PdfCommentEngineResponse engineResponse =
requestComments(sessionId, trimmedPrompt, chunks);
List<PdfCommentInstruction> instructions =
engineResponse.comments() == null ? List.of() : engineResponse.comments();
// Resolve chunk-id-referenced comments to absolute sticky-note specs, then delegate
// placement to the shared service (same primitive /api/v1/misc/add-comments uses).
List<StickyNoteSpec> specs = resolveSpecs(instructions, chunks, sessionId);
int applied = pdfAnnotationService.addStickyNotes(document, specs);
log.info(
"[pdf-comment-agent] session={} placed {}/{} sticky notes",
sessionId,
applied,
instructions.size());
byte[] annotatedBytes;
try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
document.save(baos);
annotatedBytes = baos.toByteArray();
}
String outputName = buildOutputFileName(pdfFile.getOriginalFilename());
log.info(
"[pdf-comment-agent] session={} done fileName={} bytes={}",
sessionId,
outputName,
annotatedBytes.length);
return new AnnotatedPdf(
annotatedBytes,
outputName,
applied,
instructions.size(),
engineResponse.rationale());
}
}
// -----------------------------------------------------------------------
// Engine round-trip
// -----------------------------------------------------------------------
private PdfCommentEngineResponse requestComments(
String sessionId, String prompt, List<TextChunk> chunks) throws IOException {
PdfCommentEngineRequest engineRequest =
new PdfCommentEngineRequest(sessionId, prompt, chunks);
String requestBody = objectMapper.writeValueAsString(engineRequest);
String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody);
PdfCommentEngineResponse engineResponse =
objectMapper.readValue(responseBody, PdfCommentEngineResponse.class);
List<PdfCommentInstruction> instructions =
engineResponse.comments() == null ? List.of() : engineResponse.comments();
log.info(
"[pdf-comment-agent] session={} engine returned {} comments: {}",
sessionId,
instructions.size(),
engineResponse.rationale());
return engineResponse;
}
// -----------------------------------------------------------------------
// Chunk-id → StickyNoteSpec resolution
// -----------------------------------------------------------------------
/**
* Translate each engine-returned {@link PdfCommentInstruction} (chunk-id-referenced) into an
* absolute-positioned {@link StickyNoteSpec}. Unknown or malformed ids are logged and dropped.
*/
private List<StickyNoteSpec> resolveSpecs(
List<PdfCommentInstruction> instructions, List<TextChunk> chunks, String sessionId) {
if (instructions.isEmpty()) {
return List.of();
}
Map<String, TextChunk> chunksById = new HashMap<>();
for (TextChunk chunk : chunks) {
chunksById.put(chunk.id(), chunk);
}
List<StickyNoteSpec> specs = new ArrayList<>(instructions.size());
for (PdfCommentInstruction inst : instructions) {
if (inst == null || inst.chunkId() == null || inst.commentText() == null) {
log.warn(
"[pdf-comment-agent] session={} skipping malformed instruction: {}",
sessionId,
inst);
continue;
}
TextChunk chunk = chunksById.get(inst.chunkId());
if (chunk == null) {
log.warn(
"[pdf-comment-agent] session={} unknown chunkId={} - skipping",
sessionId,
inst.chunkId());
continue;
}
// Anchor the sticky-note icon at the top-left of the chunk's bbox.
float iconX = chunk.x();
float iconY = chunk.y() + chunk.height() - ANNOTATION_SIZE;
AnnotationLocation loc =
new AnnotationLocation(
chunk.page(), iconX, iconY, ANNOTATION_SIZE, ANNOTATION_SIZE);
specs.add(new StickyNoteSpec(loc, inst.commentText(), inst.author(), inst.subject()));
}
return specs;
}
// -----------------------------------------------------------------------
// Helpers
// -----------------------------------------------------------------------
private static String buildOutputFileName(String originalFilename) {
String safe = safeName(originalFilename);
if (safe == null || safe.isBlank() || "<unnamed>".equals(safe)) {
return FALLBACK_OUTPUT_NAME;
}
String base = FilenameUtils.getBaseName(safe);
if (base == null || base.isBlank()) {
base = "document";
}
return base + "-commented.pdf";
}
private static String safeName(String originalFilename) {
return originalFilename != null
? originalFilename.replaceAll("[\\r\\n]", "_")
: "<unnamed>";
}
}
@@ -210,6 +210,12 @@ public class PdfContentExtractor {
artifact.setFiles(results.stream().map(ExtractedFileText.class::cast).toList());
yield artifact;
}
case TOOL_REPORT ->
// TOOL_REPORT artifacts don't come from PDF content extraction — they're
// built by AiWorkflowService from tool-response metadata. Never reached
// from this code path; presence in the enum is to satisfy the switch.
throw new IllegalArgumentException(
"TOOL_REPORT artifacts are not produced by PdfContentExtractor");
};
}
@@ -320,7 +326,8 @@ public class PdfContentExtractor {
* Values MUST match {@code ArtifactKind} in {@code engine/src/stirling/contracts/common.py}.
*/
enum ArtifactKind {
EXTRACTED_TEXT("extracted_text");
EXTRACTED_TEXT("extracted_text"),
TOOL_REPORT("tool_report");
private final String value;
@@ -364,4 +371,23 @@ public class PdfContentExtractor {
private final ArtifactKind kind = ArtifactKind.EXTRACTED_TEXT;
private List<ExtractedFileText> files = new ArrayList<>();
}
/**
* Carries a structured report produced by a specialist tool back to the orchestrator on a
* resume turn. Shape matches {@code engine/src/stirling/contracts/common.py ToolReportArtifact}
* — {@code sourceTool} must be a valid endpoint path string.
*/
@Data
static final class ToolReportArtifact implements WorkflowArtifact {
private final ArtifactKind kind = ArtifactKind.TOOL_REPORT;
private String sourceTool;
private tools.jackson.databind.JsonNode report;
ToolReportArtifact() {}
ToolReportArtifact(String sourceTool, tools.jackson.databind.JsonNode report) {
this.sourceTool = sourceTool;
this.report = report;
}
}
}
@@ -0,0 +1,176 @@
package stirling.software.proprietary.service;
import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.List;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.common.PDRectangle;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.pdfbox.text.TextPosition;
import org.springframework.stereotype.Service;
import lombok.extern.slf4j.Slf4j;
import stirling.software.proprietary.model.api.ai.comments.TextChunk;
/**
* Extracts positioned, line-level text chunks from a PDF so the PDF Comment Agent can decide where
* to anchor annotations. One chunk per text line, with the bounding box converted to PDF user-space
* coordinates (origin = bottom-left).
*/
@Slf4j
@Service
public class PdfTextChunkExtractor {
/** Hard cap on total chunks emitted per document. */
private static final int MAX_CHUNKS_PER_DOC = 2000;
/** Truncate each chunk's text to at most this many characters. */
private static final int MAX_CHUNK_TEXT_LENGTH = 500;
/**
* Extract line-level text chunks with bounding boxes from the given document.
*
* <p>Each chunk's coordinates are in PDF user-space (origin = bottom-left of the page). Chunks
* that are whitespace-only after trimming are skipped.
*
* @param document the open PDF
* @return positioned chunks (never {@code null})
* @throws IOException on PDF parse errors
*/
public List<TextChunk> extract(PDDocument document) throws IOException {
List<TextChunk> chunks = new ArrayList<>();
ChunkStripper stripper = new ChunkStripper(document, chunks);
stripper.setSortByPosition(true);
stripper.getText(document);
return chunks;
}
/**
* PDFTextStripper subclass that emits one {@link TextChunk} per call to {@code writeString}.
* PDFBox invokes {@code writeString} once per visual line when {@link
* PDFTextStripper#setSortByPosition(boolean)} is true, which gives us exactly the granularity
* we want.
*/
private static final class ChunkStripper extends PDFTextStripper {
private final PDDocument document;
private final List<TextChunk> chunks;
private int currentPageIdx = 0; // 0-indexed, tracked via startPage
private int chunkIdxOnPage = 0;
private boolean capWarningLogged = false;
ChunkStripper(PDDocument document, List<TextChunk> chunks) throws IOException {
super();
this.document = document;
this.chunks = chunks;
}
@Override
protected void startPage(org.apache.pdfbox.pdmodel.PDPage page) throws IOException {
super.startPage(page);
// getCurrentPageNo() is 1-based; convert to 0-based.
currentPageIdx = getCurrentPageNo() - 1;
chunkIdxOnPage = 0;
}
@Override
protected void writeString(String text, List<TextPosition> textPositions)
throws IOException {
if (chunks.size() >= MAX_CHUNKS_PER_DOC) {
if (!capWarningLogged) {
log.warn(
"[pdf-comment-agent] chunk cap of {} reached; remaining text will not"
+ " be extracted",
MAX_CHUNKS_PER_DOC);
capWarningLogged = true;
}
return;
}
if (textPositions == null || textPositions.isEmpty()) {
return;
}
String trimmed = text == null ? "" : text.trim();
if (trimmed.isEmpty()) {
return;
}
// Compute the bounding box from the min/max of TextPosition adjusted coordinates.
// getXDirAdj / getYDirAdj / getHeightDir / getWidthDirAdj already account for the
// page's rotation so we can treat them as axis-aligned in the page's display frame.
float minX = Float.POSITIVE_INFINITY;
float maxRight = Float.NEGATIVE_INFINITY;
float minYTopDown = Float.POSITIVE_INFINITY; // smallest y in top-down coords
float maxHeight = 0f;
for (TextPosition pos : textPositions) {
float x = pos.getXDirAdj();
float right = x + pos.getWidthDirAdj();
float yTop = pos.getYDirAdj();
float h = pos.getHeightDir();
if (h <= 0f) {
h = pos.getFontSizeInPt();
}
if (x < minX) minX = x;
if (right > maxRight) maxRight = right;
if (yTop < minYTopDown) minYTopDown = yTop;
if (h > maxHeight) maxHeight = h;
}
if (maxHeight <= 0f) {
// Fallback if everything was zero — small but non-zero so the rect is valid.
maxHeight = 10f;
}
float width = maxRight - minX;
if (width <= 0f) {
return;
}
// Convert y to PDF user-space (origin at bottom-left of the page).
// getYDirAdj reports the top of each glyph, measured from the top of the page.
PDRectangle mediaBox = document.getPage(currentPageIdx).getMediaBox();
float pageHeight = mediaBox.getHeight();
float bottomY = pageHeight - minYTopDown - maxHeight;
String id = "p" + currentPageIdx + "-c" + chunkIdxOnPage;
chunkIdxOnPage++;
String storedText = trimmed;
if (storedText.length() > MAX_CHUNK_TEXT_LENGTH) {
storedText = storedText.substring(0, MAX_CHUNK_TEXT_LENGTH);
}
chunks.add(
new TextChunk(id, currentPageIdx, minX, bottomY, width, maxHeight, storedText));
}
@Override
protected void writeCharacters(TextPosition text) {
// no-op: we only emit chunks at writeString granularity
}
@Override
public String getText(PDDocument doc) throws IOException {
// We don't actually need the concatenated text — just the side-effects. Return early
// to avoid building a (potentially massive) StringBuilder of the whole document.
try (Writer discard =
new Writer() {
@Override
public void write(char[] cbuf, int off, int len) {
// discard
}
@Override
public void flush() {}
@Override
public void close() {}
}) {
writeText(doc, discard);
}
return "";
}
}
}