mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +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:
@@ -237,6 +237,13 @@ public class ApplicationProperties {
|
|||||||
private boolean enabled = false;
|
private boolean enabled = false;
|
||||||
private String url = "http://localhost:5001";
|
private String url = "http://localhost:5001";
|
||||||
private int timeoutSeconds = 120;
|
private int timeoutSeconds = 120;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Longer timeout for heavy operations like RAG ingestion, which embeds the whole document
|
||||||
|
* and can take multiple minutes for large books. Applied per-call when the caller
|
||||||
|
* explicitly requests it via {@code AiEngineClient.postWithTimeout}.
|
||||||
|
*/
|
||||||
|
private int longRunningTimeoutSeconds = 600;
|
||||||
}
|
}
|
||||||
|
|
||||||
@Data
|
@Data
|
||||||
|
|||||||
@@ -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")
|
@Schema(description = "Per-file content extraction request from the AI engine")
|
||||||
public class AiWorkflowFileRequest {
|
public class AiWorkflowFileRequest {
|
||||||
|
|
||||||
@Schema(description = "Original filename of the requested file", example = "contract.pdf")
|
@Schema(description = "The file the engine wants content extracted for")
|
||||||
private String fileName;
|
private AiFile file;
|
||||||
|
|
||||||
@Schema(description = "Specific 1-based page numbers to extract from this file")
|
@Schema(description = "Specific 1-based page numbers to extract from this file")
|
||||||
private List<Integer> pageNumbers = new ArrayList<>();
|
private List<Integer> pageNumbers = new ArrayList<>();
|
||||||
|
|||||||
+1
@@ -12,6 +12,7 @@ public enum AiWorkflowOutcome {
|
|||||||
ANSWER("answer"),
|
ANSWER("answer"),
|
||||||
NOT_FOUND("not_found"),
|
NOT_FOUND("not_found"),
|
||||||
NEED_CONTENT("need_content"),
|
NEED_CONTENT("need_content"),
|
||||||
|
NEED_INGEST("need_ingest"),
|
||||||
PLAN("plan"),
|
PLAN("plan"),
|
||||||
NEED_CLARIFICATION("need_clarification"),
|
NEED_CLARIFICATION("need_clarification"),
|
||||||
CANNOT_DO("cannot_do"),
|
CANNOT_DO("cannot_do"),
|
||||||
|
|||||||
+6
-7
@@ -73,6 +73,12 @@ public class AiWorkflowResponse {
|
|||||||
@Schema(description = "Per-file text extraction requests from the AI engine")
|
@Schema(description = "Per-file text extraction requests from the AI engine")
|
||||||
private List<AiWorkflowFileRequest> files = new ArrayList<>();
|
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")
|
@Schema(description = "Maximum number of pages the AI engine wants text extracted from")
|
||||||
private Integer maxPages;
|
private Integer maxPages;
|
||||||
|
|
||||||
@@ -89,11 +95,4 @@ public class AiWorkflowResponse {
|
|||||||
+ " body or via the X-Stirling-Tool-Report header. May be null for tools"
|
+ " body or via the X-Stirling-Tool-Report header. May be null for tools"
|
||||||
+ " that produce only a file.")
|
+ " that produce only a file.")
|
||||||
private JsonNode report;
|
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 {
|
public String post(String path, String jsonBody) throws IOException {
|
||||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
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()) {
|
if (!config.isEnabled()) {
|
||||||
throw new ResponseStatusException(
|
throw new ResponseStatusException(
|
||||||
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
|
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
|
||||||
}
|
}
|
||||||
|
|
||||||
String url = config.getUrl().stripTrailing() + path;
|
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 request =
|
||||||
HttpRequest.newBuilder()
|
HttpRequest.newBuilder()
|
||||||
.uri(URI.create(url))
|
.uri(URI.create(url))
|
||||||
.header("Content-Type", "application/json")
|
.header("Content-Type", "application/json")
|
||||||
.header("Accept", "application/json")
|
.header("Accept", "application/json")
|
||||||
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
|
.timeout(timeout)
|
||||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||||
.build();
|
.build();
|
||||||
|
|
||||||
|
|||||||
+123
-57
@@ -36,7 +36,9 @@ import stirling.software.common.util.TempFile;
|
|||||||
import stirling.software.common.util.TempFileManager;
|
import stirling.software.common.util.TempFileManager;
|
||||||
import stirling.software.common.util.ZipExtractionUtils;
|
import stirling.software.common.util.ZipExtractionUtils;
|
||||||
import stirling.software.proprietary.model.api.ai.AiConversationMessage;
|
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.AiWorkflowFileInput;
|
||||||
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
|
import stirling.software.proprietary.model.api.ai.AiWorkflowFileRequest;
|
||||||
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
|
import stirling.software.proprietary.model.api.ai.AiWorkflowOutcome;
|
||||||
@@ -58,6 +60,8 @@ import tools.jackson.databind.ObjectMapper;
|
|||||||
@RequiredArgsConstructor
|
@RequiredArgsConstructor
|
||||||
public class AiWorkflowService {
|
public class AiWorkflowService {
|
||||||
|
|
||||||
|
private static final String RAG_DOCUMENTS_ENDPOINT = "/api/v1/rag/documents";
|
||||||
|
|
||||||
private final CustomPDFDocumentFactory pdfDocumentFactory;
|
private final CustomPDFDocumentFactory pdfDocumentFactory;
|
||||||
private final AiEngineClient aiEngineClient;
|
private final AiEngineClient aiEngineClient;
|
||||||
private final PdfContentExtractor pdfContentExtractor;
|
private final PdfContentExtractor pdfContentExtractor;
|
||||||
@@ -66,6 +70,7 @@ public class AiWorkflowService {
|
|||||||
private final FileStorage fileStorage;
|
private final FileStorage fileStorage;
|
||||||
private final ToolMetadataService toolMetadataService;
|
private final ToolMetadataService toolMetadataService;
|
||||||
private final TempFileManager tempFileManager;
|
private final TempFileManager tempFileManager;
|
||||||
|
private final FileIdStrategy fileIdStrategy;
|
||||||
|
|
||||||
@FunctionalInterface
|
@FunctionalInterface
|
||||||
public interface ProgressListener {
|
public interface ProgressListener {
|
||||||
@@ -98,15 +103,24 @@ public class AiWorkflowService {
|
|||||||
throws IOException {
|
throws IOException {
|
||||||
validateRequest(request);
|
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()) {
|
for (AiWorkflowFileInput fileInput : request.getFileInputs()) {
|
||||||
filesByName.put(
|
MultipartFile multipartFile = fileInput.getFileInput();
|
||||||
fileInput.getFileInput().getOriginalFilename(), fileInput.getFileInput());
|
AiFile aiFile =
|
||||||
|
new AiFile(
|
||||||
|
fileIdStrategy.idFor(multipartFile),
|
||||||
|
multipartFile.getOriginalFilename());
|
||||||
|
filesById.put(aiFile.getId(), multipartFile);
|
||||||
|
files.add(aiFile);
|
||||||
}
|
}
|
||||||
|
|
||||||
WorkflowTurnRequest initialRequest = new WorkflowTurnRequest();
|
WorkflowTurnRequest initialRequest = new WorkflowTurnRequest();
|
||||||
initialRequest.setUserMessage(request.getUserMessage().trim());
|
initialRequest.setUserMessage(request.getUserMessage().trim());
|
||||||
initialRequest.setFileNames(new ArrayList<>(filesByName.keySet()));
|
initialRequest.setFiles(files);
|
||||||
initialRequest.setConversationHistory(
|
initialRequest.setConversationHistory(
|
||||||
request.getConversationHistory() == null
|
request.getConversationHistory() == null
|
||||||
? new ArrayList<>()
|
? new ArrayList<>()
|
||||||
@@ -116,23 +130,24 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
WorkflowState state = new WorkflowState.Pending(initialRequest);
|
WorkflowState state = new WorkflowState.Pending(initialRequest);
|
||||||
while (state instanceof WorkflowState.Pending pending) {
|
while (state instanceof WorkflowState.Pending pending) {
|
||||||
state = advance(pending.request(), filesByName, listener);
|
state = advance(pending.request(), filesById, listener);
|
||||||
}
|
}
|
||||||
return ((WorkflowState.Terminal) state).response();
|
return ((WorkflowState.Terminal) state).response();
|
||||||
}
|
}
|
||||||
|
|
||||||
private WorkflowState advance(
|
private WorkflowState advance(
|
||||||
WorkflowTurnRequest request,
|
WorkflowTurnRequest request,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
ProgressListener listener)
|
ProgressListener listener)
|
||||||
throws IOException {
|
throws IOException {
|
||||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
|
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.CALLING_ENGINE));
|
||||||
AiWorkflowResponse response = invokeOrchestrator(request);
|
AiWorkflowResponse response = invokeOrchestrator(request);
|
||||||
return switch (response.getOutcome()) {
|
return switch (response.getOutcome()) {
|
||||||
case NEED_CONTENT -> onNeedContent(response, filesByName, request, listener);
|
case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
|
||||||
case TOOL_CALL -> onToolCall(response, filesByName, listener);
|
case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
|
||||||
case PLAN -> onPlan(response, filesByName, request, listener);
|
case TOOL_CALL -> onToolCall(response, filesById, listener);
|
||||||
case ANSWER -> onAnswer(response, filesByName, request, listener);
|
case PLAN -> onPlan(response, filesById, request, listener);
|
||||||
|
case ANSWER -> onAnswer(response, filesById, request, listener);
|
||||||
case NOT_FOUND,
|
case NOT_FOUND,
|
||||||
NEED_CLARIFICATION,
|
NEED_CLARIFICATION,
|
||||||
CANNOT_DO,
|
CANNOT_DO,
|
||||||
@@ -146,7 +161,7 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
private WorkflowState onNeedContent(
|
private WorkflowState onNeedContent(
|
||||||
AiWorkflowResponse response,
|
AiWorkflowResponse response,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
WorkflowTurnRequest request,
|
WorkflowTurnRequest request,
|
||||||
ProgressListener listener)
|
ProgressListener listener)
|
||||||
throws IOException {
|
throws IOException {
|
||||||
@@ -157,43 +172,42 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
List<AiWorkflowFileRequest> requestedFiles = response.getFiles();
|
List<AiWorkflowFileRequest> requestedFiles = response.getFiles();
|
||||||
|
|
||||||
// Validate requested file names before loading anything
|
// Validate requested file ids before loading anything
|
||||||
if (requestedFiles != null && !requestedFiles.isEmpty()) {
|
if (requestedFiles != null && !requestedFiles.isEmpty()) {
|
||||||
for (AiWorkflowFileRequest fileReq : requestedFiles) {
|
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(
|
return new WorkflowState.Terminal(
|
||||||
cannotContinue(
|
cannotContinue("AI engine requested unknown file: " + display));
|
||||||
"AI engine requested unknown file: " + fileReq.getFileName()));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
List<String> fileNamesToLoad =
|
List<AiFile> filesToLoad =
|
||||||
(requestedFiles == null || requestedFiles.isEmpty())
|
(requestedFiles == null || requestedFiles.isEmpty())
|
||||||
? new ArrayList<>(filesByName.keySet())
|
? new ArrayList<>(request.getFiles())
|
||||||
: requestedFiles.stream().map(AiWorkflowFileRequest::getFileName).toList();
|
: requestedFiles.stream().map(AiWorkflowFileRequest::getFile).toList();
|
||||||
|
|
||||||
Map<String, AiWorkflowFileRequest> requestedByName =
|
Map<String, AiWorkflowFileRequest> requestedById =
|
||||||
requestedFiles == null || requestedFiles.isEmpty()
|
requestedFiles == null || requestedFiles.isEmpty()
|
||||||
? Map.of()
|
? Map.of()
|
||||||
: requestedFiles.stream()
|
: requestedFiles.stream()
|
||||||
.collect(
|
.collect(Collectors.toMap(r -> r.getFile().getId(), r -> r));
|
||||||
Collectors.toMap(
|
|
||||||
AiWorkflowFileRequest::getFileName, r -> r));
|
|
||||||
|
|
||||||
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT));
|
listener.onProgress(AiWorkflowProgressEvent.of(AiWorkflowPhase.EXTRACTING_CONTENT));
|
||||||
|
|
||||||
List<LoadedFile> loadedFiles = new ArrayList<>();
|
List<LoadedFile> loadedFiles = new ArrayList<>();
|
||||||
try {
|
try {
|
||||||
for (String fileName : fileNamesToLoad) {
|
for (AiFile file : filesToLoad) {
|
||||||
PDDocument doc = pdfDocumentFactory.load(filesByName.get(fileName), true);
|
PDDocument doc = pdfDocumentFactory.load(filesById.get(file.getId()), true);
|
||||||
loadedFiles.add(new LoadedFile(fileName, doc));
|
loadedFiles.add(new LoadedFile(file.getId(), file.getName(), doc));
|
||||||
}
|
}
|
||||||
|
|
||||||
List<PdfContentResult> contentResults =
|
List<PdfContentResult> contentResults =
|
||||||
pdfContentExtractor.extractContent(
|
pdfContentExtractor.extractContent(
|
||||||
loadedFiles,
|
loadedFiles,
|
||||||
requestedByName,
|
requestedById,
|
||||||
response.getMaxPages(),
|
response.getMaxPages(),
|
||||||
response.getMaxCharacters());
|
response.getMaxCharacters());
|
||||||
|
|
||||||
@@ -201,7 +215,7 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
WorkflowTurnRequest nextRequest = new WorkflowTurnRequest();
|
WorkflowTurnRequest nextRequest = new WorkflowTurnRequest();
|
||||||
nextRequest.setUserMessage(request.getUserMessage());
|
nextRequest.setUserMessage(request.getUserMessage());
|
||||||
nextRequest.setFileNames(request.getFileNames());
|
nextRequest.setFiles(request.getFiles());
|
||||||
nextRequest.setConversationHistory(request.getConversationHistory());
|
nextRequest.setConversationHistory(request.getConversationHistory());
|
||||||
nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults));
|
nextRequest.setArtifacts(pdfContentExtractor.buildArtifacts(contentResults));
|
||||||
nextRequest.setResumeWith(response.getResumeWith());
|
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")
|
@SuppressWarnings("unchecked")
|
||||||
private WorkflowState onToolCall(
|
private WorkflowState onToolCall(
|
||||||
AiWorkflowResponse response,
|
AiWorkflowResponse response,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
ProgressListener listener) {
|
ProgressListener listener) {
|
||||||
String endpointPath = response.getTool();
|
String endpointPath = response.getTool();
|
||||||
Map<String, Object> parameters = response.getParameters();
|
Map<String, Object> parameters = response.getParameters();
|
||||||
@@ -233,14 +311,14 @@ public class AiWorkflowService {
|
|||||||
}
|
}
|
||||||
|
|
||||||
try {
|
try {
|
||||||
List<Resource> inputFiles = toResources(filesByName);
|
List<Resource> inputFiles = toResources(filesById);
|
||||||
listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1));
|
listener.onProgress(AiWorkflowProgressEvent.executingTool(endpointPath, 1, 1));
|
||||||
ToolResult result = executeStep(endpointPath, parameters, inputFiles);
|
ToolResult result = executeStep(endpointPath, parameters, inputFiles);
|
||||||
return new WorkflowState.Terminal(
|
return new WorkflowState.Terminal(
|
||||||
buildCompletedResponse(
|
buildCompletedResponse(
|
||||||
response.getRationale(),
|
response.getRationale(),
|
||||||
result.files(),
|
result.files(),
|
||||||
new ArrayList<>(filesByName.keySet()),
|
inputFileNames(filesById),
|
||||||
result.report()));
|
result.report()));
|
||||||
} catch (Exception e) {
|
} catch (Exception e) {
|
||||||
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
|
log.error("Failed to execute tool {}: {}", endpointPath, e.getMessage(), e);
|
||||||
@@ -251,36 +329,23 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
private WorkflowState onPlan(
|
private WorkflowState onPlan(
|
||||||
AiWorkflowResponse response,
|
AiWorkflowResponse response,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
WorkflowTurnRequest previousRequest,
|
WorkflowTurnRequest previousRequest,
|
||||||
ProgressListener listener) {
|
ProgressListener listener) {
|
||||||
return runPlan(
|
return runPlan(
|
||||||
response.getSteps(),
|
response.getSteps(),
|
||||||
response.getResumeWith(),
|
response.getResumeWith(),
|
||||||
response.getSummary(),
|
response.getSummary(),
|
||||||
filesByName,
|
filesById,
|
||||||
previousRequest,
|
previousRequest,
|
||||||
listener);
|
listener);
|
||||||
}
|
}
|
||||||
|
|
||||||
private WorkflowState onAnswer(
|
private WorkflowState onAnswer(
|
||||||
AiWorkflowResponse response,
|
AiWorkflowResponse response,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
WorkflowTurnRequest previousRequest,
|
WorkflowTurnRequest previousRequest,
|
||||||
ProgressListener listener) {
|
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);
|
return new WorkflowState.Terminal(response);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -289,7 +354,7 @@ public class AiWorkflowService {
|
|||||||
List<Map<String, Object>> steps,
|
List<Map<String, Object>> steps,
|
||||||
String resumeWith,
|
String resumeWith,
|
||||||
String summary,
|
String summary,
|
||||||
Map<String, MultipartFile> filesByName,
|
Map<String, MultipartFile> filesById,
|
||||||
WorkflowTurnRequest previousRequest,
|
WorkflowTurnRequest previousRequest,
|
||||||
ProgressListener listener) {
|
ProgressListener listener) {
|
||||||
if (steps == null || steps.isEmpty()) {
|
if (steps == null || steps.isEmpty()) {
|
||||||
@@ -298,7 +363,7 @@ public class AiWorkflowService {
|
|||||||
}
|
}
|
||||||
|
|
||||||
try {
|
try {
|
||||||
List<Resource> currentFiles = toResources(filesByName);
|
List<Resource> currentFiles = toResources(filesById);
|
||||||
// Propagate the *last* non-null report — the terminal step defines the output.
|
// Propagate the *last* non-null report — the terminal step defines the output.
|
||||||
JsonNode lastReport = null;
|
JsonNode lastReport = null;
|
||||||
String lastReportTool = null;
|
String lastReportTool = null;
|
||||||
@@ -331,7 +396,7 @@ public class AiWorkflowService {
|
|||||||
if (resumeWith != null && !resumeWith.isBlank() && lastReport != null) {
|
if (resumeWith != null && !resumeWith.isBlank() && lastReport != null) {
|
||||||
WorkflowTurnRequest resumeRequest = new WorkflowTurnRequest();
|
WorkflowTurnRequest resumeRequest = new WorkflowTurnRequest();
|
||||||
resumeRequest.setUserMessage(previousRequest.getUserMessage());
|
resumeRequest.setUserMessage(previousRequest.getUserMessage());
|
||||||
resumeRequest.setFileNames(previousRequest.getFileNames());
|
resumeRequest.setFiles(previousRequest.getFiles());
|
||||||
resumeRequest.setConversationHistory(previousRequest.getConversationHistory());
|
resumeRequest.setConversationHistory(previousRequest.getConversationHistory());
|
||||||
resumeRequest.setArtifacts(new ArrayList<>(previousRequest.getArtifacts()));
|
resumeRequest.setArtifacts(new ArrayList<>(previousRequest.getArtifacts()));
|
||||||
resumeRequest
|
resumeRequest
|
||||||
@@ -345,10 +410,7 @@ public class AiWorkflowService {
|
|||||||
|
|
||||||
return new WorkflowState.Terminal(
|
return new WorkflowState.Terminal(
|
||||||
buildCompletedResponse(
|
buildCompletedResponse(
|
||||||
summary,
|
summary, currentFiles, inputFileNames(filesById), lastReport));
|
||||||
currentFiles,
|
|
||||||
new ArrayList<>(filesByName.keySet()),
|
|
||||||
lastReport));
|
|
||||||
} catch (Exception e) {
|
} catch (Exception e) {
|
||||||
log.error("Failed to execute plan: {}", e.getMessage(), e);
|
log.error("Failed to execute plan: {}", e.getMessage(), e);
|
||||||
return new WorkflowState.Terminal(
|
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
|
* 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
|
* 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<>();
|
List<Resource> resources = new ArrayList<>();
|
||||||
for (MultipartFile file : filesByName.values()) {
|
for (MultipartFile file : filesById.values()) {
|
||||||
TempFile tempFile = tempFileManager.createManagedTempFile("ai-workflow");
|
TempFile tempFile = tempFileManager.createManagedTempFile("ai-workflow");
|
||||||
file.transferTo(tempFile.getPath());
|
file.transferTo(tempFile.getPath());
|
||||||
final String originalName = Filenames.toSimpleFileName(file.getOriginalFilename());
|
final String originalName = Filenames.toSimpleFileName(file.getOriginalFilename());
|
||||||
@@ -554,7 +620,7 @@ public class AiWorkflowService {
|
|||||||
@Data
|
@Data
|
||||||
private static class WorkflowTurnRequest {
|
private static class WorkflowTurnRequest {
|
||||||
private String userMessage;
|
private String userMessage;
|
||||||
private List<String> fileNames = new ArrayList<>();
|
private List<AiFile> files = new ArrayList<>();
|
||||||
private List<AiConversationMessage> conversationHistory = new ArrayList<>();
|
private List<AiConversationMessage> conversationHistory = new ArrayList<>();
|
||||||
private List<WorkflowArtifact> artifacts = new ArrayList<>();
|
private List<WorkflowArtifact> artifacts = new ArrayList<>();
|
||||||
private String resumeWith;
|
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;
|
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)
|
// Low-level extraction methods (usable by any agent)
|
||||||
@@ -126,7 +130,7 @@ public class PdfContentExtractor {
|
|||||||
*/
|
*/
|
||||||
List<PdfContentResult> extractContent(
|
List<PdfContentResult> extractContent(
|
||||||
List<LoadedFile> loadedFiles,
|
List<LoadedFile> loadedFiles,
|
||||||
Map<String, AiWorkflowFileRequest> requestedByName,
|
Map<String, AiWorkflowFileRequest> requestedById,
|
||||||
int maxPages,
|
int maxPages,
|
||||||
int maxCharacters)
|
int maxCharacters)
|
||||||
throws IOException {
|
throws IOException {
|
||||||
@@ -136,7 +140,7 @@ public class PdfContentExtractor {
|
|||||||
|
|
||||||
for (LoadedFile lf : loadedFiles) {
|
for (LoadedFile lf : loadedFiles) {
|
||||||
if (remainingPages <= 0 || remainingCharacters <= 0) break;
|
if (remainingPages <= 0 || remainingCharacters <= 0) break;
|
||||||
AiWorkflowFileRequest fileReq = requestedByName.get(lf.fileName());
|
AiWorkflowFileRequest fileReq = requestedById.get(lf.id());
|
||||||
List<AiPdfContentType> contentTypes =
|
List<AiPdfContentType> contentTypes =
|
||||||
fileReq != null && !fileReq.getContentTypes().isEmpty()
|
fileReq != null && !fileReq.getContentTypes().isEmpty()
|
||||||
? fileReq.getContentTypes()
|
? 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.assertEquals;
|
||||||
import static org.junit.jupiter.api.Assertions.assertNotNull;
|
import static org.junit.jupiter.api.Assertions.assertNotNull;
|
||||||
import static org.mockito.ArgumentMatchers.any;
|
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.anyString;
|
||||||
import static org.mockito.ArgumentMatchers.eq;
|
import static org.mockito.ArgumentMatchers.eq;
|
||||||
|
import static org.mockito.Mockito.lenient;
|
||||||
import static org.mockito.Mockito.never;
|
import static org.mockito.Mockito.never;
|
||||||
import static org.mockito.Mockito.times;
|
import static org.mockito.Mockito.times;
|
||||||
import static org.mockito.Mockito.verify;
|
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.ZipEntry;
|
||||||
import java.util.zip.ZipOutputStream;
|
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.BeforeEach;
|
||||||
import org.junit.jupiter.api.Test;
|
import org.junit.jupiter.api.Test;
|
||||||
import org.junit.jupiter.api.extension.ExtendWith;
|
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.http.ResponseEntity;
|
||||||
import org.springframework.mock.web.MockMultipartFile;
|
import org.springframework.mock.web.MockMultipartFile;
|
||||||
import org.springframework.util.MultiValueMap;
|
import org.springframework.util.MultiValueMap;
|
||||||
|
import org.springframework.web.multipart.MultipartFile;
|
||||||
|
|
||||||
import stirling.software.common.model.ApplicationProperties;
|
import stirling.software.common.model.ApplicationProperties;
|
||||||
import stirling.software.common.service.CustomPDFDocumentFactory;
|
import stirling.software.common.service.CustomPDFDocumentFactory;
|
||||||
@@ -72,6 +78,7 @@ class AiWorkflowServiceTest {
|
|||||||
@Mock private InternalApiClient internalApiClient;
|
@Mock private InternalApiClient internalApiClient;
|
||||||
@Mock private FileStorage fileStorage;
|
@Mock private FileStorage fileStorage;
|
||||||
@Mock private ToolMetadataService toolMetadataService;
|
@Mock private ToolMetadataService toolMetadataService;
|
||||||
|
@Mock private FileIdStrategy fileIdStrategy;
|
||||||
|
|
||||||
@TempDir Path tempDir;
|
@TempDir Path tempDir;
|
||||||
|
|
||||||
@@ -80,13 +87,19 @@ class AiWorkflowServiceTest {
|
|||||||
private AiWorkflowService service;
|
private AiWorkflowService service;
|
||||||
|
|
||||||
@BeforeEach
|
@BeforeEach
|
||||||
void setUp() {
|
void setUp() throws IOException {
|
||||||
ApplicationProperties props = new ApplicationProperties();
|
ApplicationProperties props = new ApplicationProperties();
|
||||||
props.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString());
|
props.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString());
|
||||||
props.getSystem().getTempFileManagement().setPrefix("ai-test-");
|
props.getSystem().getTempFileManagement().setPrefix("ai-test-");
|
||||||
tempFileManager = new TempFileManager(new TempFileRegistry(), props);
|
tempFileManager = new TempFileManager(new TempFileRegistry(), props);
|
||||||
objectMapper = JsonMapper.builder().build();
|
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 =
|
service =
|
||||||
new AiWorkflowService(
|
new AiWorkflowService(
|
||||||
pdfDocumentFactory,
|
pdfDocumentFactory,
|
||||||
@@ -96,7 +109,8 @@ class AiWorkflowServiceTest {
|
|||||||
internalApiClient,
|
internalApiClient,
|
||||||
fileStorage,
|
fileStorage,
|
||||||
toolMetadataService,
|
toolMetadataService,
|
||||||
tempFileManager);
|
tempFileManager,
|
||||||
|
fileIdStrategy);
|
||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
@@ -246,6 +260,46 @@ class AiWorkflowServiceTest {
|
|||||||
verify(internalApiClient, never()).post(anyString(), any());
|
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 ---
|
// --- helpers ---
|
||||||
|
|
||||||
private void stubOrchestrator(String responseJson) throws IOException {
|
private void stubOrchestrator(String responseJson) throws IOException {
|
||||||
|
|||||||
+6
-1
@@ -30,7 +30,12 @@ STIRLING_RAG_PGVECTOR_DSN=
|
|||||||
|
|
||||||
STIRLING_RAG_CHUNK_SIZE=512
|
STIRLING_RAG_CHUNK_SIZE=512
|
||||||
STIRLING_RAG_CHUNK_OVERLAP=64
|
STIRLING_RAG_CHUNK_OVERLAP=64
|
||||||
STIRLING_RAG_TOP_K=5
|
STIRLING_RAG_TOP_K=20
|
||||||
|
|
||||||
|
# Per-run cap on ``search_knowledge`` calls. After this many calls the tool is
|
||||||
|
# removed from the agent's toolset so it must answer from what it already retrieved
|
||||||
|
# rather than chain more searches.
|
||||||
|
STIRLING_RAG_MAX_SEARCHES=5
|
||||||
|
|
||||||
# Upper bounds on PDF page text the engine will request per extraction round.
|
# Upper bounds on PDF page text the engine will request per extraction round.
|
||||||
STIRLING_MAX_PAGES=200
|
STIRLING_MAX_PAGES=200
|
||||||
|
|||||||
@@ -18,10 +18,11 @@ from stirling.contracts import (
|
|||||||
OrchestratorRequest,
|
OrchestratorRequest,
|
||||||
OrchestratorResponse,
|
OrchestratorResponse,
|
||||||
PdfEditResponse,
|
PdfEditResponse,
|
||||||
PdfQuestionResponse,
|
PdfQuestionOrchestrateResponse,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
UnsupportedCapabilityResponse,
|
UnsupportedCapabilityResponse,
|
||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
|
format_file_names,
|
||||||
)
|
)
|
||||||
from stirling.contracts.pdf_edit import EditPlanResponse
|
from stirling.contracts.pdf_edit import EditPlanResponse
|
||||||
from stirling.services import AppRuntime
|
from stirling.services import AppRuntime
|
||||||
@@ -78,12 +79,13 @@ class OrchestratorAgent:
|
|||||||
"You are the top-level orchestrator. "
|
"You are the top-level orchestrator. "
|
||||||
"Choose exactly one output function that best handles the request. "
|
"Choose exactly one output function that best handles the request. "
|
||||||
"Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
|
"Use delegate_pdf_edit for requested modifications of single or multiple PDFs. "
|
||||||
"Use delegate_pdf_question for questions about PDF contents. "
|
"Use delegate_pdf_question for questions about the contents of the attached PDFs. "
|
||||||
"Use delegate_user_spec for requests to create or define an agent spec. "
|
"Use delegate_user_spec for requests to create or define an agent spec. "
|
||||||
"Use delegate_pdf_review when the user wants the PDF returned with review"
|
"Use delegate_pdf_review when the user wants the PDF returned with review"
|
||||||
" comments attached — anything like 'review this', 'annotate with comments',"
|
" comments attached — anything like 'review this', 'annotate with comments',"
|
||||||
" 'leave feedback on the PDF'. "
|
" 'leave feedback on the PDF'. "
|
||||||
"Use unsupported_capability only when none of the other outputs fit."
|
"Use unsupported_capability when the user asks about the assistant itself "
|
||||||
|
"or when none of the other outputs fit; supply a helpful message."
|
||||||
),
|
),
|
||||||
model_settings=runtime.fast_model_settings,
|
model_settings=runtime.fast_model_settings,
|
||||||
)
|
)
|
||||||
@@ -91,7 +93,7 @@ class OrchestratorAgent:
|
|||||||
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
|
async def handle(self, request: OrchestratorRequest) -> OrchestratorResponse:
|
||||||
logger.info(
|
logger.info(
|
||||||
"[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r",
|
"[orchestrator] handle: files=%s resume_with=%s artifacts=%s msg=%r",
|
||||||
request.file_names,
|
[file.name for file in request.files],
|
||||||
request.resume_with,
|
request.resume_with,
|
||||||
[type(a).__name__ for a in request.artifacts],
|
[type(a).__name__ for a in request.artifacts],
|
||||||
request.user_message,
|
request.user_message,
|
||||||
@@ -137,10 +139,10 @@ class OrchestratorAgent:
|
|||||||
async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
|
async def _run_pdf_edit(self, request: OrchestratorRequest) -> PdfEditResponse:
|
||||||
return await PdfEditAgent(self.runtime).orchestrate(request)
|
return await PdfEditAgent(self.runtime).orchestrate(request)
|
||||||
|
|
||||||
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionResponse:
|
async def delegate_pdf_question(self, ctx: RunContext[OrchestratorDeps]) -> PdfQuestionOrchestrateResponse:
|
||||||
return await self._run_pdf_question(ctx.deps.request)
|
return await self._run_pdf_question(ctx.deps.request)
|
||||||
|
|
||||||
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionResponse:
|
async def _run_pdf_question(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
|
||||||
return await PdfQuestionAgent(self.runtime).orchestrate(request)
|
return await PdfQuestionAgent(self.runtime).orchestrate(request)
|
||||||
|
|
||||||
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
|
async def delegate_user_spec(self, ctx: RunContext[OrchestratorDeps]) -> AgentDraftWorkflowResponse:
|
||||||
@@ -165,12 +167,11 @@ class OrchestratorAgent:
|
|||||||
|
|
||||||
def _build_prompt(self, request: OrchestratorRequest) -> str:
|
def _build_prompt(self, request: OrchestratorRequest) -> str:
|
||||||
artifact_summary = self._describe_artifacts(request)
|
artifact_summary = self._describe_artifacts(request)
|
||||||
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
|
|
||||||
history = format_conversation_history(request.conversation_history)
|
history = format_conversation_history(request.conversation_history)
|
||||||
return (
|
return (
|
||||||
f"Conversation history:\n{history}\n"
|
f"Conversation history:\n{history}\n"
|
||||||
f"User message: {request.user_message}\n"
|
f"User message: {request.user_message}\n"
|
||||||
f"Files: {file_names}\n"
|
f"Files: {format_file_names(request.files)}\n"
|
||||||
f"Available artifacts:\n{artifact_summary}"
|
f"Available artifacts:\n{artifact_summary}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ from stirling.contracts import (
|
|||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
ToolOperationStep,
|
ToolOperationStep,
|
||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
|
format_file_names,
|
||||||
)
|
)
|
||||||
from stirling.logging import Pretty
|
from stirling.logging import Pretty
|
||||||
from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
|
from stirling.models import OPERATIONS, ApiModel, ParamToolModel, ToolEndpoint
|
||||||
@@ -116,7 +117,6 @@ class PdfEditParameterSelector:
|
|||||||
) -> str:
|
) -> str:
|
||||||
operation_id = operation_plan[operation_index]
|
operation_id = operation_plan[operation_index]
|
||||||
operation_list = ", ".join(operation.name for operation in operation_plan)
|
operation_list = ", ".join(operation.name for operation in operation_plan)
|
||||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
|
||||||
generated_steps_text = (
|
generated_steps_text = (
|
||||||
"\n".join(
|
"\n".join(
|
||||||
f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
|
f"- Step {step_index + 1}: {step.model_dump_json()}" for step_index, step in enumerate(generated_steps)
|
||||||
@@ -127,7 +127,7 @@ class PdfEditParameterSelector:
|
|||||||
return (
|
return (
|
||||||
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
|
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
|
||||||
f"User request: {request.user_message}\n"
|
f"User request: {request.user_message}\n"
|
||||||
f"Files: {file_names}\n"
|
f"Files: {format_file_names(request.files)}\n"
|
||||||
f"Operation plan: {operation_list}\n"
|
f"Operation plan: {operation_list}\n"
|
||||||
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
|
f"Selected operation index: {operation_index + 1} of {len(operation_plan)}\n"
|
||||||
f"Selected operation: {operation_id.name}\n"
|
f"Selected operation: {operation_id.name}\n"
|
||||||
@@ -153,7 +153,7 @@ class PdfEditAgent:
|
|||||||
return await self.handle(
|
return await self.handle(
|
||||||
PdfEditRequest(
|
PdfEditRequest(
|
||||||
user_message=request.user_message,
|
user_message=request.user_message,
|
||||||
file_names=request.file_names,
|
files=request.files,
|
||||||
conversation_history=request.conversation_history,
|
conversation_history=request.conversation_history,
|
||||||
page_text=extracted_text.files if extracted_text is not None else [],
|
page_text=extracted_text.files if extracted_text is not None else [],
|
||||||
)
|
)
|
||||||
@@ -166,7 +166,7 @@ class PdfEditAgent:
|
|||||||
async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse:
|
async def handle(self, request: PdfEditRequest, allow_need_content: bool = True) -> PdfEditResponse:
|
||||||
logger.info(
|
logger.info(
|
||||||
"[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r",
|
"[pdf-edit] handle: files=%s has_text=%s allow_need_content=%s msg=%r",
|
||||||
request.file_names,
|
[file.name for file in request.files],
|
||||||
has_page_text(request.page_text),
|
has_page_text(request.page_text),
|
||||||
allow_need_content,
|
allow_need_content,
|
||||||
request.user_message,
|
request.user_message,
|
||||||
@@ -225,11 +225,10 @@ class PdfEditAgent:
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
|
def _build_selection_prompt(self, request: PdfEditRequest) -> str:
|
||||||
file_names = ", ".join(request.file_names) if request.file_names else "No file names were provided."
|
|
||||||
return (
|
return (
|
||||||
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
|
f"Conversation history:\n{format_conversation_history(request.conversation_history)}\n"
|
||||||
f"User request: {request.user_message}\n"
|
f"User request: {request.user_message}\n"
|
||||||
f"Files: {file_names}\n"
|
f"Files: {format_file_names(request.files)}\n"
|
||||||
f"Supported operations: {self._supported_operations_prompt()}\n"
|
f"Supported operations: {self._supported_operations_prompt()}\n"
|
||||||
f"Extracted page text:\n{format_page_text(request.page_text)}"
|
f"Extracted page text:\n{format_page_text(request.page_text)}"
|
||||||
)
|
)
|
||||||
@@ -243,8 +242,7 @@ class PdfEditAgent:
|
|||||||
request: PdfEditRequest,
|
request: PdfEditRequest,
|
||||||
) -> NeedContentResponse:
|
) -> NeedContentResponse:
|
||||||
files = selection.files or [
|
files = selection.files or [
|
||||||
NeedContentFileRequest(file_name=file_name, content_types=[PdfContentType.PAGE_TEXT])
|
NeedContentFileRequest(file=file, content_types=[PdfContentType.PAGE_TEXT]) for file in request.files
|
||||||
for file_name in request.file_names
|
|
||||||
]
|
]
|
||||||
return NeedContentResponse(
|
return NeedContentResponse(
|
||||||
resume_with=SupportedCapability.PDF_EDIT,
|
resume_with=SupportedCapability.PDF_EDIT,
|
||||||
|
|||||||
@@ -1,32 +1,67 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
from pydantic_ai import Agent
|
from pydantic_ai import Agent
|
||||||
from pydantic_ai.output import NativeOutput
|
from pydantic_ai.output import NativeOutput
|
||||||
|
|
||||||
from stirling.agents._page_text import (
|
|
||||||
format_page_text,
|
|
||||||
get_extracted_text_artifact,
|
|
||||||
has_page_text,
|
|
||||||
)
|
|
||||||
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
|
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
|
AiFile,
|
||||||
EditPlanResponse,
|
EditPlanResponse,
|
||||||
NeedContentFileRequest,
|
NeedIngestResponse,
|
||||||
NeedContentResponse,
|
|
||||||
OrchestratorRequest,
|
OrchestratorRequest,
|
||||||
PdfContentType,
|
PdfContentType,
|
||||||
PdfQuestionAnswerResponse,
|
PdfQuestionAnswerResponse,
|
||||||
PdfQuestionNotFoundResponse,
|
PdfQuestionNotFoundResponse,
|
||||||
|
PdfQuestionOrchestrateResponse,
|
||||||
PdfQuestionRequest,
|
PdfQuestionRequest,
|
||||||
PdfQuestionResponse,
|
PdfQuestionResponse,
|
||||||
|
PdfQuestionTerminalResponse,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
ToolOperationStep,
|
ToolOperationStep,
|
||||||
Verdict,
|
Verdict,
|
||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
|
format_file_names,
|
||||||
)
|
)
|
||||||
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
|
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
|
||||||
|
from stirling.rag import RagCapability
|
||||||
from stirling.services import AppRuntime
|
from stirling.services import AppRuntime
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
PDF_QUESTION_SYSTEM_PROMPT = (
|
||||||
|
"You answer questions about PDF documents by retrieving relevant content with the "
|
||||||
|
"search_knowledge tool. Use it before answering. Do not guess or use outside knowledge.\n"
|
||||||
|
"\n"
|
||||||
|
"The search_knowledge tool has a finite call budget per run. When it is no longer "
|
||||||
|
"available, answer from what you have already retrieved.\n"
|
||||||
|
"\n"
|
||||||
|
"Guidelines:\n"
|
||||||
|
"- Make targeted search_knowledge calls. Typically one or two is enough.\n"
|
||||||
|
"- Answer from the retrieved text. If the retrieved content doesn't support a confident "
|
||||||
|
"answer, return not_found.\n"
|
||||||
|
"- For questions that would require reading the entire document end-to-end (e.g. "
|
||||||
|
"'what's the shortest chapter', 'how many X are there'), return not_found.\n"
|
||||||
|
"- Include a short list of evidence snippets (with page numbers where available) drawn "
|
||||||
|
"from what search_knowledge returned.\n"
|
||||||
|
"\n"
|
||||||
|
"Writing the not_found reason:\n"
|
||||||
|
"- The reason is shown directly to the end user, so write it in plain, friendly "
|
||||||
|
"language. One or two short sentences.\n"
|
||||||
|
"- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
|
||||||
|
"'search_knowledge', or other implementation details.\n"
|
||||||
|
"- Be honest about the actual limitation. For questions that require full-document "
|
||||||
|
"analysis (shortest chapter, word counts, etc.), explain that the document is too "
|
||||||
|
"long to analyse end-to-end: you can only look up specific passages, and that's "
|
||||||
|
"not enough to compare every part of the document against every other.\n"
|
||||||
|
"- For questions where the answer just isn't in the document, say so directly: "
|
||||||
|
"'I couldn't find that information in the document.'\n"
|
||||||
|
"- Do not make it sound like you're choosing not to answer. Be clear that it's "
|
||||||
|
"a genuine constraint."
|
||||||
|
)
|
||||||
|
|
||||||
_MATH_SYNTH_SYSTEM_PROMPT = (
|
_MATH_SYNTH_SYSTEM_PROMPT = (
|
||||||
"You are given a math-audit Verdict (structured JSON) and the user's "
|
"You are given a math-audit Verdict (structured JSON) and the user's "
|
||||||
"original question. Answer the question in plain prose using only "
|
"original question. Answer the question in plain prose using only "
|
||||||
@@ -41,26 +76,6 @@ _MATH_SYNTH_SYSTEM_PROMPT = (
|
|||||||
class PdfQuestionAgent:
|
class PdfQuestionAgent:
|
||||||
def __init__(self, runtime: AppRuntime) -> None:
|
def __init__(self, runtime: AppRuntime) -> None:
|
||||||
self.runtime = runtime
|
self.runtime = runtime
|
||||||
rag = runtime.rag_capability
|
|
||||||
self.agent = Agent(
|
|
||||||
model=runtime.smart_model,
|
|
||||||
output_type=NativeOutput(
|
|
||||||
[
|
|
||||||
PdfQuestionAnswerResponse,
|
|
||||||
PdfQuestionNotFoundResponse,
|
|
||||||
]
|
|
||||||
),
|
|
||||||
system_prompt=(
|
|
||||||
"Answer questions about PDFs using only the extracted page text provided in the prompt. "
|
|
||||||
"Do not guess or use outside knowledge. "
|
|
||||||
"If the answer is not supported by the provided text, return not_found. "
|
|
||||||
"When answering, include a short list of evidence snippets with their page numbers. "
|
|
||||||
"Reply in the SAME LANGUAGE as the question."
|
|
||||||
),
|
|
||||||
instructions=rag.instructions,
|
|
||||||
toolsets=[rag.toolset],
|
|
||||||
model_settings=runtime.smart_model_settings,
|
|
||||||
)
|
|
||||||
self._math_synth_agent: Agent[None, str] = Agent(
|
self._math_synth_agent: Agent[None, str] = Agent(
|
||||||
model=runtime.fast_model,
|
model=runtime.fast_model,
|
||||||
output_type=str,
|
output_type=str,
|
||||||
@@ -70,30 +85,31 @@ class PdfQuestionAgent:
|
|||||||
self._math_intent_classifier = MathIntentClassifier(runtime)
|
self._math_intent_classifier = MathIntentClassifier(runtime)
|
||||||
|
|
||||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||||
if not has_page_text(request.page_text):
|
logger.info(
|
||||||
return NeedContentResponse(
|
"[pdf-question] handle: files=%s question=%r",
|
||||||
|
[file.name for file in request.files],
|
||||||
|
request.question,
|
||||||
|
)
|
||||||
|
missing = await self._find_missing_files(request.files)
|
||||||
|
if missing:
|
||||||
|
logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
|
||||||
|
return NeedIngestResponse(
|
||||||
resume_with=SupportedCapability.PDF_QUESTION,
|
resume_with=SupportedCapability.PDF_QUESTION,
|
||||||
reason="No extracted PDF page text was provided, so the question cannot be answered yet.",
|
reason="Some files have not been ingested into RAG yet.",
|
||||||
files=[
|
files_to_ingest=missing,
|
||||||
NeedContentFileRequest(
|
content_types=[PdfContentType.PAGE_TEXT],
|
||||||
file_name=file_name,
|
|
||||||
content_types=[PdfContentType.PAGE_TEXT],
|
|
||||||
)
|
|
||||||
for file_name in request.file_names
|
|
||||||
],
|
|
||||||
max_pages=self.runtime.settings.max_pages,
|
|
||||||
max_characters=self.runtime.settings.max_characters,
|
|
||||||
)
|
)
|
||||||
return await self._run_answer_agent(request)
|
return await self._run_answer_agent(request)
|
||||||
|
|
||||||
async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionResponse:
|
async def orchestrate(self, request: OrchestratorRequest) -> PdfQuestionOrchestrateResponse:
|
||||||
"""Entry point for the orchestrator delegate.
|
"""Entry point for the orchestrator delegate.
|
||||||
|
|
||||||
Decides math intent locally via a small classifier LLM (language-agnostic).
|
Decides math intent locally via a small classifier LLM (language-agnostic).
|
||||||
On a math first turn, embeds an :class:`EditPlanResponse` in the answer
|
On a math first turn, returns an :class:`EditPlanResponse` (``outcome=PLAN``)
|
||||||
response; on the resume turn, digests the captured :class:`Verdict` into
|
with ``resume_with=PDF_QUESTION`` so the caller runs the math specialist
|
||||||
a localised prose answer. Non-math first turns fall through to the
|
and re-invokes the orchestrator. On the resume turn, the captured
|
||||||
text-grounded :meth:`handle` pipeline.
|
:class:`Verdict` is digested into a localised prose answer. Non-math
|
||||||
|
first turns fall through to the text-grounded :meth:`handle` pipeline.
|
||||||
"""
|
"""
|
||||||
verdict = extract_math_verdict(request)
|
verdict = extract_math_verdict(request)
|
||||||
if verdict is not None:
|
if verdict is not None:
|
||||||
@@ -104,36 +120,53 @@ class PdfQuestionAgent:
|
|||||||
return PdfQuestionAnswerResponse(answer=answer, evidence=[])
|
return PdfQuestionAnswerResponse(answer=answer, evidence=[])
|
||||||
|
|
||||||
if await self._math_intent_classifier.classify(request.user_message):
|
if await self._math_intent_classifier.classify(request.user_message):
|
||||||
# First turn — ask the caller to run the math specialist and come back.
|
# First turn — emit a one-step plan calling the math specialist,
|
||||||
# The plan rides on the answer response as a nullable member; ``answer``
|
# with resume_with set so the caller comes back with the verdict
|
||||||
# is empty on this turn and the caller resumes once the plan is run.
|
# in artifacts (handled by the resume branch above).
|
||||||
return PdfQuestionAnswerResponse(
|
return EditPlanResponse(
|
||||||
answer="",
|
summary="",
|
||||||
evidence=[],
|
steps=[
|
||||||
edit_plan=EditPlanResponse(
|
ToolOperationStep(
|
||||||
summary="",
|
tool=AgentToolId.MATH_AUDITOR_AGENT,
|
||||||
steps=[
|
parameters=MathAuditorAgentParams(),
|
||||||
ToolOperationStep(
|
)
|
||||||
tool=AgentToolId.MATH_AUDITOR_AGENT,
|
],
|
||||||
parameters=MathAuditorAgentParams(),
|
resume_with=SupportedCapability.PDF_QUESTION,
|
||||||
)
|
|
||||||
],
|
|
||||||
resume_with=SupportedCapability.PDF_QUESTION,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
extracted_text = get_extracted_text_artifact(request)
|
|
||||||
return await self.handle(
|
return await self.handle(
|
||||||
PdfQuestionRequest(
|
PdfQuestionRequest(
|
||||||
question=request.user_message,
|
question=request.user_message,
|
||||||
file_names=request.file_names,
|
files=request.files,
|
||||||
page_text=extracted_text.files if extracted_text is not None else [],
|
|
||||||
conversation_history=request.conversation_history,
|
conversation_history=request.conversation_history,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
|
||||||
result = await self.agent.run(self._build_prompt(request))
|
missing: list[AiFile] = []
|
||||||
|
for file in files:
|
||||||
|
if not await self.runtime.rag_service.has_collection(file.id):
|
||||||
|
missing.append(file)
|
||||||
|
return missing
|
||||||
|
|
||||||
|
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
|
||||||
|
rag = RagCapability(
|
||||||
|
rag_service=self.runtime.rag_service,
|
||||||
|
collections=[file.id for file in request.files],
|
||||||
|
top_k=self.runtime.settings.rag_default_top_k,
|
||||||
|
max_searches=self.runtime.settings.rag_max_searches,
|
||||||
|
)
|
||||||
|
agent = Agent(
|
||||||
|
model=self.runtime.smart_model,
|
||||||
|
output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
|
||||||
|
system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
|
||||||
|
instructions=rag.instructions,
|
||||||
|
toolsets=[rag.toolset],
|
||||||
|
model_settings=self.runtime.smart_model_settings,
|
||||||
|
)
|
||||||
|
prompt = self._build_prompt(request)
|
||||||
|
logger.debug("[pdf-question] prompt:\n%s", prompt)
|
||||||
|
result = await agent.run(prompt)
|
||||||
return result.output
|
return result.output
|
||||||
|
|
||||||
async def _synthesise_math_answer(self, user_message: str, verdict: Verdict) -> str:
|
async def _synthesise_math_answer(self, user_message: str, verdict: Verdict) -> str:
|
||||||
@@ -146,12 +179,10 @@ class PdfQuestionAgent:
|
|||||||
return result.output
|
return result.output
|
||||||
|
|
||||||
def _build_prompt(self, request: PdfQuestionRequest) -> str:
|
def _build_prompt(self, request: PdfQuestionRequest) -> str:
|
||||||
file_names = ", ".join(request.file_names) if request.file_names else "Unknown files"
|
|
||||||
pages = format_page_text(request.page_text, empty="")
|
|
||||||
history = format_conversation_history(request.conversation_history)
|
history = format_conversation_history(request.conversation_history)
|
||||||
return (
|
return (
|
||||||
f"Conversation history:\n{history}\n"
|
f"Conversation history:\n{history}\n"
|
||||||
f"Files: {file_names}\n"
|
f"Files: {format_file_names(request.files)}\n"
|
||||||
f"Question: {request.question}\n"
|
f"Question: {request.question}\n"
|
||||||
f"Extracted page text:\n{pages}"
|
"Use search_knowledge to retrieve the relevant content, then answer."
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -7,7 +7,13 @@ from fastapi import Depends, FastAPI
|
|||||||
from pydantic_ai import Agent
|
from pydantic_ai import Agent
|
||||||
from pydantic_ai.models.instrumented import InstrumentationSettings
|
from pydantic_ai.models.instrumented import InstrumentationSettings
|
||||||
|
|
||||||
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
|
from stirling.agents import (
|
||||||
|
ExecutionPlanningAgent,
|
||||||
|
OrchestratorAgent,
|
||||||
|
PdfEditAgent,
|
||||||
|
PdfQuestionAgent,
|
||||||
|
UserSpecAgent,
|
||||||
|
)
|
||||||
from stirling.agents.ledger import MathAuditorAgent
|
from stirling.agents.ledger import MathAuditorAgent
|
||||||
from stirling.agents.pdf_comment import PdfCommentAgent
|
from stirling.agents.pdf_comment import PdfCommentAgent
|
||||||
from stirling.api.middleware import UserIdMiddleware
|
from stirling.api.middleware import UserIdMiddleware
|
||||||
@@ -51,6 +57,7 @@ async def lifespan(fast_api: FastAPI):
|
|||||||
if tracer_provider:
|
if tracer_provider:
|
||||||
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
|
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
|
||||||
yield
|
yield
|
||||||
|
await runtime.rag_service.close()
|
||||||
if tracer_provider:
|
if tracer_provider:
|
||||||
tracer_provider.shutdown()
|
tracer_provider.shutdown()
|
||||||
|
|
||||||
|
|||||||
@@ -2,7 +2,13 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from fastapi import Request
|
from fastapi import Request
|
||||||
|
|
||||||
from stirling.agents import ExecutionPlanningAgent, OrchestratorAgent, PdfEditAgent, PdfQuestionAgent, UserSpecAgent
|
from stirling.agents import (
|
||||||
|
ExecutionPlanningAgent,
|
||||||
|
OrchestratorAgent,
|
||||||
|
PdfEditAgent,
|
||||||
|
PdfQuestionAgent,
|
||||||
|
UserSpecAgent,
|
||||||
|
)
|
||||||
from stirling.agents.ledger import MathAuditorAgent
|
from stirling.agents.ledger import MathAuditorAgent
|
||||||
from stirling.agents.pdf_comment import PdfCommentAgent
|
from stirling.agents.pdf_comment import PdfCommentAgent
|
||||||
from stirling.rag import RagService
|
from stirling.rag import RagService
|
||||||
@@ -37,10 +43,6 @@ def get_rag_service(request: Request) -> RagService:
|
|||||||
return request.app.state.runtime.rag_service
|
return request.app.state.runtime.rag_service
|
||||||
|
|
||||||
|
|
||||||
def get_rag_embedding_model(request: Request) -> str:
|
|
||||||
return request.app.state.runtime.settings.rag_embedding_model
|
|
||||||
|
|
||||||
|
|
||||||
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
|
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
|
||||||
return request.app.state.math_auditor_agent
|
return request.app.state.math_auditor_agent
|
||||||
|
|
||||||
|
|||||||
@@ -4,75 +4,60 @@ from typing import Annotated
|
|||||||
|
|
||||||
from fastapi import APIRouter, Depends
|
from fastapi import APIRouter, Depends
|
||||||
|
|
||||||
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
|
from stirling.api.dependencies import get_rag_service
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
RagCollectionsResponse,
|
DeleteDocumentResponse,
|
||||||
RagDeleteCollectionResponse,
|
IngestDocumentRequest,
|
||||||
RagIndexRequest,
|
IngestDocumentResponse,
|
||||||
RagIndexResponse,
|
PdfContentType,
|
||||||
RagSearchRequest,
|
|
||||||
RagSearchResponse,
|
|
||||||
RagSearchResultItem,
|
|
||||||
RagStatusResponse,
|
|
||||||
)
|
)
|
||||||
from stirling.rag import RagService
|
from stirling.models import FileId
|
||||||
|
from stirling.rag import Document, RagService
|
||||||
|
|
||||||
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
|
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
|
||||||
|
|
||||||
|
|
||||||
@router.get("/status", response_model=RagStatusResponse)
|
@router.post("/documents", response_model=IngestDocumentResponse)
|
||||||
async def rag_status(
|
async def ingest_document(
|
||||||
|
request: IngestDocumentRequest,
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
rag: Annotated[RagService, Depends(get_rag_service)],
|
||||||
embedding_model: Annotated[str, Depends(get_rag_embedding_model)],
|
) -> IngestDocumentResponse:
|
||||||
) -> RagStatusResponse:
|
"""Replace-ingest a document's content under ``document_id``.
|
||||||
collections = await rag.list_collections()
|
|
||||||
return RagStatusResponse(embedding_model=embedding_model, collections=collections)
|
Any previously-stored content for this document is removed and the
|
||||||
|
provided content replaces it wholesale. All pages are chunked up front
|
||||||
|
and then embedded in a single batched call so large documents (e.g. a
|
||||||
|
500-page book) don't fan out into hundreds of embedding requests.
|
||||||
|
"""
|
||||||
|
await rag.delete_collection(request.document_id)
|
||||||
|
|
||||||
|
chunks: list[Document] = []
|
||||||
|
if request.page_text:
|
||||||
|
for page in request.page_text:
|
||||||
|
if not page.text.strip():
|
||||||
|
continue
|
||||||
|
chunks.extend(
|
||||||
|
rag.chunk_text(
|
||||||
|
text=page.text,
|
||||||
|
source=f"{request.source}:page:{page.page_number}",
|
||||||
|
base_metadata={
|
||||||
|
"page_number": str(page.page_number),
|
||||||
|
"content_type": PdfContentType.PAGE_TEXT.value,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
indexed = await rag.index_documents(request.document_id, chunks) if chunks else 0
|
||||||
|
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=indexed)
|
||||||
|
|
||||||
|
|
||||||
@router.post("/index", response_model=RagIndexResponse)
|
@router.delete("/documents/{document_id}", response_model=DeleteDocumentResponse)
|
||||||
async def rag_index(
|
async def delete_document(
|
||||||
request: RagIndexRequest,
|
document_id: FileId,
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
rag: Annotated[RagService, Depends(get_rag_service)],
|
||||||
) -> RagIndexResponse:
|
) -> DeleteDocumentResponse:
|
||||||
count = await rag.index_text(
|
"""Remove a document's content from RAG. Idempotent."""
|
||||||
collection=request.collection,
|
existed = await rag.has_collection(document_id)
|
||||||
text=request.text,
|
if existed:
|
||||||
source=request.source,
|
await rag.delete_collection(document_id)
|
||||||
metadata=request.metadata,
|
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
|
||||||
)
|
|
||||||
return RagIndexResponse(collection=request.collection, chunks_indexed=count)
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/search", response_model=RagSearchResponse)
|
|
||||||
async def rag_search(
|
|
||||||
request: RagSearchRequest,
|
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
|
||||||
) -> RagSearchResponse:
|
|
||||||
results = await rag.search(query=request.query, collection=request.collection, top_k=request.top_k)
|
|
||||||
items = [
|
|
||||||
RagSearchResultItem(
|
|
||||||
text=r.document.text,
|
|
||||||
source=r.document.metadata.get("source", ""),
|
|
||||||
chunk_id=r.document.metadata.get("chunk_index", ""),
|
|
||||||
score=r.score,
|
|
||||||
)
|
|
||||||
for r in results
|
|
||||||
]
|
|
||||||
return RagSearchResponse(query=request.query, results=items)
|
|
||||||
|
|
||||||
|
|
||||||
@router.get("/collections", response_model=RagCollectionsResponse)
|
|
||||||
async def rag_collections(
|
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
|
||||||
) -> RagCollectionsResponse:
|
|
||||||
collections = await rag.list_collections()
|
|
||||||
return RagCollectionsResponse(collections=collections)
|
|
||||||
|
|
||||||
|
|
||||||
@router.delete("/collections/{name}", response_model=RagDeleteCollectionResponse)
|
|
||||||
async def rag_delete_collection(
|
|
||||||
name: str,
|
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
|
||||||
) -> RagDeleteCollectionResponse:
|
|
||||||
await rag.delete_collection(name)
|
|
||||||
return RagDeleteCollectionResponse(status="deleted", collection=name)
|
|
||||||
|
|||||||
@@ -36,6 +36,7 @@ class AppSettings(BaseSettings):
|
|||||||
rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
|
rag_chunk_size: int = Field(validation_alias="STIRLING_RAG_CHUNK_SIZE")
|
||||||
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
|
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
|
||||||
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
|
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
|
||||||
|
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
|
||||||
|
|
||||||
max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
|
max_pages: int = Field(validation_alias="STIRLING_MAX_PAGES")
|
||||||
max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS")
|
max_characters: int = Field(validation_alias="STIRLING_MAX_CHARACTERS")
|
||||||
|
|||||||
@@ -10,12 +10,14 @@ from .agent_drafts import (
|
|||||||
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
|
from .agent_specs import AgentSpec, AgentSpecStep, AiToolAgentStep
|
||||||
from .comments import CommentSpec
|
from .comments import CommentSpec
|
||||||
from .common import (
|
from .common import (
|
||||||
|
AiFile,
|
||||||
ArtifactKind,
|
ArtifactKind,
|
||||||
ConversationMessage,
|
ConversationMessage,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
MathAuditorToolReportArtifact,
|
MathAuditorToolReportArtifact,
|
||||||
NeedContentFileRequest,
|
NeedContentFileRequest,
|
||||||
NeedContentResponse,
|
NeedContentResponse,
|
||||||
|
NeedIngestResponse,
|
||||||
PdfContentType,
|
PdfContentType,
|
||||||
PdfTextSelection,
|
PdfTextSelection,
|
||||||
StepKind,
|
StepKind,
|
||||||
@@ -24,6 +26,7 @@ from .common import (
|
|||||||
ToolReportArtifact,
|
ToolReportArtifact,
|
||||||
WorkflowOutcome,
|
WorkflowOutcome,
|
||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
|
format_file_names,
|
||||||
)
|
)
|
||||||
from .execution import (
|
from .execution import (
|
||||||
AgentExecutionRequest,
|
AgentExecutionRequest,
|
||||||
@@ -71,24 +74,20 @@ from .pdf_edit import (
|
|||||||
from .pdf_questions import (
|
from .pdf_questions import (
|
||||||
PdfQuestionAnswerResponse,
|
PdfQuestionAnswerResponse,
|
||||||
PdfQuestionNotFoundResponse,
|
PdfQuestionNotFoundResponse,
|
||||||
|
PdfQuestionOrchestrateResponse,
|
||||||
PdfQuestionRequest,
|
PdfQuestionRequest,
|
||||||
PdfQuestionResponse,
|
PdfQuestionResponse,
|
||||||
PdfQuestionTerminalResponse,
|
PdfQuestionTerminalResponse,
|
||||||
)
|
)
|
||||||
from .rag import (
|
from .rag import (
|
||||||
MAX_INDEX_TEXT_LENGTH,
|
DeleteDocumentResponse,
|
||||||
RagCollectionsResponse,
|
IngestDocumentRequest,
|
||||||
RagDeleteCollectionResponse,
|
IngestDocumentResponse,
|
||||||
RagIndexRequest,
|
IngestedPageText,
|
||||||
RagIndexResponse,
|
|
||||||
RagSearchRequest,
|
|
||||||
RagSearchResponse,
|
|
||||||
RagSearchResultItem,
|
|
||||||
RagStatusResponse,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"MAX_INDEX_TEXT_LENGTH",
|
"AiFile",
|
||||||
"AgentDraft",
|
"AgentDraft",
|
||||||
"AgentDraftRequest",
|
"AgentDraftRequest",
|
||||||
"AgentDraftResponse",
|
"AgentDraftResponse",
|
||||||
@@ -105,6 +104,7 @@ __all__ = [
|
|||||||
"CommentSpec",
|
"CommentSpec",
|
||||||
"CompletedExecutionAction",
|
"CompletedExecutionAction",
|
||||||
"ConversationMessage",
|
"ConversationMessage",
|
||||||
|
"DeleteDocumentResponse",
|
||||||
"Discrepancy",
|
"Discrepancy",
|
||||||
"DiscrepancyKind",
|
"DiscrepancyKind",
|
||||||
"EditCannotDoResponse",
|
"EditCannotDoResponse",
|
||||||
@@ -119,10 +119,15 @@ __all__ = [
|
|||||||
"FolioManifest",
|
"FolioManifest",
|
||||||
"FolioType",
|
"FolioType",
|
||||||
"format_conversation_history",
|
"format_conversation_history",
|
||||||
|
"format_file_names",
|
||||||
"HealthResponse",
|
"HealthResponse",
|
||||||
|
"IngestDocumentRequest",
|
||||||
|
"IngestDocumentResponse",
|
||||||
|
"IngestedPageText",
|
||||||
"MathAuditorToolReportArtifact",
|
"MathAuditorToolReportArtifact",
|
||||||
"NeedContentFileRequest",
|
"NeedContentFileRequest",
|
||||||
"NeedContentResponse",
|
"NeedContentResponse",
|
||||||
|
"NeedIngestResponse",
|
||||||
"NextExecutionAction",
|
"NextExecutionAction",
|
||||||
"OrchestratorRequest",
|
"OrchestratorRequest",
|
||||||
"OrchestratorResponse",
|
"OrchestratorResponse",
|
||||||
@@ -136,18 +141,11 @@ __all__ = [
|
|||||||
"PdfEditTerminalResponse",
|
"PdfEditTerminalResponse",
|
||||||
"PdfQuestionAnswerResponse",
|
"PdfQuestionAnswerResponse",
|
||||||
"PdfQuestionNotFoundResponse",
|
"PdfQuestionNotFoundResponse",
|
||||||
|
"PdfQuestionOrchestrateResponse",
|
||||||
"PdfQuestionRequest",
|
"PdfQuestionRequest",
|
||||||
"PdfQuestionResponse",
|
"PdfQuestionResponse",
|
||||||
"PdfQuestionTerminalResponse",
|
"PdfQuestionTerminalResponse",
|
||||||
"PdfTextSelection",
|
"PdfTextSelection",
|
||||||
"RagCollectionsResponse",
|
|
||||||
"RagDeleteCollectionResponse",
|
|
||||||
"RagIndexRequest",
|
|
||||||
"RagIndexResponse",
|
|
||||||
"RagSearchRequest",
|
|
||||||
"RagSearchResponse",
|
|
||||||
"RagSearchResultItem",
|
|
||||||
"RagStatusResponse",
|
|
||||||
"Requisition",
|
"Requisition",
|
||||||
"Severity",
|
"Severity",
|
||||||
"StepKind",
|
"StepKind",
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ from typing import Literal, assert_never
|
|||||||
from pydantic import Field, model_validator
|
from pydantic import Field, model_validator
|
||||||
|
|
||||||
from stirling.contracts.ledger import Verdict
|
from stirling.contracts.ledger import Verdict
|
||||||
from stirling.models import OPERATIONS, ApiModel, ToolEndpoint
|
from stirling.models import OPERATIONS, ApiModel, FileId, ToolEndpoint
|
||||||
from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
|
from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
|
||||||
|
|
||||||
|
|
||||||
@@ -50,6 +50,7 @@ class WorkflowOutcome(StrEnum):
|
|||||||
|
|
||||||
ANSWER = "answer"
|
ANSWER = "answer"
|
||||||
NEED_CONTENT = "need_content"
|
NEED_CONTENT = "need_content"
|
||||||
|
NEED_INGEST = "need_ingest"
|
||||||
NOT_FOUND = "not_found"
|
NOT_FOUND = "not_found"
|
||||||
PLAN = "plan"
|
PLAN = "plan"
|
||||||
NEED_CLARIFICATION = "need_clarification"
|
NEED_CLARIFICATION = "need_clarification"
|
||||||
@@ -94,12 +95,30 @@ class ConversationMessage(ApiModel):
|
|||||||
content: str
|
content: str
|
||||||
|
|
||||||
|
|
||||||
|
class AiFile(ApiModel):
|
||||||
|
"""A file the user has supplied, identified by both a stable id and a display name.
|
||||||
|
|
||||||
|
The id is opaque to the engine: Java generates it (content hash, file path, UUID, etc.)
|
||||||
|
and the engine uses it as the RAG collection key for any agent that indexes content.
|
||||||
|
The name is used in user-facing prompts and responses.
|
||||||
|
"""
|
||||||
|
|
||||||
|
id: FileId = Field(min_length=1)
|
||||||
|
name: str = Field(min_length=1)
|
||||||
|
|
||||||
|
|
||||||
def format_conversation_history(conversation_history: list[ConversationMessage]) -> str:
|
def format_conversation_history(conversation_history: list[ConversationMessage]) -> str:
|
||||||
if not conversation_history:
|
if not conversation_history:
|
||||||
return "None"
|
return "None"
|
||||||
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
|
return "\n".join(f"- {message.role}: {message.content}" for message in conversation_history)
|
||||||
|
|
||||||
|
|
||||||
|
def format_file_names(files: list[AiFile]) -> str:
|
||||||
|
if not files:
|
||||||
|
return "No file names were provided."
|
||||||
|
return ", ".join(file.name for file in files)
|
||||||
|
|
||||||
|
|
||||||
class PdfTextSelection(ApiModel):
|
class PdfTextSelection(ApiModel):
|
||||||
page_number: int | None = None
|
page_number: int | None = None
|
||||||
text: str
|
text: str
|
||||||
@@ -111,7 +130,7 @@ class ExtractedFileText(ApiModel):
|
|||||||
|
|
||||||
|
|
||||||
class NeedContentFileRequest(ApiModel):
|
class NeedContentFileRequest(ApiModel):
|
||||||
file_name: str
|
file: AiFile
|
||||||
page_numbers: list[int] = Field(default_factory=list)
|
page_numbers: list[int] = Field(default_factory=list)
|
||||||
content_types: list[PdfContentType]
|
content_types: list[PdfContentType]
|
||||||
|
|
||||||
@@ -146,6 +165,20 @@ class MathAuditorToolReportArtifact(ApiModel):
|
|||||||
ToolReportArtifact = MathAuditorToolReportArtifact
|
ToolReportArtifact = MathAuditorToolReportArtifact
|
||||||
|
|
||||||
|
|
||||||
|
class NeedIngestResponse(ApiModel):
|
||||||
|
"""Signal that the listed files must be ingested into RAG before the agent can continue.
|
||||||
|
|
||||||
|
Java's handling: for each file, extract the requested content types, POST to
|
||||||
|
``/api/v1/rag/documents`` keyed by ``file.id``, then retry the original request.
|
||||||
|
"""
|
||||||
|
|
||||||
|
outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST
|
||||||
|
resume_with: SupportedCapability
|
||||||
|
reason: str
|
||||||
|
files_to_ingest: list[AiFile]
|
||||||
|
content_types: list[PdfContentType] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
class ToolOperationStep(ApiModel):
|
class ToolOperationStep(ApiModel):
|
||||||
kind: Literal[StepKind.TOOL] = StepKind.TOOL
|
kind: Literal[StepKind.TOOL] = StepKind.TOOL
|
||||||
tool: AnyToolId
|
tool: AnyToolId
|
||||||
|
|||||||
@@ -8,10 +8,12 @@ from stirling.models import ApiModel
|
|||||||
|
|
||||||
from .agent_drafts import AgentDraftResponse
|
from .agent_drafts import AgentDraftResponse
|
||||||
from .common import (
|
from .common import (
|
||||||
|
AiFile,
|
||||||
ArtifactKind,
|
ArtifactKind,
|
||||||
ConversationMessage,
|
ConversationMessage,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
NeedContentResponse,
|
NeedContentResponse,
|
||||||
|
NeedIngestResponse,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
ToolReportArtifact,
|
ToolReportArtifact,
|
||||||
WorkflowOutcome,
|
WorkflowOutcome,
|
||||||
@@ -31,7 +33,7 @@ WorkflowArtifact = Annotated[ExtractedTextArtifact | ToolReportArtifact, Field(d
|
|||||||
|
|
||||||
class OrchestratorRequest(ApiModel):
|
class OrchestratorRequest(ApiModel):
|
||||||
user_message: str
|
user_message: str
|
||||||
file_names: list[str]
|
files: list[AiFile] = Field(default_factory=list)
|
||||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||||
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
|
artifacts: list[WorkflowArtifact] = Field(default_factory=list)
|
||||||
resume_with: SupportedCapability | None = None
|
resume_with: SupportedCapability | None = None
|
||||||
@@ -47,6 +49,7 @@ type OrchestratorResponse = Annotated[
|
|||||||
PdfEditTerminalResponse
|
PdfEditTerminalResponse
|
||||||
| PdfQuestionTerminalResponse
|
| PdfQuestionTerminalResponse
|
||||||
| NeedContentResponse
|
| NeedContentResponse
|
||||||
|
| NeedIngestResponse
|
||||||
| AgentDraftResponse
|
| AgentDraftResponse
|
||||||
| NextExecutionAction
|
| NextExecutionAction
|
||||||
| UnsupportedCapabilityResponse,
|
| UnsupportedCapabilityResponse,
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ from pydantic import Field
|
|||||||
from stirling.models import ApiModel
|
from stirling.models import ApiModel
|
||||||
|
|
||||||
from .common import (
|
from .common import (
|
||||||
|
AiFile,
|
||||||
ConversationMessage,
|
ConversationMessage,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
NeedContentResponse,
|
NeedContentResponse,
|
||||||
@@ -18,7 +19,7 @@ from .common import (
|
|||||||
|
|
||||||
class PdfEditRequest(ApiModel):
|
class PdfEditRequest(ApiModel):
|
||||||
user_message: str
|
user_message: str
|
||||||
file_names: list[str] = Field(default_factory=list)
|
files: list[AiFile] = Field(default_factory=list)
|
||||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||||
page_text: list[ExtractedFileText] = Field(default_factory=list)
|
page_text: list[ExtractedFileText] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|||||||
@@ -7,9 +7,10 @@ from pydantic import Field
|
|||||||
from stirling.models import ApiModel
|
from stirling.models import ApiModel
|
||||||
|
|
||||||
from .common import (
|
from .common import (
|
||||||
|
AiFile,
|
||||||
ConversationMessage,
|
ConversationMessage,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
NeedContentResponse,
|
NeedIngestResponse,
|
||||||
WorkflowOutcome,
|
WorkflowOutcome,
|
||||||
)
|
)
|
||||||
from .pdf_edit import EditPlanResponse
|
from .pdf_edit import EditPlanResponse
|
||||||
@@ -17,8 +18,7 @@ from .pdf_edit import EditPlanResponse
|
|||||||
|
|
||||||
class PdfQuestionRequest(ApiModel):
|
class PdfQuestionRequest(ApiModel):
|
||||||
question: str
|
question: str
|
||||||
page_text: list[ExtractedFileText] = Field(default_factory=list)
|
files: list[AiFile] = Field(default_factory=list)
|
||||||
file_names: list[str]
|
|
||||||
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
conversation_history: list[ConversationMessage] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
@@ -26,15 +26,6 @@ class PdfQuestionAnswerResponse(ApiModel):
|
|||||||
outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
|
outcome: Literal[WorkflowOutcome.ANSWER] = WorkflowOutcome.ANSWER
|
||||||
answer: str
|
answer: str
|
||||||
evidence: list[ExtractedFileText] = Field(default_factory=list)
|
evidence: list[ExtractedFileText] = Field(default_factory=list)
|
||||||
edit_plan: EditPlanResponse | None = Field(
|
|
||||||
default=None,
|
|
||||||
description=(
|
|
||||||
"Optional plan the caller must run before the answer is final. When"
|
|
||||||
" populated, ``answer`` is empty on this turn — the caller executes"
|
|
||||||
" the plan and re-invokes the orchestrator with ``resume_with`` set"
|
|
||||||
" to PDF_QUESTION; the real answer arrives on the resume turn."
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class PdfQuestionNotFoundResponse(ApiModel):
|
class PdfQuestionNotFoundResponse(ApiModel):
|
||||||
@@ -44,6 +35,14 @@ class PdfQuestionNotFoundResponse(ApiModel):
|
|||||||
|
|
||||||
type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse
|
type PdfQuestionTerminalResponse = PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse
|
||||||
type PdfQuestionResponse = Annotated[
|
type PdfQuestionResponse = Annotated[
|
||||||
PdfQuestionTerminalResponse | NeedContentResponse,
|
PdfQuestionTerminalResponse | NeedIngestResponse,
|
||||||
Field(discriminator="outcome"),
|
Field(discriminator="outcome"),
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
# ``orchestrate`` may also emit an ``EditPlanResponse`` on the math-routing
|
||||||
|
# first turn (``outcome=PLAN`` with ``resume_with=PDF_QUESTION``). It's not in
|
||||||
|
# ``PdfQuestionTerminalResponse`` because that alias would otherwise duplicate
|
||||||
|
# the PLAN branch already provided by ``PdfEditTerminalResponse`` in the
|
||||||
|
# top-level :class:`OrchestratorResponse` discriminated union.
|
||||||
|
type PdfQuestionOrchestrateResponse = PdfQuestionResponse | EditPlanResponse
|
||||||
|
|||||||
@@ -4,48 +4,36 @@ from pydantic import Field
|
|||||||
|
|
||||||
from stirling.models import ApiModel
|
from stirling.models import ApiModel
|
||||||
|
|
||||||
MAX_INDEX_TEXT_LENGTH = 1_000_000 # 1MB text limit per index request
|
from .common import FileId
|
||||||
|
|
||||||
|
|
||||||
class RagStatusResponse(ApiModel):
|
class IngestedPageText(ApiModel):
|
||||||
embedding_model: str
|
page_number: int = Field(ge=1)
|
||||||
collections: list[str]
|
text: str
|
||||||
|
|
||||||
|
|
||||||
class RagIndexRequest(ApiModel):
|
class IngestDocumentRequest(ApiModel):
|
||||||
collection: str = Field(min_length=1)
|
"""Replace-ingest a document's content into RAG under the given document_id.
|
||||||
text: str = Field(max_length=MAX_INDEX_TEXT_LENGTH)
|
|
||||||
source: str = ""
|
Each content-type field is optional; the endpoint replaces the document's entire
|
||||||
metadata: dict[str, str] = Field(default_factory=dict)
|
stored content with whatever is provided. To add a content type later, call again
|
||||||
|
with all content types the document should have (incremental-add-without-replace
|
||||||
|
will be a separate endpoint if/when we need it).
|
||||||
|
|
||||||
|
``source`` is a human-readable label (typically the original filename) that flows
|
||||||
|
into chunk metadata so search results are readable when document_id is a hash.
|
||||||
|
"""
|
||||||
|
|
||||||
|
document_id: FileId = Field(min_length=1)
|
||||||
|
source: str = Field(min_length=1)
|
||||||
|
page_text: list[IngestedPageText] | None = None
|
||||||
|
|
||||||
|
|
||||||
class RagIndexResponse(ApiModel):
|
class IngestDocumentResponse(ApiModel):
|
||||||
collection: str
|
document_id: FileId
|
||||||
chunks_indexed: int
|
chunks_indexed: int
|
||||||
|
|
||||||
|
|
||||||
class RagSearchRequest(ApiModel):
|
class DeleteDocumentResponse(ApiModel):
|
||||||
query: str
|
document_id: FileId
|
||||||
collection: str | None = Field(default=None, min_length=1)
|
deleted: bool
|
||||||
top_k: int = 5
|
|
||||||
|
|
||||||
|
|
||||||
class RagSearchResultItem(ApiModel):
|
|
||||||
text: str
|
|
||||||
source: str
|
|
||||||
chunk_id: str
|
|
||||||
score: float
|
|
||||||
|
|
||||||
|
|
||||||
class RagSearchResponse(ApiModel):
|
|
||||||
query: str
|
|
||||||
results: list[RagSearchResultItem]
|
|
||||||
|
|
||||||
|
|
||||||
class RagCollectionsResponse(ApiModel):
|
|
||||||
collections: list[str]
|
|
||||||
|
|
||||||
|
|
||||||
class RagDeleteCollectionResponse(ApiModel):
|
|
||||||
status: str
|
|
||||||
collection: str
|
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
from . import tool_models
|
from . import tool_models
|
||||||
from .base import ApiModel
|
from .base import ApiModel, FileId
|
||||||
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
|
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"ApiModel",
|
"ApiModel",
|
||||||
|
"FileId",
|
||||||
"OPERATIONS",
|
"OPERATIONS",
|
||||||
"ParamToolModel",
|
"ParamToolModel",
|
||||||
"ToolEndpoint",
|
"ToolEndpoint",
|
||||||
|
|||||||
@@ -1,8 +1,15 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import NewType
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict
|
from pydantic import BaseModel, ConfigDict
|
||||||
from pydantic.alias_generators import to_camel
|
from pydantic.alias_generators import to_camel
|
||||||
|
|
||||||
|
# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's
|
||||||
|
# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key
|
||||||
|
# throughout the engine.
|
||||||
|
FileId = NewType("FileId", str)
|
||||||
|
|
||||||
|
|
||||||
class ApiModel(BaseModel):
|
class ApiModel(BaseModel):
|
||||||
model_config = ConfigDict(
|
model_config = ConfigDict(
|
||||||
|
|||||||
@@ -1,11 +1,16 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
|
|
||||||
from pydantic_ai import FunctionToolset
|
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
||||||
from pydantic_ai.toolsets import AbstractToolset
|
from pydantic_ai.toolsets import AbstractToolset
|
||||||
|
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.rag.service import RagService
|
from stirling.rag.service import RagService
|
||||||
|
from stirling.rag.store import SearchResult
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class RagCapability:
|
class RagCapability:
|
||||||
@@ -22,19 +27,29 @@ class RagCapability:
|
|||||||
|
|
||||||
When no collections are pinned, the instructions are generated dynamically at
|
When no collections are pinned, the instructions are generated dynamically at
|
||||||
run time so the agent sees the current list of collections in the store.
|
run time so the agent sees the current list of collections in the store.
|
||||||
|
|
||||||
|
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a
|
||||||
|
single agent run.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
rag_service: RagService,
|
rag_service: RagService,
|
||||||
collections: list[str] | None = None,
|
collections: list[FileId] | None = None,
|
||||||
top_k: int = 5,
|
top_k: int = 5,
|
||||||
|
max_searches: int = 5,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._rag_service = rag_service
|
self._rag_service = rag_service
|
||||||
self._collections = collections
|
self._collections = collections
|
||||||
self._top_k = top_k
|
self._top_k = top_k
|
||||||
|
self._max_searches = max_searches
|
||||||
|
self._search_count = 0
|
||||||
toolset: FunctionToolset[None] = FunctionToolset()
|
toolset: FunctionToolset[None] = FunctionToolset()
|
||||||
toolset.add_function(self._search_knowledge, name="search_knowledge")
|
toolset.add_function(
|
||||||
|
self._search_knowledge,
|
||||||
|
name="search_knowledge",
|
||||||
|
prepare=self._prepare_search_knowledge,
|
||||||
|
)
|
||||||
self._toolset = toolset
|
self._toolset = toolset
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -48,7 +63,7 @@ class RagCapability:
|
|||||||
return self._toolset
|
return self._toolset
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _static_instructions_text(collections: list[str]) -> str:
|
def _static_instructions_text(collections: list[FileId]) -> str:
|
||||||
collection_desc = f"collections: {', '.join(collections)}"
|
collection_desc = f"collections: {', '.join(collections)}"
|
||||||
return (
|
return (
|
||||||
"You have access to a knowledge base search tool called 'search_knowledge'. "
|
"You have access to a knowledge base search tool called 'search_knowledge'. "
|
||||||
@@ -73,6 +88,18 @@ class RagCapability:
|
|||||||
"You do not have to use it if the answer is already clear from the provided text."
|
"You do not have to use it if the answer is already clear from the provided text."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _prepare_search_knowledge(
|
||||||
|
self,
|
||||||
|
ctx: RunContext[None],
|
||||||
|
tool_def: ToolDefinition,
|
||||||
|
) -> ToolDefinition | None:
|
||||||
|
"""Remove the search tool from the agent's toolset once the per-run search
|
||||||
|
budget is exhausted. The agent then has no choice but to answer from what it
|
||||||
|
has already retrieved, which prevents runaway search loops."""
|
||||||
|
if self._search_count >= self._max_searches:
|
||||||
|
return None
|
||||||
|
return tool_def
|
||||||
|
|
||||||
async def _search_knowledge(self, query: str, max_results: int | None = None) -> str:
|
async def _search_knowledge(self, query: str, max_results: int | None = None) -> str:
|
||||||
"""Search the knowledge base for information relevant to the query.
|
"""Search the knowledge base for information relevant to the query.
|
||||||
|
|
||||||
@@ -83,6 +110,7 @@ class RagCapability:
|
|||||||
Returns:
|
Returns:
|
||||||
Formatted text with the most relevant knowledge base excerpts.
|
Formatted text with the most relevant knowledge base excerpts.
|
||||||
"""
|
"""
|
||||||
|
self._search_count += 1
|
||||||
k = max_results if max_results is not None else self._top_k
|
k = max_results if max_results is not None else self._top_k
|
||||||
if self._collections:
|
if self._collections:
|
||||||
all_results = []
|
all_results = []
|
||||||
@@ -95,8 +123,21 @@ class RagCapability:
|
|||||||
results = await self._rag_service.search(query, top_k=k)
|
results = await self._rag_service.search(query, top_k=k)
|
||||||
|
|
||||||
if not results:
|
if not results:
|
||||||
|
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
|
||||||
return "No relevant results found in the knowledge base."
|
return "No relevant results found in the knowledge base."
|
||||||
|
|
||||||
|
formatted = self._format_results(results)
|
||||||
|
logger.info(
|
||||||
|
"[rag] search_knowledge query=%r -> %d results, %d chars",
|
||||||
|
query,
|
||||||
|
len(results),
|
||||||
|
len(formatted),
|
||||||
|
)
|
||||||
|
logger.debug("[rag] search_knowledge query=%r returned:\n%s", query, formatted)
|
||||||
|
return formatted
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _format_results(results: list[SearchResult]) -> str:
|
||||||
sections = []
|
sections = []
|
||||||
for i, result in enumerate(results, 1):
|
for i, result in enumerate(results, 1):
|
||||||
source = result.document.metadata.get("source", "unknown")
|
source = result.document.metadata.get("source", "unknown")
|
||||||
|
|||||||
@@ -5,14 +5,27 @@ from pydantic_ai import Embedder
|
|||||||
from stirling.rag.chunker import chunk_text
|
from stirling.rag.chunker import chunk_text
|
||||||
from stirling.rag.store import Document
|
from stirling.rag.store import Document
|
||||||
|
|
||||||
|
# Keep each upstream embed request under every major provider's per-call limit while
|
||||||
|
# still batching large enough that a book-sized document ingests in a reasonable number
|
||||||
|
# of round trips. VoyageAI caps at 1000, OpenAI at 2048, Cohere at 96; 256 is a good
|
||||||
|
# default for Voyage/OpenAI. Cohere users should pass a lower value via construction.
|
||||||
|
DEFAULT_EMBED_BATCH_SIZE = 256
|
||||||
|
|
||||||
|
|
||||||
class EmbeddingService:
|
class EmbeddingService:
|
||||||
"""Wraps Pydantic AI's Embedder to provide document chunking and embedding."""
|
"""Wraps Pydantic AI's Embedder to provide document chunking and embedding."""
|
||||||
|
|
||||||
def __init__(self, model_name: str, chunk_size: int = 512, chunk_overlap: int = 64) -> None:
|
def __init__(
|
||||||
|
self,
|
||||||
|
model_name: str,
|
||||||
|
chunk_size: int = 512,
|
||||||
|
chunk_overlap: int = 64,
|
||||||
|
embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
|
||||||
|
) -> None:
|
||||||
self._embedder = Embedder(model_name)
|
self._embedder = Embedder(model_name)
|
||||||
self._chunk_size = chunk_size
|
self._chunk_size = chunk_size
|
||||||
self._chunk_overlap = chunk_overlap
|
self._chunk_overlap = chunk_overlap
|
||||||
|
self._embed_batch_size = embed_batch_size
|
||||||
|
|
||||||
async def embed_query(self, text: str) -> list[float]:
|
async def embed_query(self, text: str) -> list[float]:
|
||||||
"""Embed a search query, optimised for retrieval."""
|
"""Embed a search query, optimised for retrieval."""
|
||||||
@@ -20,11 +33,19 @@ class EmbeddingService:
|
|||||||
return list(result.embeddings[0])
|
return list(result.embeddings[0])
|
||||||
|
|
||||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||||
"""Embed multiple document texts for indexing."""
|
"""Embed multiple document texts for indexing.
|
||||||
|
|
||||||
|
Splits the input into batches of ``embed_batch_size`` so callers can hand us
|
||||||
|
any number of chunks without hitting provider per-request limits.
|
||||||
|
"""
|
||||||
if not texts:
|
if not texts:
|
||||||
return []
|
return []
|
||||||
result = await self._embedder.embed_documents(texts)
|
all_embeddings: list[list[float]] = []
|
||||||
return [list(emb) for emb in result.embeddings]
|
for start in range(0, len(texts), self._embed_batch_size):
|
||||||
|
batch = texts[start : start + self._embed_batch_size]
|
||||||
|
result = await self._embedder.embed_documents(batch)
|
||||||
|
all_embeddings.extend(list(emb) for emb in result.embeddings)
|
||||||
|
return all_embeddings
|
||||||
|
|
||||||
def chunk_and_prepare(
|
def chunk_and_prepare(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -131,3 +131,7 @@ class PgVectorStore(VectorStore):
|
|||||||
)
|
)
|
||||||
row = await cur.fetchone()
|
row = await cur.fetchone()
|
||||||
return row is not None
|
return row is not None
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
# Connections are opened and closed per call, so nothing persistent to release.
|
||||||
|
return None
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.rag.embedder import EmbeddingService
|
from stirling.rag.embedder import EmbeddingService
|
||||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
from stirling.rag.store import Document, SearchResult, VectorStore
|
||||||
|
|
||||||
@@ -18,7 +19,7 @@ class RagService:
|
|||||||
|
|
||||||
async def index_text(
|
async def index_text(
|
||||||
self,
|
self,
|
||||||
collection: str,
|
collection: FileId,
|
||||||
text: str,
|
text: str,
|
||||||
source: str = "",
|
source: str = "",
|
||||||
metadata: dict[str, str] | None = None,
|
metadata: dict[str, str] | None = None,
|
||||||
@@ -31,7 +32,7 @@ class RagService:
|
|||||||
await self._store.add_documents(collection, documents, embeddings)
|
await self._store.add_documents(collection, documents, embeddings)
|
||||||
return len(documents)
|
return len(documents)
|
||||||
|
|
||||||
async def index_documents(self, collection: str, documents: list[Document]) -> int:
|
async def index_documents(self, collection: FileId, documents: list[Document]) -> int:
|
||||||
"""Embed and store pre-chunked documents. Returns the number stored."""
|
"""Embed and store pre-chunked documents. Returns the number stored."""
|
||||||
if not documents:
|
if not documents:
|
||||||
return 0
|
return 0
|
||||||
@@ -39,10 +40,23 @@ class RagService:
|
|||||||
await self._store.add_documents(collection, documents, embeddings)
|
await self._store.add_documents(collection, documents, embeddings)
|
||||||
return len(documents)
|
return len(documents)
|
||||||
|
|
||||||
|
def chunk_text(
|
||||||
|
self,
|
||||||
|
text: str,
|
||||||
|
source: str = "",
|
||||||
|
base_metadata: dict[str, str] | None = None,
|
||||||
|
) -> list[Document]:
|
||||||
|
"""Chunk text into Document objects ready for indexing. Does NOT embed.
|
||||||
|
|
||||||
|
Exposed so callers that ingest many chunks can accumulate them across calls
|
||||||
|
and then pass the full batch to ``index_documents`` for a single embedding pass.
|
||||||
|
"""
|
||||||
|
return self._embedder.chunk_and_prepare(text, source=source, base_metadata=base_metadata)
|
||||||
|
|
||||||
async def search(
|
async def search(
|
||||||
self,
|
self,
|
||||||
query: str,
|
query: str,
|
||||||
collection: str | None = None,
|
collection: FileId | None = None,
|
||||||
top_k: int | None = None,
|
top_k: int | None = None,
|
||||||
) -> list[SearchResult]:
|
) -> list[SearchResult]:
|
||||||
"""Embed query and search across one or all collections.
|
"""Embed query and search across one or all collections.
|
||||||
@@ -71,10 +85,18 @@ class RagService:
|
|||||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||||
return all_results[:k]
|
return all_results[:k]
|
||||||
|
|
||||||
async def delete_collection(self, collection: str) -> None:
|
async def delete_collection(self, collection: FileId) -> None:
|
||||||
"""Remove a collection and all its documents."""
|
"""Remove a collection and all its documents."""
|
||||||
await self._store.delete_collection(collection)
|
await self._store.delete_collection(collection)
|
||||||
|
|
||||||
async def list_collections(self) -> list[str]:
|
async def has_collection(self, collection: FileId) -> bool:
|
||||||
|
"""Check whether a collection exists."""
|
||||||
|
return await self._store.has_collection(collection)
|
||||||
|
|
||||||
|
async def list_collections(self) -> list[FileId]:
|
||||||
"""List all available collections."""
|
"""List all available collections."""
|
||||||
return await self._store.list_collections()
|
return [FileId(name) for name in await self._store.list_collections()]
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
"""Release the underlying vector store's resources."""
|
||||||
|
await self._store.close()
|
||||||
|
|||||||
@@ -225,3 +225,19 @@ class SqliteVecStore(VectorStore):
|
|||||||
def _sync_has_collection(self, collection: str) -> bool:
|
def _sync_has_collection(self, collection: str) -> bool:
|
||||||
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
|
row = self._conn.execute("SELECT 1 FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||||
return row is not None
|
return row is not None
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
async with self._lock:
|
||||||
|
await asyncio.to_thread(self._sync_close)
|
||||||
|
|
||||||
|
def _sync_close(self) -> None:
|
||||||
|
"""Checkpoint the WAL into the main database file and close the connection so
|
||||||
|
the .db-shm and .db-wal files are cleaned up on graceful shutdown."""
|
||||||
|
if self._db_path is not None:
|
||||||
|
try:
|
||||||
|
self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
||||||
|
self._conn.commit()
|
||||||
|
except sqlite3.Error:
|
||||||
|
# Best effort: if checkpointing fails we still want to close the connection.
|
||||||
|
pass
|
||||||
|
self._conn.close()
|
||||||
|
|||||||
@@ -57,3 +57,7 @@ class VectorStore(ABC):
|
|||||||
@abstractmethod
|
@abstractmethod
|
||||||
async def has_collection(self, collection: str) -> bool:
|
async def has_collection(self, collection: str) -> bool:
|
||||||
"""Check whether a collection exists."""
|
"""Check whether a collection exists."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def close(self) -> None:
|
||||||
|
"""Release any resources held by the store (connections, handles, etc.)."""
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ from pydantic import ValidationError
|
|||||||
|
|
||||||
from stirling.agents.math_presentation import extract_math_verdict
|
from stirling.agents.math_presentation import extract_math_verdict
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
|
AiFile,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
ExtractedTextArtifact,
|
ExtractedTextArtifact,
|
||||||
MathAuditorToolReportArtifact,
|
MathAuditorToolReportArtifact,
|
||||||
@@ -21,6 +22,7 @@ from stirling.contracts import (
|
|||||||
WorkflowArtifact,
|
WorkflowArtifact,
|
||||||
)
|
)
|
||||||
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
||||||
|
from stirling.models import FileId
|
||||||
|
|
||||||
|
|
||||||
def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
|
def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
|
||||||
@@ -42,7 +44,7 @@ def _make_verdict(discrepancies: list[Discrepancy]) -> Verdict:
|
|||||||
def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest:
|
def _orchestrator_request_with_artifacts(artifacts: list[WorkflowArtifact]) -> OrchestratorRequest:
|
||||||
return OrchestratorRequest(
|
return OrchestratorRequest(
|
||||||
user_message="review the math",
|
user_message="review the math",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
artifacts=artifacts,
|
artifacts=artifacts,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -23,8 +23,9 @@ from unittest.mock import AsyncMock, patch
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from stirling.agents import OrchestratorAgent
|
from stirling.agents import OrchestratorAgent
|
||||||
from stirling.contracts import OrchestratorRequest
|
from stirling.contracts import AiFile, OrchestratorRequest
|
||||||
from stirling.contracts.pdf_edit import EditPlanResponse
|
from stirling.contracts.pdf_edit import EditPlanResponse
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
|
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
@@ -39,7 +40,7 @@ async def test_delegate_pdf_review_wires_prompt_to_tool_step(runtime: AppRuntime
|
|||||||
orchestrator = OrchestratorAgent(runtime)
|
orchestrator = OrchestratorAgent(runtime)
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="please add review comments flagging ambiguous dates",
|
user_message="please add review comments flagging ambiguous dates",
|
||||||
file_names=["contract.pdf"],
|
files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
|
||||||
)
|
)
|
||||||
ctx = SimpleNamespace(deps=_FakeDeps(request=request))
|
ctx = SimpleNamespace(deps=_FakeDeps(request=request))
|
||||||
|
|
||||||
|
|||||||
@@ -12,12 +12,15 @@ import pytest
|
|||||||
|
|
||||||
from stirling.agents.pdf_questions import _MATH_SYNTH_SYSTEM_PROMPT, PdfQuestionAgent
|
from stirling.agents.pdf_questions import _MATH_SYNTH_SYSTEM_PROMPT, PdfQuestionAgent
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
|
AiFile,
|
||||||
|
EditPlanResponse,
|
||||||
MathAuditorToolReportArtifact,
|
MathAuditorToolReportArtifact,
|
||||||
OrchestratorRequest,
|
OrchestratorRequest,
|
||||||
PdfQuestionAnswerResponse,
|
PdfQuestionAnswerResponse,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
)
|
)
|
||||||
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.models.agent_tool_models import AgentToolId
|
from stirling.models.agent_tool_models import AgentToolId
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
@@ -49,25 +52,23 @@ def _make_verdict() -> Verdict:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_orchestrate_classifier_true_embeds_plan_in_answer(runtime: AppRuntime) -> None:
|
async def test_orchestrate_classifier_true_returns_math_audit_plan(runtime: AppRuntime) -> None:
|
||||||
"""First turn — classifier says math; the response is a PdfQuestionAnswerResponse
|
"""First turn — classifier says math; the response is an EditPlanResponse
|
||||||
with the math-auditor plan attached as a nullable ``edit_plan`` field. The
|
(``outcome=PLAN``) with ``resume_with=PDF_QUESTION``. The caller runs the
|
||||||
answer is empty on this turn; the caller runs the embedded plan and resumes."""
|
plan and re-invokes the orchestrator with the verdict in artifacts."""
|
||||||
agent = PdfQuestionAgent(runtime)
|
agent = PdfQuestionAgent(runtime)
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="ist die mathematik korrekt?",
|
user_message="ist die mathematik korrekt?",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
)
|
)
|
||||||
|
|
||||||
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
|
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
|
||||||
response = await agent.orchestrate(request)
|
response = await agent.orchestrate(request)
|
||||||
|
|
||||||
assert isinstance(response, PdfQuestionAnswerResponse)
|
assert isinstance(response, EditPlanResponse)
|
||||||
assert response.answer == ""
|
assert response.resume_with == SupportedCapability.PDF_QUESTION
|
||||||
assert response.edit_plan is not None
|
assert len(response.steps) == 1
|
||||||
assert response.edit_plan.resume_with == SupportedCapability.PDF_QUESTION
|
assert response.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
|
||||||
assert len(response.edit_plan.steps) == 1
|
|
||||||
assert response.edit_plan.steps[0].tool == AgentToolId.MATH_AUDITOR_AGENT
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
@@ -81,7 +82,7 @@ async def test_orchestrate_resume_synthesises_answer_without_calling_classifier(
|
|||||||
verdict = _make_verdict()
|
verdict = _make_verdict()
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="ist die mathematik korrekt?",
|
user_message="ist die mathematik korrekt?",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
|
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
|
||||||
)
|
)
|
||||||
canned_answer = "Die Summe stimmt nicht: angegeben $215,000, erwartet $215,500."
|
canned_answer = "Die Summe stimmt nicht: angegeben $215,000, erwartet $215,500."
|
||||||
|
|||||||
@@ -20,9 +20,9 @@ from stirling.agents.pdf_review import (
|
|||||||
_LocalisedComment,
|
_LocalisedComment,
|
||||||
_LocalisedVerdict,
|
_LocalisedVerdict,
|
||||||
)
|
)
|
||||||
from stirling.contracts import EditPlanResponse, OrchestratorRequest, SupportedCapability
|
from stirling.contracts import AiFile, EditPlanResponse, OrchestratorRequest, SupportedCapability
|
||||||
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
from stirling.contracts.ledger import Discrepancy, DiscrepancyKind, Severity, Verdict
|
||||||
from stirling.models import ToolEndpoint
|
from stirling.models import FileId, ToolEndpoint
|
||||||
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
|
from stirling.models.agent_tool_models import AgentToolId, PdfCommentAgentParams
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
@@ -144,7 +144,10 @@ async def test_orchestrate_classifier_true_emits_math_audit_plan(runtime: AppRun
|
|||||||
"""First turn — when the math-intent classifier says yes, emit a one-step plan
|
"""First turn — when the math-intent classifier says yes, emit a one-step plan
|
||||||
calling the math auditor with resume_with=PDF_REVIEW."""
|
calling the math auditor with resume_with=PDF_REVIEW."""
|
||||||
agent = PdfReviewAgent(runtime)
|
agent = PdfReviewAgent(runtime)
|
||||||
request = OrchestratorRequest(user_message="vérifie les totaux", file_names=["report.pdf"])
|
request = OrchestratorRequest(
|
||||||
|
user_message="vérifie les totaux",
|
||||||
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
|
)
|
||||||
|
|
||||||
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
|
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=True)):
|
||||||
response = await agent.orchestrate(request)
|
response = await agent.orchestrate(request)
|
||||||
@@ -161,7 +164,7 @@ async def test_orchestrate_classifier_false_routes_to_pdf_comment_agent(runtime:
|
|||||||
agent = PdfReviewAgent(runtime)
|
agent = PdfReviewAgent(runtime)
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="review the invoices for ambiguous wording",
|
user_message="review the invoices for ambiguous wording",
|
||||||
file_names=["contract.pdf"],
|
files=[AiFile(id=FileId("contract-id"), name="contract.pdf")],
|
||||||
)
|
)
|
||||||
|
|
||||||
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=False)):
|
with patch.object(agent._math_intent_classifier, "classify", AsyncMock(return_value=False)):
|
||||||
@@ -187,7 +190,7 @@ async def test_orchestrate_resume_uses_verdict_without_calling_classifier(
|
|||||||
verdict = _make_verdict([_discrepancy(page=0, stated="$100")])
|
verdict = _make_verdict([_discrepancy(page=0, stated="$100")])
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="flag math errors",
|
user_message="flag math errors",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
|
artifacts=[MathAuditorToolReportArtifact(report=verdict)],
|
||||||
)
|
)
|
||||||
canned = _LocalisedVerdict(
|
canned = _LocalisedVerdict(
|
||||||
|
|||||||
@@ -30,6 +30,7 @@ def build_app_settings() -> AppSettings:
|
|||||||
rag_chunk_size=512,
|
rag_chunk_size=512,
|
||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
|
rag_max_searches=5,
|
||||||
max_pages=200,
|
max_pages=200,
|
||||||
max_characters=200_000,
|
max_characters=200_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
@@ -42,6 +42,7 @@ class StubSettingsProvider:
|
|||||||
rag_chunk_size=512,
|
rag_chunk_size=512,
|
||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
|
rag_max_searches=5,
|
||||||
max_pages=100,
|
max_pages=100,
|
||||||
max_characters=100_000,
|
max_characters=100_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import pytest
|
|||||||
from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
|
from stirling.agents import PdfEditAgent, PdfEditParameterSelector, PdfEditPlanSelection
|
||||||
from stirling.agents.pdf_edit import PdfEditPlanOutput
|
from stirling.agents.pdf_edit import PdfEditPlanOutput
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
|
AiFile,
|
||||||
EditCannotDoResponse,
|
EditCannotDoResponse,
|
||||||
EditClarificationRequest,
|
EditClarificationRequest,
|
||||||
EditPlanResponse,
|
EditPlanResponse,
|
||||||
@@ -19,6 +20,7 @@ from stirling.contracts import (
|
|||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
ToolOperationStep,
|
ToolOperationStep,
|
||||||
)
|
)
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
|
from stirling.models.tool_models import Angle, FlattenParams, RotatePdfParams, ToolEndpoint
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
@@ -91,7 +93,7 @@ async def test_pdf_edit_agent_builds_multi_step_plan(runtime: AppRuntime) -> Non
|
|||||||
response = await agent.handle(
|
response = await agent.handle(
|
||||||
PdfEditRequest(
|
PdfEditRequest(
|
||||||
user_message="Rotate the PDF clockwise and then compress it.",
|
user_message="Rotate the PDF clockwise and then compress it.",
|
||||||
file_names=["scan.pdf"],
|
files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -117,7 +119,7 @@ async def test_pdf_edit_agent_passes_previous_steps_to_parameter_selector(runtim
|
|||||||
|
|
||||||
request = PdfEditRequest(
|
request = PdfEditRequest(
|
||||||
user_message="Rotate the PDF clockwise and then compress it.",
|
user_message="Rotate the PDF clockwise and then compress it.",
|
||||||
file_names=["scan.pdf"],
|
files=[AiFile(id=FileId("scan-id"), name="scan.pdf")],
|
||||||
)
|
)
|
||||||
response = await agent.handle(request)
|
response = await agent.handle(request)
|
||||||
|
|
||||||
@@ -181,13 +183,18 @@ async def test_pdf_edit_agent_returns_need_content_without_building_plan(runtime
|
|||||||
response = await agent.handle(
|
response = await agent.handle(
|
||||||
PdfEditRequest(
|
PdfEditRequest(
|
||||||
user_message="Split after every page that says 'NEW PAGE'.",
|
user_message="Split after every page that says 'NEW PAGE'.",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
assert isinstance(response, NeedContentResponse)
|
assert isinstance(response, NeedContentResponse)
|
||||||
assert response.resume_with == SupportedCapability.PDF_EDIT
|
assert response.resume_with == SupportedCapability.PDF_EDIT
|
||||||
assert response.files == [NeedContentFileRequest(file_name="report.pdf", content_types=[PdfContentType.PAGE_TEXT])]
|
assert response.files == [
|
||||||
|
NeedContentFileRequest(
|
||||||
|
file=AiFile(id=FileId("report-id"), name="report.pdf"),
|
||||||
|
content_types=[PdfContentType.PAGE_TEXT],
|
||||||
|
)
|
||||||
|
]
|
||||||
assert response.max_pages == runtime.settings.max_pages
|
assert response.max_pages == runtime.settings.max_pages
|
||||||
assert response.max_characters == runtime.settings.max_characters
|
assert response.max_characters == runtime.settings.max_characters
|
||||||
assert parameter_selector.calls == []
|
assert parameter_selector.calls == []
|
||||||
@@ -269,7 +276,7 @@ async def test_pdf_edit_agent_passes_page_text_to_parameter_selector(runtime: Ap
|
|||||||
await agent.handle(
|
await agent.handle(
|
||||||
PdfEditRequest(
|
PdfEditRequest(
|
||||||
user_message="Rotate clockwise.",
|
user_message="Rotate clockwise.",
|
||||||
file_names=["report.pdf"],
|
files=[AiFile(id=FileId("report-id"), name="report.pdf")],
|
||||||
page_text=page_text,
|
page_text=page_text,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,64 +1,133 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import replace
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from stirling.agents import PdfQuestionAgent
|
from stirling.agents import PdfQuestionAgent
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
|
AiFile,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
NeedContentResponse,
|
NeedIngestResponse,
|
||||||
|
PdfContentType,
|
||||||
PdfQuestionAnswerResponse,
|
PdfQuestionAnswerResponse,
|
||||||
PdfQuestionNotFoundResponse,
|
PdfQuestionNotFoundResponse,
|
||||||
PdfQuestionRequest,
|
PdfQuestionRequest,
|
||||||
|
PdfQuestionTerminalResponse,
|
||||||
PdfTextSelection,
|
PdfTextSelection,
|
||||||
|
SupportedCapability,
|
||||||
)
|
)
|
||||||
|
from stirling.models import FileId
|
||||||
|
from stirling.rag import Document, RagService, SqliteVecStore
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
|
|
||||||
class StubPdfQuestionAgent(PdfQuestionAgent):
|
class StubEmbedder:
|
||||||
def __init__(self, runtime: AppRuntime, response: PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse) -> None:
|
"""Deterministic embeddings so RAG lookups work in tests without network."""
|
||||||
super().__init__(runtime)
|
|
||||||
self.response = response
|
|
||||||
|
|
||||||
async def _run_answer_agent(
|
def __init__(self, dim: int = 8) -> None:
|
||||||
|
self._dim = dim
|
||||||
|
|
||||||
|
async def embed_query(self, text: str) -> list[float]:
|
||||||
|
h = hash(text) % 1000
|
||||||
|
return [(h + i) / 1000.0 for i in range(self._dim)]
|
||||||
|
|
||||||
|
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||||
|
return [await self.embed_query(t) for t in texts]
|
||||||
|
|
||||||
|
def chunk_and_prepare(
|
||||||
self,
|
self,
|
||||||
request: PdfQuestionRequest,
|
text: str,
|
||||||
) -> PdfQuestionAnswerResponse | PdfQuestionNotFoundResponse:
|
source: str = "",
|
||||||
return self.response
|
base_metadata: dict[str, str] | None = None,
|
||||||
|
) -> list[Document]:
|
||||||
|
from stirling.rag.chunker import chunk_text
|
||||||
|
|
||||||
|
chunks = chunk_text(text, 100, 10)
|
||||||
|
docs: list[Document] = []
|
||||||
|
for i, chunk in enumerate(chunks):
|
||||||
|
meta = dict(base_metadata) if base_metadata else {}
|
||||||
|
meta["source"] = source
|
||||||
|
meta["chunk_index"] = str(i)
|
||||||
|
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
|
||||||
|
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
|
||||||
|
return docs
|
||||||
|
|
||||||
|
|
||||||
def invoice_page() -> ExtractedFileText:
|
class StubPdfQuestionAgent(PdfQuestionAgent):
|
||||||
return ExtractedFileText(
|
def __init__(self, runtime: AppRuntime, response: PdfQuestionTerminalResponse) -> None:
|
||||||
file_name="invoice.pdf",
|
super().__init__(runtime)
|
||||||
pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
|
self._response = response
|
||||||
|
|
||||||
|
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
|
||||||
|
return self._response
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
|
||||||
|
"""A runtime whose RAG service uses a stub embedder + ephemeral store."""
|
||||||
|
stub = RagService(
|
||||||
|
embedder=StubEmbedder(), # type: ignore[arg-type]
|
||||||
|
store=SqliteVecStore.ephemeral(),
|
||||||
|
default_top_k=runtime.settings.rag_default_top_k,
|
||||||
)
|
)
|
||||||
|
return replace(runtime, rag_service=stub)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_pdf_question_agent_requires_extracted_text(runtime: AppRuntime) -> None:
|
async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
agent = PdfQuestionAgent(runtime)
|
agent = PdfQuestionAgent(runtime_with_stub_rag)
|
||||||
|
|
||||||
response = await agent.handle(
|
missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf")
|
||||||
PdfQuestionRequest(question="What is the total?", page_text=[], file_names=["test.pdf"])
|
response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[missing_file]))
|
||||||
)
|
|
||||||
|
|
||||||
assert isinstance(response, NeedContentResponse)
|
assert isinstance(response, NeedIngestResponse)
|
||||||
|
assert response.resume_with == SupportedCapability.PDF_QUESTION
|
||||||
|
assert response.files_to_ingest == [missing_file]
|
||||||
|
assert PdfContentType.PAGE_TEXT in response.content_types
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -> None:
|
async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
|
await runtime_with_stub_rag.rag_service.index_text(
|
||||||
|
collection=FileId("present-id"),
|
||||||
|
text="Invoice total: 120.00.",
|
||||||
|
source="present.pdf",
|
||||||
|
)
|
||||||
|
agent = PdfQuestionAgent(runtime_with_stub_rag)
|
||||||
|
|
||||||
|
present_file = AiFile(id=FileId("present-id"), name="present.pdf")
|
||||||
|
missing_file = AiFile(id=FileId("missing-id"), name="missing.pdf")
|
||||||
|
response = await agent.handle(PdfQuestionRequest(question="What is the total?", files=[present_file, missing_file]))
|
||||||
|
|
||||||
|
assert isinstance(response, NeedIngestResponse)
|
||||||
|
assert response.files_to_ingest == [missing_file]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
|
await runtime_with_stub_rag.rag_service.index_text(
|
||||||
|
collection=FileId("invoice-id"),
|
||||||
|
text="Invoice total: 120.00.",
|
||||||
|
source="invoice.pdf",
|
||||||
|
)
|
||||||
agent = StubPdfQuestionAgent(
|
agent = StubPdfQuestionAgent(
|
||||||
runtime,
|
runtime_with_stub_rag,
|
||||||
PdfQuestionAnswerResponse(
|
PdfQuestionAnswerResponse(
|
||||||
answer="The invoice total is 120.00.",
|
answer="The invoice total is 120.00.",
|
||||||
evidence=[invoice_page()],
|
evidence=[
|
||||||
|
ExtractedFileText(
|
||||||
|
file_name="invoice.pdf",
|
||||||
|
pages=[PdfTextSelection(page_number=1, text="Invoice total: 120.00")],
|
||||||
|
)
|
||||||
|
],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
response = await agent.handle(
|
response = await agent.handle(
|
||||||
PdfQuestionRequest(
|
PdfQuestionRequest(
|
||||||
question="What is the total?",
|
question="What is the total?",
|
||||||
page_text=[invoice_page()],
|
files=[AiFile(id=FileId("invoice-id"), name="invoice.pdf")],
|
||||||
file_names=["invoice.pdf"],
|
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -67,19 +136,21 @@ async def test_pdf_question_agent_returns_grounded_answer(runtime: AppRuntime) -
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_pdf_question_agent_returns_not_found_when_text_is_insufficient(runtime: AppRuntime) -> None:
|
async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
agent = StubPdfQuestionAgent(runtime, PdfQuestionNotFoundResponse(reason="The answer is not present in the text."))
|
await runtime_with_stub_rag.rag_service.index_text(
|
||||||
|
collection=FileId("shipping-id"),
|
||||||
|
text="This page contains only a shipping address.",
|
||||||
|
source="shipping.pdf",
|
||||||
|
)
|
||||||
|
agent = StubPdfQuestionAgent(
|
||||||
|
runtime_with_stub_rag,
|
||||||
|
PdfQuestionNotFoundResponse(reason="The answer is not present in the text."),
|
||||||
|
)
|
||||||
|
|
||||||
response = await agent.handle(
|
response = await agent.handle(
|
||||||
PdfQuestionRequest(
|
PdfQuestionRequest(
|
||||||
question="What is the total?",
|
question="What is the total?",
|
||||||
page_text=[
|
files=[AiFile(id=FileId("shipping-id"), name="shipping.pdf")],
|
||||||
ExtractedFileText(
|
|
||||||
file_name="invoice.pdf",
|
|
||||||
pages=[PdfTextSelection(page_number=1, text="This page contains only a shipping address.")],
|
|
||||||
)
|
|
||||||
],
|
|
||||||
file_names=["invoice.pdf"],
|
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
+41
-16
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.rag.capability import RagCapability
|
from stirling.rag.capability import RagCapability
|
||||||
from stirling.rag.chunker import chunk_text
|
from stirling.rag.chunker import chunk_text
|
||||||
from stirling.rag.service import RagService
|
from stirling.rag.service import RagService
|
||||||
@@ -163,38 +164,38 @@ class TestRagService:
|
|||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_index_and_search(self, rag_service: RagService) -> None:
|
async def test_index_and_search(self, rag_service: RagService) -> None:
|
||||||
text = "Python is great for data science. It has many libraries like pandas and numpy."
|
text = "Python is great for data science. It has many libraries like pandas and numpy."
|
||||||
count = await rag_service.index_text("docs", text, source="guide.pdf")
|
count = await rag_service.index_text(FileId("docs"), text, source="guide.pdf")
|
||||||
assert count > 0
|
assert count > 0
|
||||||
|
|
||||||
results = await rag_service.search("Python libraries", collection="docs")
|
results = await rag_service.search("Python libraries", collection=FileId("docs"))
|
||||||
assert len(results) > 0
|
assert len(results) > 0
|
||||||
assert results[0].document.text # non-empty text
|
assert results[0].document.text # non-empty text
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
|
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
|
||||||
count = await rag_service.index_text("docs", "", source="empty.pdf")
|
count = await rag_service.index_text(FileId("docs"), "", source="empty.pdf")
|
||||||
assert count == 0
|
assert count == 0
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
|
async def test_search_nonexistent_collection_returns_empty(self, rag_service: RagService) -> None:
|
||||||
results = await rag_service.search("anything", collection="nonexistent")
|
results = await rag_service.search("anything", collection=FileId("nonexistent"))
|
||||||
assert results == []
|
assert results == []
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_all_collections(self, rag_service: RagService) -> None:
|
async def test_search_all_collections(self, rag_service: RagService) -> None:
|
||||||
await rag_service.index_text("col-a", "Machine learning overview.", source="ml.pdf")
|
await rag_service.index_text(FileId("col-a"), "Machine learning overview.", source="ml.pdf")
|
||||||
await rag_service.index_text("col-b", "Deep learning with neural networks.", source="dl.pdf")
|
await rag_service.index_text(FileId("col-b"), "Deep learning with neural networks.", source="dl.pdf")
|
||||||
|
|
||||||
results = await rag_service.search("neural networks")
|
results = await rag_service.search("neural networks")
|
||||||
assert len(results) > 0
|
assert len(results) > 0
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_delete_collection(self, rag_service: RagService) -> None:
|
async def test_delete_collection(self, rag_service: RagService) -> None:
|
||||||
await rag_service.index_text("temp", "Temporary data.", source="tmp.pdf")
|
await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf")
|
||||||
collections = await rag_service.list_collections()
|
collections = await rag_service.list_collections()
|
||||||
assert "temp" in collections
|
assert "temp" in collections
|
||||||
|
|
||||||
await rag_service.delete_collection("temp")
|
await rag_service.delete_collection(FileId("temp"))
|
||||||
collections = await rag_service.list_collections()
|
collections = await rag_service.list_collections()
|
||||||
assert "temp" not in collections
|
assert "temp" not in collections
|
||||||
|
|
||||||
@@ -214,7 +215,7 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
|
|||||||
|
|
||||||
class TestRagCapability:
|
class TestRagCapability:
|
||||||
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
|
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
|
||||||
cap = RagCapability(rag_service, collections=["docs", "manuals"])
|
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
|
||||||
instructions = cap.instructions
|
instructions = cap.instructions
|
||||||
assert isinstance(instructions, str)
|
assert isinstance(instructions, str)
|
||||||
assert "docs, manuals" in instructions
|
assert "docs, manuals" in instructions
|
||||||
@@ -227,8 +228,8 @@ class TestRagCapability:
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
|
async def test_dynamic_instructions_list_available_collections(self, rag_service: RagService) -> None:
|
||||||
await rag_service.index_text("col-a", "Alpha content.", source="a.pdf")
|
await rag_service.index_text(FileId("col-a"), "Alpha content.", source="a.pdf")
|
||||||
await rag_service.index_text("col-b", "Beta content.", source="b.pdf")
|
await rag_service.index_text(FileId("col-b"), "Beta content.", source="b.pdf")
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(rag_service)
|
||||||
instructions_fn = cap.instructions
|
instructions_fn = cap.instructions
|
||||||
assert callable(instructions_fn)
|
assert callable(instructions_fn)
|
||||||
@@ -252,7 +253,7 @@ class TestRagCapability:
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
|
async def test_search_knowledge_formats_results_with_source_and_score(self, rag_service: RagService) -> None:
|
||||||
await rag_service.index_text("docs", "Python is a programming language.", source="guide.pdf")
|
await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf")
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(rag_service)
|
||||||
output = await _invoke_search_knowledge(cap, "Python")
|
output = await _invoke_search_knowledge(cap, "Python")
|
||||||
assert "[Result 1" in output
|
assert "[Result 1" in output
|
||||||
@@ -262,10 +263,10 @@ class TestRagCapability:
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
|
async def test_search_knowledge_restricts_to_pinned_collections(self, rag_service: RagService) -> None:
|
||||||
await rag_service.index_text("pinned", "Pinned collection content.", source="pinned.pdf")
|
await rag_service.index_text(FileId("pinned"), "Pinned collection content.", source="pinned.pdf")
|
||||||
await rag_service.index_text("other", "Content in another collection.", source="other.pdf")
|
await rag_service.index_text(FileId("other"), "Content in another collection.", source="other.pdf")
|
||||||
|
|
||||||
cap = RagCapability(rag_service, collections=["pinned"])
|
cap = RagCapability(rag_service, collections=[FileId("pinned")])
|
||||||
output = await _invoke_search_knowledge(cap, "content")
|
output = await _invoke_search_knowledge(cap, "content")
|
||||||
assert "pinned.pdf" in output
|
assert "pinned.pdf" in output
|
||||||
assert "other.pdf" not in output
|
assert "other.pdf" not in output
|
||||||
@@ -273,7 +274,7 @@ class TestRagCapability:
|
|||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
|
async def test_search_knowledge_respects_max_results(self, rag_service: RagService) -> None:
|
||||||
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
|
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
|
||||||
await rag_service.index_text("bulk", paragraphs, source="bulk.pdf")
|
await rag_service.index_text(FileId("bulk"), paragraphs, source="bulk.pdf")
|
||||||
|
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(rag_service)
|
||||||
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
|
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
|
||||||
@@ -281,3 +282,27 @@ class TestRagCapability:
|
|||||||
assert "[Result 1" in output
|
assert "[Result 1" in output
|
||||||
assert "[Result 2" in output
|
assert "[Result 2" in output
|
||||||
assert "[Result 3" not in output
|
assert "[Result 3" not in output
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, rag_service: RagService) -> None:
|
||||||
|
"""The prepare callback must return None once max_searches has been reached
|
||||||
|
so the agent can no longer call the tool on subsequent turns."""
|
||||||
|
await rag_service.index_text(FileId("docs"), "Some content.", source="x.pdf")
|
||||||
|
cap = RagCapability(rag_service, max_searches=2)
|
||||||
|
tool_def = _dummy_tool_def()
|
||||||
|
|
||||||
|
# Budget intact: prepare returns the tool definition.
|
||||||
|
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
|
||||||
|
|
||||||
|
# Use the budget.
|
||||||
|
await _invoke_search_knowledge(cap, "content")
|
||||||
|
await _invoke_search_knowledge(cap, "content")
|
||||||
|
|
||||||
|
# Budget spent: prepare returns None, removing the tool from the agent's next turn.
|
||||||
|
assert await cap._prepare_search_knowledge(None, tool_def) is None # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
|
def _dummy_tool_def() -> object:
|
||||||
|
"""Sentinel passed to ``_prepare_search_knowledge``. The callback only inspects
|
||||||
|
``_search_count``; it doesn't read anything off the tool_def or context."""
|
||||||
|
return object()
|
||||||
|
|||||||
+108
-108
@@ -6,14 +6,13 @@ import pytest
|
|||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
from stirling.api import app
|
from stirling.api import app
|
||||||
from stirling.api.dependencies import get_rag_embedding_model, get_rag_service
|
from stirling.api.dependencies import get_rag_service
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.rag import Document, RagService, SqliteVecStore
|
from stirling.rag import Document, RagService, SqliteVecStore
|
||||||
|
|
||||||
TEST_EMBEDDING_MODEL = "test-embedder"
|
|
||||||
|
|
||||||
|
|
||||||
class StubEmbedder:
|
class StubEmbedder:
|
||||||
"""Deterministic embeddings for route tests — no network, no provider needed."""
|
"""Deterministic embeddings for route tests: no network, no provider needed."""
|
||||||
|
|
||||||
def __init__(self, dim: int = 8) -> None:
|
def __init__(self, dim: int = 8) -> None:
|
||||||
self._dim = dim
|
self._dim = dim
|
||||||
@@ -53,153 +52,154 @@ def _build_service() -> RagService:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def client() -> Iterator[TestClient]:
|
def service() -> RagService:
|
||||||
service = _build_service()
|
return _build_service()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def client(service: RagService) -> Iterator[TestClient]:
|
||||||
app.dependency_overrides[get_rag_service] = lambda: service
|
app.dependency_overrides[get_rag_service] = lambda: service
|
||||||
app.dependency_overrides[get_rag_embedding_model] = lambda: TEST_EMBEDDING_MODEL
|
|
||||||
try:
|
try:
|
||||||
yield TestClient(app)
|
yield TestClient(app)
|
||||||
finally:
|
finally:
|
||||||
app.dependency_overrides.pop(get_rag_service, None)
|
app.dependency_overrides.pop(get_rag_service, None)
|
||||||
app.dependency_overrides.pop(get_rag_embedding_model, None)
|
|
||||||
|
|
||||||
|
|
||||||
# ── /status ─────────────────────────────────────────────────────────────
|
# ── POST /documents ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
def test_status_reports_embedding_model_and_collections(client: TestClient) -> None:
|
def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None:
|
||||||
|
response = client.post(
|
||||||
|
"/api/v1/rag/documents",
|
||||||
|
json={
|
||||||
|
"documentId": "doc-123",
|
||||||
|
"source": "report.pdf",
|
||||||
|
"pageText": [
|
||||||
|
{"pageNumber": 1, "text": "The introduction covers the main topic."},
|
||||||
|
{"pageNumber": 2, "text": "The conclusion summarises the findings."},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert response.status_code == 200
|
||||||
|
body = response.json()
|
||||||
|
assert body["documentId"] == "doc-123"
|
||||||
|
assert body["chunksIndexed"] >= 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None:
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "my-docs", "text": "Hello world.", "source": "a.pdf"},
|
json={
|
||||||
|
"documentId": "replace-me",
|
||||||
|
"source": "replace-me.pdf",
|
||||||
|
"pageText": [{"pageNumber": 1, "text": "Original content that existed before."}],
|
||||||
|
},
|
||||||
)
|
)
|
||||||
response = client.get("/api/v1/rag/status")
|
# Second ingest with different content should replace the first entirely
|
||||||
assert response.status_code == 200
|
|
||||||
body = response.json()
|
|
||||||
assert body["embeddingModel"] == TEST_EMBEDDING_MODEL
|
|
||||||
assert "my-docs" in body["collections"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_status_when_empty(client: TestClient) -> None:
|
|
||||||
response = client.get("/api/v1/rag/status")
|
|
||||||
assert response.status_code == 200
|
|
||||||
body = response.json()
|
|
||||||
assert body == {"embeddingModel": TEST_EMBEDDING_MODEL, "collections": []}
|
|
||||||
|
|
||||||
|
|
||||||
# ── /index ──────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
|
|
||||||
def test_index_returns_chunk_count(client: TestClient) -> None:
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "indexed", "text": "Short text.", "source": "doc.pdf"},
|
json={
|
||||||
|
"documentId": "replace-me",
|
||||||
|
"source": "replace-me.pdf",
|
||||||
|
"pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}],
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
body = response.json()
|
|
||||||
assert body["collection"] == "indexed"
|
results = await service.search("New content", collection=FileId("replace-me"), top_k=5)
|
||||||
assert body["chunksIndexed"] >= 1
|
texts = [r.document.text for r in results]
|
||||||
|
assert any("New content" in t for t in texts)
|
||||||
|
assert not any("Original content" in t for t in texts)
|
||||||
|
|
||||||
|
|
||||||
def test_index_rejects_empty_collection_name(client: TestClient) -> None:
|
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "", "text": "Text.", "source": "x.pdf"},
|
json={
|
||||||
|
"documentId": "mixed",
|
||||||
|
"source": "mixed.pdf",
|
||||||
|
"pageText": [
|
||||||
|
{"pageNumber": 1, "text": " "},
|
||||||
|
{"pageNumber": 2, "text": "Real content on page 2."},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert response.json()["chunksIndexed"] >= 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
|
||||||
|
response = client.post("/api/v1/rag/documents", json={"documentId": "empty", "source": "empty.pdf"})
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert response.json()["chunksIndexed"] == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
||||||
|
response = client.post(
|
||||||
|
"/api/v1/rag/documents",
|
||||||
|
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
|
|
||||||
|
|
||||||
def test_index_rejects_oversized_text(client: TestClient) -> None:
|
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
||||||
huge = "x" * 1_000_001 # Just over the 1MB cap
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "toobig", "text": huge},
|
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
|
|
||||||
|
|
||||||
# ── /search ─────────────────────────────────────────────────────────────
|
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
||||||
|
|
||||||
|
|
||||||
def test_search_returns_results(client: TestClient) -> None:
|
|
||||||
client.post(
|
|
||||||
"/api/v1/rag/index",
|
|
||||||
json={"collection": "search-test", "text": "Python is fun.", "source": "guide.pdf"},
|
|
||||||
)
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/search",
|
"/api/v1/rag/documents",
|
||||||
json={"query": "Python", "collection": "search-test", "topK": 3},
|
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
|
||||||
assert response.status_code == 200
|
|
||||||
body = response.json()
|
|
||||||
assert body["query"] == "Python"
|
|
||||||
assert len(body["results"]) >= 1
|
|
||||||
first = body["results"][0]
|
|
||||||
assert first["source"] == "guide.pdf"
|
|
||||||
assert "score" in first
|
|
||||||
|
|
||||||
|
|
||||||
def test_search_rejects_empty_collection_name(client: TestClient) -> None:
|
|
||||||
response = client.post(
|
|
||||||
"/api/v1/rag/search",
|
|
||||||
json={"query": "anything", "collection": ""},
|
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
|
|
||||||
|
|
||||||
def test_search_without_collection_searches_all(client: TestClient) -> None:
|
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
|
||||||
client.post(
|
|
||||||
"/api/v1/rag/index",
|
|
||||||
json={"collection": "col-one", "text": "Alpha content.", "source": "one.pdf"},
|
|
||||||
)
|
|
||||||
client.post(
|
|
||||||
"/api/v1/rag/index",
|
|
||||||
json={"collection": "col-two", "text": "Beta content.", "source": "two.pdf"},
|
|
||||||
)
|
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/search",
|
"/api/v1/rag/documents",
|
||||||
json={"query": "content"},
|
json={
|
||||||
|
"documentId": "bad-page",
|
||||||
|
"source": "bad-page.pdf",
|
||||||
|
"pageText": [{"pageNumber": 0, "text": "something"}],
|
||||||
|
},
|
||||||
)
|
)
|
||||||
assert response.status_code == 200
|
assert response.status_code == 422
|
||||||
body = response.json()
|
|
||||||
assert len(body["results"]) >= 1
|
|
||||||
|
|
||||||
|
|
||||||
# ── /collections ────────────────────────────────────────────────────────
|
# ── DELETE /documents/{id} ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
def test_collections_empty_when_no_data(client: TestClient) -> None:
|
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
|
||||||
response = client.get("/api/v1/rag/collections")
|
|
||||||
assert response.status_code == 200
|
|
||||||
assert response.json() == {"collections": []}
|
|
||||||
|
|
||||||
|
|
||||||
def test_collections_lists_indexed(client: TestClient) -> None:
|
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "list-me", "text": "Text.", "source": "x.pdf"},
|
json={
|
||||||
|
"documentId": "to-delete",
|
||||||
|
"source": "to-delete.pdf",
|
||||||
|
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
||||||
|
},
|
||||||
)
|
)
|
||||||
response = client.get("/api/v1/rag/collections")
|
response = client.delete("/api/v1/rag/documents/to-delete")
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert "list-me" in response.json()["collections"]
|
assert response.json() == {"documentId": "to-delete", "deleted": True}
|
||||||
|
|
||||||
|
|
||||||
# ── DELETE /collections/{name} ──────────────────────────────────────────
|
def test_delete_document_is_idempotent(client: TestClient) -> None:
|
||||||
|
response = client.delete("/api/v1/rag/documents/never-existed")
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert response.json() == {"documentId": "never-existed", "deleted": False}
|
||||||
|
|
||||||
|
|
||||||
def test_delete_collection_removes_it(client: TestClient) -> None:
|
@pytest.mark.anyio
|
||||||
|
async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None:
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/index",
|
"/api/v1/rag/documents",
|
||||||
json={"collection": "to-delete", "text": "Text.", "source": "x.pdf"},
|
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
|
||||||
)
|
)
|
||||||
response = client.delete("/api/v1/rag/collections/to-delete")
|
assert await service.has_collection(FileId("gone"))
|
||||||
assert response.status_code == 200
|
client.delete("/api/v1/rag/documents/gone")
|
||||||
assert response.json() == {"status": "deleted", "collection": "to-delete"}
|
assert not await service.has_collection(FileId("gone"))
|
||||||
|
|
||||||
listing = client.get("/api/v1/rag/collections").json()
|
|
||||||
assert "to-delete" not in listing["collections"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_delete_nonexistent_collection_is_idempotent(client: TestClient) -> None:
|
|
||||||
response = client.delete("/api/v1/rag/collections/never-existed")
|
|
||||||
assert response.status_code == 200
|
|
||||||
assert response.json() == {"status": "deleted", "collection": "never-existed"}
|
|
||||||
|
|||||||
@@ -88,7 +88,10 @@ def test_health_route() -> None:
|
|||||||
|
|
||||||
|
|
||||||
def test_orchestrator_route() -> None:
|
def test_orchestrator_route() -> None:
|
||||||
response = client.post("/api/v1/orchestrator", json={"userMessage": "route this", "fileNames": ["test.pdf"]})
|
response = client.post(
|
||||||
|
"/api/v1/orchestrator",
|
||||||
|
json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
|
||||||
|
)
|
||||||
|
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert response.json()["outcome"] == "need_content"
|
assert response.json()["outcome"] == "need_content"
|
||||||
@@ -106,8 +109,7 @@ def test_pdf_questions_route() -> None:
|
|||||||
"/api/v1/pdf/questions",
|
"/api/v1/pdf/questions",
|
||||||
json={
|
json={
|
||||||
"question": "what is this?",
|
"question": "what is this?",
|
||||||
"fileNames": ["test.pdf"],
|
"files": [{"id": "test-id", "name": "test.pdf"}],
|
||||||
"pageText": [{"fileName": "test.pdf", "pages": [{"pageNumber": 1, "text": "Example"}]}],
|
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ from stirling.contracts import (
|
|||||||
AgentExecutionRequest,
|
AgentExecutionRequest,
|
||||||
AgentSpec,
|
AgentSpec,
|
||||||
AgentSpecStep,
|
AgentSpecStep,
|
||||||
|
AiFile,
|
||||||
EditPlanResponse,
|
EditPlanResponse,
|
||||||
ExecutionContext,
|
ExecutionContext,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
@@ -12,13 +13,14 @@ from stirling.contracts import (
|
|||||||
PdfTextSelection,
|
PdfTextSelection,
|
||||||
ToolOperationStep,
|
ToolOperationStep,
|
||||||
)
|
)
|
||||||
|
from stirling.models import FileId
|
||||||
from stirling.models.tool_models import Angle, RotatePdfParams, ToolEndpoint
|
from stirling.models.tool_models import Angle, RotatePdfParams, ToolEndpoint
|
||||||
|
|
||||||
|
|
||||||
def test_orchestrator_request_accepts_user_message() -> None:
|
def test_orchestrator_request_accepts_user_message() -> None:
|
||||||
request = OrchestratorRequest(
|
request = OrchestratorRequest(
|
||||||
user_message="Rotate the PDF",
|
user_message="Rotate the PDF",
|
||||||
file_names=["test.pdf"],
|
files=[AiFile(id=FileId("test-id"), name="test.pdf")],
|
||||||
artifacts=[
|
artifacts=[
|
||||||
ExtractedTextArtifact(
|
ExtractedTextArtifact(
|
||||||
files=[
|
files=[
|
||||||
@@ -89,6 +91,7 @@ def test_app_settings_accepts_model_configuration() -> None:
|
|||||||
rag_chunk_size=512,
|
rag_chunk_size=512,
|
||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
|
rag_max_searches=5,
|
||||||
max_pages=200,
|
max_pages=200,
|
||||||
max_characters=200_000,
|
max_characters=200_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
Reference in New Issue
Block a user