mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-16 03:20:46 +02:00
Flesh out RAG system (#6197)
# Description of Changes
Flesh out the RAG system and connect it to the PDF Question Agent so it
can respond to questions about PDFs of an extremely large size.
I'd expect lots more work will need to be done to finish off the RAG
system to really be what we need, but this should be a reasonable start
which will let us connect it to tools and have the ingestion mostly
handled automatically. I'm leaving file deletion and proper file ID
management to be done in a future PR. We also need to consider whether
all tools should retrieve content exclusively via RAG, or whether it's
beneficial to have tools sometimes fetch the direct content and other
times fetch it from RAG.
A diagram of the expected interaction is as follows:
```mermaid
sequenceDiagram
autonumber
actor U as User
participant FE as Frontend<br/>(ChatPanel)
participant J as Java<br/>(AiWorkflowService)
participant O as Engine:<br/>OrchestratorAgent
participant QA as Engine:<br/>PdfQuestionAgent
participant RAG as Engine:<br/>RagService + SqliteVecStore
participant V as VoyageAI<br/>(embeddings)
participant L as LLM<br/>(Claude / etc.)
U->>FE: types "Summarise this PDF"<br/>(PDF already uploaded)
FE->>J: POST /api/v1/ai/orchestrate/stream<br/>multipart: fileInputs[], userMessage
Note over J: ByteHashFileIdStrategy<br/>id = sha256(bytes)[:16]
J->>O: POST /api/v1/orchestrator<br/>{ files:[{id,name}], userMessage }
O->>L: route via fast model
L-->>O: delegate_pdf_question
O->>QA: PdfQuestionRequest
loop for each file
QA->>RAG: has_collection(file.id)
RAG-->>QA: false
end
QA-->>O: NeedIngestResponse(files_to_ingest)
O-->>J: { outcome:"need_ingest", filesToIngest:[...] }
Note over J: onNeedIngest
loop per file
J->>J: PDFBox: extract page text
J->>O: POST /api/v1/rag/documents<br/>(long-running timeout)
O->>RAG: chunk + stage documents
O->>V: embed_documents (batches of 256)
V-->>O: embeddings
O->>RAG: add_documents
O-->>J: { chunks_indexed: N }
end
Note over J: retry with resumeWith=pdf_question
J->>O: POST /api/v1/orchestrator
Note over O: fast-path to PdfQuestionAgent
O->>QA: PdfQuestionRequest
Note over QA: build RagCapability<br/>pinned to file IDs
QA->>L: run(prompt) with search_knowledge tool
loop up to max_searches
L->>QA: search_knowledge(query)
QA->>V: embed_query
V-->>QA: query vector
QA->>RAG: search(vector, collections=[file.id])
RAG-->>QA: top-k chunks
QA-->>L: formatted chunks
end
Note over QA: once budget spent,<br/>prepare() hides the tool
L-->>QA: PdfQuestionAnswerResponse
QA-->>O: answer
O-->>J: { outcome:"answer", answer, evidence }
J-->>FE: SSE "result"
FE->>U: assistant bubble
```
This commit is contained in:
James Brunton
@@ -0,0 +1,27 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import io.swagger.v3.oas.annotations.media.Schema;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
/**
|
||||
* A file supplied to the AI engine, identified by a stable opaque id plus a display name.
|
||||
*
|
||||
* <p>Values MUST match {@code AiFile} in {@code engine/src/stirling/contracts/common.py}.
|
||||
*/
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
@Schema(description = "File reference sent to the AI engine")
|
||||
public class AiFile {
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"Opaque, stable identifier. Owned by Java; used as the RAG collection key.")
|
||||
private String id;
|
||||
|
||||
@Schema(description = "Original filename, used by agents in user-facing prompts and responses.")
|
||||
private String name;
|
||||
}
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
/**
|
||||
* Body for {@code POST /api/v1/rag/documents} on the AI engine. Sent by Java when the engine
|
||||
* reports {@code need_ingest} and the requested document's extracted content must be stored before
|
||||
* the workflow can continue.
|
||||
*/
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
public class AiRagIngestRequest {
|
||||
|
||||
private String documentId;
|
||||
|
||||
private String source;
|
||||
|
||||
private List<AiRagPageText> pageText;
|
||||
}
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
/** A single page of extracted text for RAG ingest requests. */
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
public class AiRagPageText {
|
||||
|
||||
private int pageNumber;
|
||||
|
||||
private String text;
|
||||
}
|
||||
-35
@@ -1,35 +0,0 @@
|
||||
package stirling.software.proprietary.model.api.ai;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import io.swagger.v3.oas.annotations.media.Schema;
|
||||
|
||||
import lombok.Data;
|
||||
|
||||
/**
|
||||
* Embedded plan optionally carried inside a question answer response. When present, the consumer
|
||||
* (Java) runs the plan steps before delivering the answer; on the resume turn the engine returns
|
||||
* the real answer using the captured tool reports.
|
||||
*
|
||||
* <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;
|
||||
}
|
||||
+2
-2
@@ -11,8 +11,8 @@ import lombok.Data;
|
||||
@Schema(description = "Per-file content extraction request from the AI engine")
|
||||
public class AiWorkflowFileRequest {
|
||||
|
||||
@Schema(description = "Original filename of the requested file", example = "contract.pdf")
|
||||
private String fileName;
|
||||
@Schema(description = "The file the engine wants content extracted for")
|
||||
private AiFile file;
|
||||
|
||||
@Schema(description = "Specific 1-based page numbers to extract from this file")
|
||||
private List<Integer> pageNumbers = new ArrayList<>();
|
||||
|
||||
+1
@@ -12,6 +12,7 @@ public enum AiWorkflowOutcome {
|
||||
ANSWER("answer"),
|
||||
NOT_FOUND("not_found"),
|
||||
NEED_CONTENT("need_content"),
|
||||
NEED_INGEST("need_ingest"),
|
||||
PLAN("plan"),
|
||||
NEED_CLARIFICATION("need_clarification"),
|
||||
CANNOT_DO("cannot_do"),
|
||||
|
||||
+6
-7
@@ -73,6 +73,12 @@ public class AiWorkflowResponse {
|
||||
@Schema(description = "Per-file text extraction requests from the AI engine")
|
||||
private List<AiWorkflowFileRequest> files = new ArrayList<>();
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"Files the AI engine requires to be ingested into RAG before it can continue"
|
||||
+ " the workflow. Populated on need_ingest outcomes.")
|
||||
private List<AiFile> filesToIngest = new ArrayList<>();
|
||||
|
||||
@Schema(description = "Maximum number of pages the AI engine wants text extracted from")
|
||||
private Integer maxPages;
|
||||
|
||||
@@ -89,11 +95,4 @@ public class AiWorkflowResponse {
|
||||
+ " body or via the X-Stirling-Tool-Report header. May be null for tools"
|
||||
+ " that produce only a file.")
|
||||
private JsonNode report;
|
||||
|
||||
@Schema(
|
||||
description =
|
||||
"Optional plan attached to an answer outcome. When non-null on outcome=ANSWER,"
|
||||
+ " run the plan steps before delivering the answer; the resumed call"
|
||||
+ " produces the real answer.")
|
||||
private AiWorkflowEditPlan editPlan;
|
||||
}
|
||||
|
||||
+18
-2
@@ -43,20 +43,36 @@ public class AiEngineClient {
|
||||
|
||||
public String post(String path, String jsonBody) throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
return postWithTimeout(path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds()));
|
||||
}
|
||||
|
||||
/**
|
||||
* POST with an explicit per-call timeout, for heavy operations (e.g. RAG ingestion of a large
|
||||
* document) that legitimately take longer than the default timeout.
|
||||
*/
|
||||
public String postLongRunning(String path, String jsonBody) throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
return postWithTimeout(
|
||||
path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds()));
|
||||
}
|
||||
|
||||
private String postWithTimeout(String path, String jsonBody, Duration timeout)
|
||||
throws IOException {
|
||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||
if (!config.isEnabled()) {
|
||||
throw new ResponseStatusException(
|
||||
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
|
||||
}
|
||||
|
||||
String url = config.getUrl().stripTrailing() + path;
|
||||
log.debug("Proxying AI engine request to {}", url);
|
||||
log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds());
|
||||
|
||||
HttpRequest request =
|
||||
HttpRequest.newBuilder()
|
||||
.uri(URI.create(url))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
|
||||
.timeout(timeout)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||
.build();
|
||||
|
||||
|
||||
+123
-57
@@ -36,7 +36,9 @@ import stirling.software.common.util.TempFile;
|
||||
import stirling.software.common.util.TempFileManager;
|
||||
import stirling.software.common.util.ZipExtractionUtils;
|
||||
import stirling.software.proprietary.model.api.ai.AiConversationMessage;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowEditPlan;
|
||||
import stirling.software.proprietary.model.api.ai.AiFile;
|
||||
import stirling.software.proprietary.model.api.ai.AiRagIngestRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiRagPageText;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileInput;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
|
||||
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
|
||||
@@ -58,6 +60,8 @@ import tools.jackson.databind.ObjectMapper;
|
||||
@RequiredArgsConstructor
|
||||
public class AiWorkflowService {
|
||||
|
||||
private static final String RAG_DOCUMENTS_ENDPOINT = "/api/v1/rag/documents";
|
||||
|
||||
private final CustomPDFDocumentFactory pdfDocumentFactory;
|
||||
private final AiEngineClient aiEngineClient;
|
||||
private final PdfContentExtractor pdfContentExtractor;
|
||||
@@ -66,6 +70,7 @@ public class AiWorkflowService {
|
||||
private final FileStorage fileStorage;
|
||||
private final ToolMetadataService toolMetadataService;
|
||||
private final TempFileManager tempFileManager;
|
||||
private final FileIdStrategy fileIdStrategy;
|
||||
|
||||
@FunctionalInterface
|
||||
public interface ProgressListener {
|
||||
@@ -98,15 +103,24 @@ public class AiWorkflowService {
|
||||
throws IOException {
|
||||
validateRequest(request);
|
||||
|
||||
Map<String, MultipartFile> filesByName = new LinkedHashMap<>();
|
||||
// Key by opaque file id, not filename. Filenames aren't guaranteed unique across an
|
||||
// upload (users can rotate the same 'scan.pdf' twice), and the engine identifies files
|
||||
// by id in every response shape that asks Java to look a file up again.
|
||||
Map<String, MultipartFile> filesById = new LinkedHashMap<>();
|
||||
List<AiFile> files = new ArrayList<>();
|
||||
for (AiWorkflowFileInput fileInput : request.getFileInputs()) {
|
||||
filesByName.put(
|
||||
fileInput.getFileInput().getOriginalFilename(), fileInput.getFileInput());
|
||||
MultipartFile multipartFile = fileInput.getFileInput();
|
||||
AiFile aiFile =
|
||||
new AiFile(
|
||||
fileIdStrategy.idFor(multipartFile),
|
||||
multipartFile.getOriginalFilename());
|
||||
filesById.put(aiFile.getId(), multipartFile);
|
||||
files.add(aiFile);
|
||||
}
|
||||
|
||||
WorkflowTurnRequest initialRequest = new WorkflowTurnRequest();
|
||||
initialRequest.setUserMessage(request.getUserMessage().trim());
|
||||
initialRequest.setFileNames(new ArrayList<>(filesByName.keySet()));
|
||||
initialRequest.setFiles(files);
|
||||
initialRequest.setConversationHistory(
|
||||
request.getConversationHistory() == null
|
||||
? new ArrayList<>()
|
||||
@@ -116,23 +130,24 @@ public class AiWorkflowService {
|
||||
|
||||
WorkflowState state = new WorkflowState.Pending(initialRequest);
|
||||
while (state instanceof WorkflowState.Pending pending) {
|
||||
state = advance(pending.request(), filesByName, listener);
|
||||
state = advance(pending.request(), filesById, listener);
|
||||
}
|
||||
return ((WorkflowState.Terminal) state).response();
|
||||
}
|
||||
|
||||
private WorkflowState advance(
|
||||
WorkflowTurnRequest request,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
ProgressListener listener)
|
||||
throws IOException {
|
||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
|
||||
AiWorkflowResponse response = invokeOrchestrator(request);
|
||||
return switch (response.getOutcome()) {
|
||||
case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener);
|
||||
case TOOL_CALL -> onToolCall(response, filesByName, listener);
|
||||
case PLAN -> onPlan(response, filesByName, request, listener);
|
||||
case ANSWER -> onAnswer(response, filesByName, request, listener);
|
||||
case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
|
||||
case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
|
||||
case TOOL_CALL -> onToolCall(response, filesById, listener);
|
||||
case PLAN -> onPlan(response, filesById, request, listener);
|
||||
case ANSWER -> onAnswer(response, filesById, request, listener);
|
||||
case NOT_FOUND,
|
||||
NEED_CLARIFICATION,
|
||||
CANNOT_DO,
|
||||
@@ -146,7 +161,7 @@ public class AiWorkflowService {
|
||||
|
||||
private WorkflowState onNeedContent(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
WorkflowTurnRequest request,
|
||||
ProgressListener listener)
|
||||
throws IOException {
|
||||
@@ -157,43 +172,42 @@ public class AiWorkflowService {
|
||||
|
||||
List<AiWorkflowFileRequest> requestedFiles = response.getFiles();
|
||||
|
||||
// Validate requested file names before loading anything
|
||||
// Validate requested file ids before loading anything
|
||||
if (requestedFiles != null && !requestedFiles.isEmpty()) {
|
||||
for (AiWorkflowFileRequest fileReq : requestedFiles) {
|
||||
if (!filesByName.containsKey(fileReq.getFileName())) {
|
||||
AiFile file = fileReq.getFile();
|
||||
if (file == null || !filesById.containsKey(file.getId())) {
|
||||
String display = file == null ? "<missing file>" : file.getName();
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue(
|
||||
"AI engine requested unknown file: " + fileReq.getFileName()));
|
||||
cannotContinue("AI engine requested unknown file: " + display));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
List<String> fileNamesToLoad =
|
||||
List<AiFile> filesToLoad =
|
||||
(requestedFiles == null || requestedFiles.isEmpty())
|
||||
? new ArrayList<>(filesByName.keySet())
|
||||
: requestedFiles.stream().map(AiWorkflowFileRequest::getFileName).toList();
|
||||
? new ArrayList<>(request.getFiles())
|
||||
: requestedFiles.stream().map(AiWorkflowFileRequest::getFile).toList();
|
||||
|
||||
Map<String, AiWorkflowFileRequest> requestedByName =
|
||||
Map<String, AiWorkflowFileRequest> requestedById =
|
||||
requestedFiles == null || requestedFiles.isEmpty()
|
||||
? Map.of()
|
||||
: requestedFiles.stream()
|
||||
.collect(
|
||||
Collectors.toMap(
|
||||
AiWorkflowFileRequest::getFileName, r -> r));
|
||||
.collect(Collectors.toMap(r -> r.getFile().getId(), r -> r));
|
||||
|
||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT));
|
||||
|
||||
List<LoadedFile> loadedFiles = new ArrayList<>();
|
||||
try {
|
||||
for (String fileName : fileNamesToLoad) {
|
||||
PDDocument doc = pdfDocumentFactory.load(filesByName.get(fileName), true);
|
||||
loadedFiles.add(new LoadedFile(fileName, doc));
|
||||
for (AiFile file : filesToLoad) {
|
||||
PDDocument doc = pdfDocumentFactory.load(filesById.get(file.getId()), true);
|
||||
loadedFiles.add(new LoadedFile(file.getId(), file.getName(), doc));
|
||||
}
|
||||
|
||||
List<PdfContentResult> contentResults =
|
||||
pdfContentExtractor.extractContent(
|
||||
loadedFiles,
|
||||
requestedByName,
|
||||
requestedById,
|
||||
response.getMaxPages(),
|
||||
response.getMaxCharacters());
|
||||
|
||||
@@ -201,7 +215,7 @@ public class AiWorkflowService {
|
||||
|
||||
WorkflowTurnRequest nextRequest = new WorkflowTurnRequest();
|
||||
nextRequest.setUserMessage(request.getUserMessage());
|
||||
nextRequest.setFileNames(request.getFileNames());
|
||||
nextRequest.setFiles(request.getFiles());
|
||||
nextRequest.setConversationHistory(request.getConversationHistory());
|
||||
nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults));
|
||||
nextRequest.setResumeWith(response.getResumeWith());
|
||||
@@ -217,10 +231,74 @@ public class AiWorkflowService {
|
||||
}
|
||||
}
|
||||
|
||||
private WorkflowState onNeedIngest(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesById,
|
||||
WorkflowTurnRequest request,
|
||||
ProgressListener listener)
|
||||
throws IOException {
|
||||
List<AiFile> filesToIngest = response.getFilesToIngest();
|
||||
if (filesToIngest == null || filesToIngest.isEmpty()) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue(
|
||||
"AI engine returned need_ingest without listing any files to ingest."));
|
||||
}
|
||||
// Guard against a retry loop: if we've already ingested this turn and the engine still
|
||||
// asks for more, something is wrong on its side.
|
||||
if (!request.getArtifacts().isEmpty() || request.getResumeWith() != null) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue(
|
||||
"AI engine requested ingest after the workflow had already been resumed."));
|
||||
}
|
||||
|
||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT));
|
||||
|
||||
for (AiFile file : filesToIngest) {
|
||||
MultipartFile multipartFile = filesById.get(file.getId());
|
||||
if (multipartFile == null) {
|
||||
return new WorkflowState.Terminal(
|
||||
cannotContinue(
|
||||
"AI engine requested ingest for unknown file: " + file.getName()));
|
||||
}
|
||||
ingestFile(file, multipartFile);
|
||||
}
|
||||
|
||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.PROCESSING));
|
||||
|
||||
WorkflowTurnRequest nextRequest = new WorkflowTurnRequest();
|
||||
nextRequest.setUserMessage(request.getUserMessage());
|
||||
nextRequest.setFiles(request.getFiles());
|
||||
nextRequest.setConversationHistory(request.getConversationHistory());
|
||||
nextRequest.setResumeWith(response.getResumeWith());
|
||||
return new WorkflowState.Pending(nextRequest);
|
||||
}
|
||||
|
||||
private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException {
|
||||
List<AiRagPageText> pages = new ArrayList<>();
|
||||
try (PDDocument document = pdfDocumentFactory.load(multipartFile, true)) {
|
||||
int pageCount = document.getNumberOfPages();
|
||||
for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
|
||||
String pageText = pdfContentExtractor.extractPageTextRaw(document, pageNumber);
|
||||
if (pageText != null && !pageText.isBlank()) {
|
||||
pages.add(new AiRagPageText(pageNumber, pageText));
|
||||
}
|
||||
}
|
||||
}
|
||||
AiRagIngestRequest ingestRequest =
|
||||
new AiRagIngestRequest(file.getId(), file.getName(), pages);
|
||||
String body = objectMapper.writeValueAsString(ingestRequest);
|
||||
aiEngineClient.postLongRunning(RAG_DOCUMENTS_ENDPOINT, body);
|
||||
log.debug(
|
||||
"Ingested file into RAG: id={}, name={}, pages={}",
|
||||
file.getId(),
|
||||
file.getName(),
|
||||
pages.size());
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private WorkflowState onToolCall(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
ProgressListener listener) {
|
||||
String endpointPath = response.getTool();
|
||||
Map<String, Object> parameters = response.getParameters();
|
||||
@@ -233,14 +311,14 @@ public class AiWorkflowService {
|
||||
}
|
||||
|
||||
try {
|
||||
List<Resource> inputFiles = toResources(filesByName);
|
||||
List<Resource> inputFiles = toResources(filesById);
|
||||
listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1));
|
||||
ToolResult result = executeStep(endpointPath, parameters, inputFiles);
|
||||
return new WorkflowState.Terminal(
|
||||
buildCompletedResponse(
|
||||
response.getRationale(),
|
||||
result.files(),
|
||||
new ArrayList<>(filesByName.keySet()),
|
||||
inputFileNames(filesById),
|
||||
result.report()));
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
|
||||
@@ -251,36 +329,23 @@ public class AiWorkflowService {
|
||||
|
||||
private WorkflowState onPlan(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
WorkflowTurnRequest previousRequest,
|
||||
ProgressListener listener) {
|
||||
return runPlan(
|
||||
response.getSteps(),
|
||||
response.getResumeWith(),
|
||||
response.getSummary(),
|
||||
filesByName,
|
||||
filesById,
|
||||
previousRequest,
|
||||
listener);
|
||||
}
|
||||
|
||||
private WorkflowState onAnswer(
|
||||
AiWorkflowResponse response,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
WorkflowTurnRequest previousRequest,
|
||||
ProgressListener listener) {
|
||||
AiWorkflowEditPlan plan = response.getEditPlan();
|
||||
if (plan != null) {
|
||||
// The engine wants us to run a side-quest before the answer is final.
|
||||
// Run the embedded plan and resume the orchestrator with the captured
|
||||
// report; the real answer arrives on the resume turn.
|
||||
return runPlan(
|
||||
plan.getSteps(),
|
||||
plan.getResumeWith(),
|
||||
plan.getSummary(),
|
||||
filesByName,
|
||||
previousRequest,
|
||||
listener);
|
||||
}
|
||||
return new WorkflowState.Terminal(response);
|
||||
}
|
||||
|
||||
@@ -289,7 +354,7 @@ public class AiWorkflowService {
|
||||
List<Map<String, Object>> steps,
|
||||
String resumeWith,
|
||||
String summary,
|
||||
Map<String, MultipartFile> filesByName,
|
||||
Map<String, MultipartFile> filesById,
|
||||
WorkflowTurnRequest previousRequest,
|
||||
ProgressListener listener) {
|
||||
if (steps == null || steps.isEmpty()) {
|
||||
@@ -298,7 +363,7 @@ public class AiWorkflowService {
|
||||
}
|
||||
|
||||
try {
|
||||
List<Resource> currentFiles = toResources(filesByName);
|
||||
List<Resource> currentFiles = toResources(filesById);
|
||||
// Propagate the *last* non-null report — the terminal step defines the output.
|
||||
JsonNode lastReport = null;
|
||||
String lastReportTool = null;
|
||||
@@ -331,7 +396,7 @@ public class AiWorkflowService {
|
||||
if (resumeWith != null && !resumeWith.isBlank() && lastReport != null) {
|
||||
WorkflowTurnRequest resumeRequest = new WorkflowTurnRequest();
|
||||
resumeRequest.setUserMessage(previousRequest.getUserMessage());
|
||||
resumeRequest.setFileNames(previousRequest.getFileNames());
|
||||
resumeRequest.setFiles(previousRequest.getFiles());
|
||||
resumeRequest.setConversationHistory(previousRequest.getConversationHistory());
|
||||
resumeRequest.setArtifacts(new ArrayList<>(previousRequest.getArtifacts()));
|
||||
resumeRequest
|
||||
@@ -345,10 +410,7 @@ public class AiWorkflowService {
|
||||
|
||||
return new WorkflowState.Terminal(
|
||||
buildCompletedResponse(
|
||||
summary,
|
||||
currentFiles,
|
||||
new ArrayList<>(filesByName.keySet()),
|
||||
lastReport));
|
||||
summary, currentFiles, inputFileNames(filesById), lastReport));
|
||||
} catch (Exception e) {
|
||||
log.error("Failed to execute plan: {}", e.getMessage(), e);
|
||||
return new WorkflowState.Terminal(
|
||||
@@ -356,6 +418,10 @@ public class AiWorkflowService {
|
||||
}
|
||||
}
|
||||
|
||||
private static List<String> inputFileNames(Map<String, MultipartFile> filesById) {
|
||||
return filesById.values().stream().map(MultipartFile::getOriginalFilename).toList();
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a single tool step. If the endpoint accepts multiple files, all files are sent in one
|
||||
* call. Otherwise, the endpoint is called once per file. ZIP responses are unpacked so each
|
||||
@@ -460,9 +526,9 @@ public class AiWorkflowService {
|
||||
}
|
||||
}
|
||||
|
||||
private List<Resource> toResources(Map<String, MultipartFile> filesByName) throws IOException {
|
||||
private List<Resource> toResources(Map<String, MultipartFile> filesById) throws IOException {
|
||||
List<Resource> resources = new ArrayList<>();
|
||||
for (MultipartFile file : filesByName.values()) {
|
||||
for (MultipartFile file : filesById.values()) {
|
||||
TempFile tempFile = tempFileManager.createManagedTempFile("ai-workflow");
|
||||
file.transferTo(tempFile.getPath());
|
||||
final String originalName = Filenames.toSimpleFileName(file.getOriginalFilename());
|
||||
@@ -554,7 +620,7 @@ public class AiWorkflowService {
|
||||
@Data
|
||||
private static class WorkflowTurnRequest {
|
||||
private String userMessage;
|
||||
private List<String> fileNames = new ArrayList<>();
|
||||
private List<AiFile> files = new ArrayList<>();
|
||||
private List<AiConversationMessage> conversationHistory = new ArrayList<>();
|
||||
private List<WorkflowArtifact> artifacts = new ArrayList<>();
|
||||
private String resumeWith;
|
||||
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.security.MessageDigest;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
/**
|
||||
* Content-addressable id derived from the SHA-256 hash of the uploaded bytes. Same content always
|
||||
* hashes to the same id, so re-uploads dedupe naturally in RAG. Suitable for session and SaaS
|
||||
* deployments; a folder-watch deployment would use a different strategy keyed by path.
|
||||
*/
|
||||
@Component
|
||||
public class ByteHashFileIdStrategy implements FileIdStrategy {
|
||||
|
||||
/**
|
||||
* Hex-char length of the returned id. 16 chars = 64 bits of collision space, which is plenty
|
||||
* for per-user document sets and keeps RAG collection names short.
|
||||
*/
|
||||
private static final int ID_HEX_LENGTH = 16;
|
||||
|
||||
private static final int BUFFER_SIZE = 64 * 1024;
|
||||
|
||||
@Override
|
||||
public String idFor(MultipartFile file) throws IOException {
|
||||
MessageDigest digest = sha256();
|
||||
try (InputStream in = file.getInputStream()) {
|
||||
byte[] buffer = new byte[BUFFER_SIZE];
|
||||
int read;
|
||||
while ((read = in.read(buffer)) != -1) {
|
||||
digest.update(buffer, 0, read);
|
||||
}
|
||||
}
|
||||
byte[] hash = digest.digest();
|
||||
StringBuilder hex = new StringBuilder(ID_HEX_LENGTH);
|
||||
for (int i = 0; hex.length() < ID_HEX_LENGTH; i++) {
|
||||
hex.append(String.format("%02x", hash[i]));
|
||||
}
|
||||
return hex.toString();
|
||||
}
|
||||
|
||||
private static MessageDigest sha256() {
|
||||
try {
|
||||
return MessageDigest.getInstance("SHA-256");
|
||||
} catch (NoSuchAlgorithmException e) {
|
||||
// SHA-256 is mandated by the JDK; absent only if the platform is broken.
|
||||
throw new IllegalStateException("SHA-256 not available", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
package stirling.software.proprietary.service;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
/**
|
||||
* Produces stable identifiers for uploaded files. The identifier is opaque to the AI engine and
|
||||
* serves as the RAG collection key when content is ingested. Swapping implementations (content
|
||||
* hash, filesystem path, tenant-scoped id, etc.) is how the system adapts to different deployment
|
||||
* models without any engine-side change.
|
||||
*/
|
||||
public interface FileIdStrategy {
|
||||
|
||||
String idFor(MultipartFile file) throws IOException;
|
||||
}
|
||||
+7
-3
@@ -44,7 +44,11 @@ public class PdfContentExtractor {
|
||||
|
||||
private static final int TEXT_PRESENCE_THRESHOLD = 20;
|
||||
|
||||
record LoadedFile(String fileName, PDDocument document) {}
|
||||
/**
|
||||
* A loaded PDF alongside the opaque file id used by the AI engine as its RAG collection key.
|
||||
* Keyed by id (not name) because filenames aren't unique across an upload.
|
||||
*/
|
||||
record LoadedFile(String id, String fileName, PDDocument document) {}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Low-level extraction methods (usable by any agent)
|
||||
@@ -126,7 +130,7 @@ public class PdfContentExtractor {
|
||||
*/
|
||||
List<PdfContentResult> extractContent(
|
||||
List<LoadedFile> loadedFiles,
|
||||
Map<String, AiWorkflowFileRequest> requestedByName,
|
||||
Map<String, AiWorkflowFileRequest> requestedById,
|
||||
int maxPages,
|
||||
int maxCharacters)
|
||||
throws IOException {
|
||||
@@ -136,7 +140,7 @@ public class PdfContentExtractor {
|
||||
|
||||
for (LoadedFile lf : loadedFiles) {
|
||||
if (remainingPages <= 0 || remainingCharacters <= 0) break;
|
||||
AiWorkflowFileRequest fileReq = requestedByName.get(lf.fileName());
|
||||
AiWorkflowFileRequest fileReq = requestedById.get(lf.id());
|
||||
List<AiPdfContentType> contentTypes =
|
||||
fileReq != null && !fileReq.getContentTypes().isEmpty()
|
||||
? fileReq.getContentTypes()
|
||||
|
||||
+56
-2
@@ -3,8 +3,11 @@ package stirling.software.proprietary.service;
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertNotNull;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyBoolean;
|
||||
import static org.mockito.ArgumentMatchers.anyInt;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.lenient;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
@@ -21,6 +24,8 @@ import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.zip.ZipEntry;
|
||||
import java.util.zip.ZipOutputStream;
|
||||
|
||||
import org.apache.pdfbox.pdmodel.PDDocument;
|
||||
import org.apache.pdfbox.pdmodel.PDPage;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
@@ -32,6 +37,7 @@ import org.springframework.core.io.Resource;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.mock.web.MockMultipartFile;
|
||||
import org.springframework.util.MultiValueMap;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import stirling.software.common.model.ApplicationProperties;
|
||||
import stirling.software.common.service.CustomPDFDocumentFactory;
|
||||
@@ -72,6 +78,7 @@ class AiWorkflowServiceTest {
|
||||
@Mock private InternalApiClient internalApiClient;
|
||||
@Mock private FileStorage fileStorage;
|
||||
@Mock private ToolMetadataService toolMetadataService;
|
||||
@Mock private FileIdStrategy fileIdStrategy;
|
||||
|
||||
@TempDir Path tempDir;
|
||||
|
||||
@@ -80,13 +87,19 @@ class AiWorkflowServiceTest {
|
||||
private AiWorkflowService service;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
void setUp() throws IOException {
|
||||
ApplicationProperties props = new ApplicationProperties();
|
||||
props.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString());
|
||||
props.getSystem().getTempFileManagement().setPrefix("ai-test-");
|
||||
tempFileManager = new TempFileManager(new TempFileRegistry(), props);
|
||||
objectMapper = JsonMapper.builder().build();
|
||||
|
||||
// Mock strategy yields the filename as id so each MockMultipartFile in a test gets a
|
||||
// distinct collection key. Real strategy (ByteHashFileIdStrategy) hashes bytes.
|
||||
lenient()
|
||||
.when(fileIdStrategy.idFor(any(MultipartFile.class)))
|
||||
.thenAnswer(inv -> ((MultipartFile) inv.getArgument(0)).getOriginalFilename());
|
||||
|
||||
service =
|
||||
new AiWorkflowService(
|
||||
pdfDocumentFactory,
|
||||
@@ -96,7 +109,8 @@ class AiWorkflowServiceTest {
|
||||
internalApiClient,
|
||||
fileStorage,
|
||||
toolMetadataService,
|
||||
tempFileManager);
|
||||
tempFileManager,
|
||||
fileIdStrategy);
|
||||
}
|
||||
|
||||
@Test
|
||||
@@ -246,6 +260,46 @@ class AiWorkflowServiceTest {
|
||||
verify(internalApiClient, never()).post(anyString(), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
void needIngestExtractsPageTextAndPostsToRagThenRetries() throws IOException {
|
||||
MockMultipartFile input = pdf("report.pdf", "bytes");
|
||||
when(fileIdStrategy.idFor(any())).thenReturn("report-id");
|
||||
|
||||
PDDocument document = new PDDocument();
|
||||
document.addPage(new PDPage());
|
||||
document.addPage(new PDPage());
|
||||
when(pdfDocumentFactory.load(any(MultipartFile.class), anyBoolean())).thenReturn(document);
|
||||
when(pdfContentExtractor.extractPageTextRaw(eq(document), anyInt()))
|
||||
.thenReturn("page content");
|
||||
|
||||
int[] orchestratorCalls = {0};
|
||||
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString()))
|
||||
.thenAnswer(
|
||||
inv -> {
|
||||
orchestratorCalls[0]++;
|
||||
if (orchestratorCalls[0] == 1) {
|
||||
return """
|
||||
{
|
||||
"outcome":"need_ingest",
|
||||
"resumeWith":"pdf_question",
|
||||
"reason":"ingest first",
|
||||
"filesToIngest":[{"id":"report-id","name":"report.pdf"}],
|
||||
"contentTypes":["page_text"]
|
||||
}
|
||||
""";
|
||||
}
|
||||
return """
|
||||
{"outcome":"answer","answer":"done","evidence":[]}
|
||||
""";
|
||||
});
|
||||
|
||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
|
||||
|
||||
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
|
||||
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/rag/documents"), anyString());
|
||||
verify(aiEngineClient, times(2)).post(eq("/api/v1/orchestrator"), anyString());
|
||||
}
|
||||
|
||||
// --- helpers ---
|
||||
|
||||
private void stubOrchestrator(String responseJson) throws IOException {
|
||||
|
||||
Reference in New Issue
Block a user