mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 10:34:06 +02:00
Add ability for Stirling engine to reason across large documents (#6314)
# Description of Changes Adds storage in the database for full document content alongside the RAG content (and changes the service to `DocumentService` instead of `RagService`). Then adds a generic capability that should be usable by any agent (currently just used by the Question Agent) which allows the agent to pull out the full contents of the doc, chunks it into various sections that will fit in the context window, and then processes them in parallel to create an intermediate result, and then processes the intermediate result into a final answer. It will re-chunk as many times as necessary to get the content small enough for the actual answer to be analysed (I've tested on PDFs ~3500 pages long, which is well above the context limit and requires maybe 3 rounds of compression to get an answer). The new full doc analysis stuff is heavier than the RAG lookup so both remain. The agents should use RAG for targeted info and the chunked reasoner for info that requires reading the full doc.
This commit is contained in:
+38
-4
@@ -30,6 +30,7 @@ import lombok.extern.slf4j.Slf4j;
|
|||||||
import stirling.software.common.model.job.ResultFile;
|
import stirling.software.common.model.job.ResultFile;
|
||||||
import stirling.software.common.service.JobOwnershipService;
|
import stirling.software.common.service.JobOwnershipService;
|
||||||
import stirling.software.common.service.TaskManager;
|
import stirling.software.common.service.TaskManager;
|
||||||
|
import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
|
||||||
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
|
||||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
|
||||||
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
|
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
|
||||||
@@ -140,13 +141,32 @@ public class AiEngineController {
|
|||||||
}
|
}
|
||||||
|
|
||||||
private void runOrchestrationStream(AiWorkflowRequest request, SseEmitter emitter) {
|
private void runOrchestrationStream(AiWorkflowRequest request, SseEmitter emitter) {
|
||||||
|
AiWorkflowService.ProgressListener listener =
|
||||||
|
new AiWorkflowService.ProgressListener() {
|
||||||
|
@Override
|
||||||
|
public void onProgress(AiWorkflowProgressEvent event) {
|
||||||
|
sendEvent(emitter, "progress", event);
|
||||||
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public void onHeartbeat() {
|
||||||
|
// Forward upstream heartbeats so the SSE pipe stays visibly alive between
|
||||||
|
// real progress events; if the frontend has gone away, sendEvent throws,
|
||||||
|
// which propagates up through the stream consumer and closes our upstream
|
||||||
|
// engine connection so the engine can cancel its in-flight workflow.
|
||||||
|
sendEvent(emitter, "heartbeat", Map.of());
|
||||||
|
}
|
||||||
|
};
|
||||||
try {
|
try {
|
||||||
AiWorkflowResponse result =
|
AiWorkflowResponse result = aiWorkflowService.orchestrate(request, listener);
|
||||||
aiWorkflowService.orchestrate(
|
|
||||||
request, progress -> sendEvent(emitter, "progress", progress));
|
|
||||||
registerFileResultAsJob(result);
|
registerFileResultAsJob(result);
|
||||||
sendEvent(emitter, "result", result);
|
sendEvent(emitter, "result", result);
|
||||||
emitter.complete();
|
emitter.complete();
|
||||||
|
} catch (ClientDisconnectedException e) {
|
||||||
|
// The frontend gave up mid-stream. The exception unwinding through orchestrate()
|
||||||
|
// already closed the upstream engine connection (engine sees disconnect and cancels).
|
||||||
|
// The emitter is already toast; nothing useful left to send.
|
||||||
|
log.debug("Client disconnected mid-stream; aborting workflow", e);
|
||||||
} catch (Exception e) {
|
} catch (Exception e) {
|
||||||
log.error("AI orchestration stream failed", e);
|
log.error("AI orchestration stream failed", e);
|
||||||
// Emit an error frame for the frontend and then complete normally. Using
|
// Emit an error frame for the frontend and then complete normally. Using
|
||||||
@@ -192,7 +212,21 @@ public class AiEngineController {
|
|||||||
try {
|
try {
|
||||||
emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
|
emitter.send(SseEmitter.event().name(name).data(data, MediaType.APPLICATION_JSON));
|
||||||
} catch (IOException e) {
|
} catch (IOException e) {
|
||||||
log.debug("Failed to send SSE event (client may have disconnected)", e);
|
// Surface the disconnect so the streaming pipeline unwinds: callers higher up close
|
||||||
|
// the upstream engine connection, which lets the engine cancel its in-flight workflow.
|
||||||
|
// Without this, the engine would keep producing (and billing for) tokens whose results
|
||||||
|
// nobody is reading.
|
||||||
|
throw new ClientDisconnectedException("Client disconnected from SSE stream", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Thrown by {@link #sendEvent} when the SSE emitter's underlying connection is gone. Treated as
|
||||||
|
* a signal to abort the workflow, not as an error to report.
|
||||||
|
*/
|
||||||
|
private static final class ClientDisconnectedException extends RuntimeException {
|
||||||
|
ClientDisconnectedException(String message, Throwable cause) {
|
||||||
|
super(message, cause);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+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/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 AiDocumentIngestRequest {
|
||||||
|
|
||||||
|
private String documentId;
|
||||||
|
|
||||||
|
private String source;
|
||||||
|
|
||||||
|
private List<AiPageText> pageText;
|
||||||
|
}
|
||||||
+59
@@ -0,0 +1,59 @@
|
|||||||
|
package stirling.software.proprietary.model.api.ai;
|
||||||
|
|
||||||
|
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||||
|
import com.fasterxml.jackson.annotation.JsonSubTypes;
|
||||||
|
import com.fasterxml.jackson.annotation.JsonTypeInfo;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Typed engine-emitted progress detail, mirroring the Python {@code ProgressEvent} discriminated
|
||||||
|
* union (see {@code engine/src/stirling/contracts/progress.py}). Carried inside {@link
|
||||||
|
* AiWorkflowProgressEvent#getEngineDetail()} for {@link AiWorkflowPhase#ENGINE_PROGRESS} events.
|
||||||
|
*
|
||||||
|
* <p>Sealed so adding a new engine-side phase forces a matching subtype on the Java side instead of
|
||||||
|
* silently passing through as an opaque map. The {@code phase} string is the discriminator and
|
||||||
|
* stays on the wire so the frontend (which doesn't know about Java's class hierarchy) can switch on
|
||||||
|
* it.
|
||||||
|
*/
|
||||||
|
@JsonTypeInfo(
|
||||||
|
use = JsonTypeInfo.Id.NAME,
|
||||||
|
include = JsonTypeInfo.As.EXISTING_PROPERTY,
|
||||||
|
property = "phase",
|
||||||
|
visible = true)
|
||||||
|
@JsonSubTypes({
|
||||||
|
@JsonSubTypes.Type(
|
||||||
|
value = AiEngineProgressDetail.WholeDocReadStarted.class,
|
||||||
|
name = "whole_doc_read_started"),
|
||||||
|
@JsonSubTypes.Type(
|
||||||
|
value = AiEngineProgressDetail.WholeDocSliceDone.class,
|
||||||
|
name = "whole_doc_slice_done"),
|
||||||
|
@JsonSubTypes.Type(
|
||||||
|
value = AiEngineProgressDetail.WholeDocCompressionRound.class,
|
||||||
|
name = "whole_doc_compression_round"),
|
||||||
|
@JsonSubTypes.Type(
|
||||||
|
value = AiEngineProgressDetail.WholeDocReadDone.class,
|
||||||
|
name = "whole_doc_read_done"),
|
||||||
|
})
|
||||||
|
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||||
|
public sealed interface AiEngineProgressDetail {
|
||||||
|
|
||||||
|
String phase();
|
||||||
|
|
||||||
|
record WholeDocReadStarted(String phase, String question, int pages, int slices)
|
||||||
|
implements AiEngineProgressDetail {}
|
||||||
|
|
||||||
|
record WholeDocSliceDone(
|
||||||
|
String phase,
|
||||||
|
int completed,
|
||||||
|
int total,
|
||||||
|
String pages,
|
||||||
|
int durationMs,
|
||||||
|
int excerpts,
|
||||||
|
int facts)
|
||||||
|
implements AiEngineProgressDetail {}
|
||||||
|
|
||||||
|
record WholeDocCompressionRound(String phase, int roundNumber, int notesIn, int groups)
|
||||||
|
implements AiEngineProgressDetail {}
|
||||||
|
|
||||||
|
record WholeDocReadDone(String phase, int completed, int slices, double durationSeconds)
|
||||||
|
implements AiEngineProgressDetail {}
|
||||||
|
}
|
||||||
+2
-2
@@ -4,11 +4,11 @@ import lombok.AllArgsConstructor;
|
|||||||
import lombok.Data;
|
import lombok.Data;
|
||||||
import lombok.NoArgsConstructor;
|
import lombok.NoArgsConstructor;
|
||||||
|
|
||||||
/** A single page of extracted text for RAG ingest requests. */
|
/** A single page of extracted text for document ingest requests. */
|
||||||
@Data
|
@Data
|
||||||
@NoArgsConstructor
|
@NoArgsConstructor
|
||||||
@AllArgsConstructor
|
@AllArgsConstructor
|
||||||
public class AiRagPageText {
|
public class AiPageText {
|
||||||
|
|
||||||
private int pageNumber;
|
private int pageNumber;
|
||||||
|
|
||||||
-24
@@ -1,24 +0,0 @@
|
|||||||
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;
|
|
||||||
}
|
|
||||||
+7
-1
@@ -9,7 +9,13 @@ public enum AiWorkflowPhase {
|
|||||||
CALLING_ENGINE("calling_engine"),
|
CALLING_ENGINE("calling_engine"),
|
||||||
EXTRACTING_CONTENT("extracting_content"),
|
EXTRACTING_CONTENT("extracting_content"),
|
||||||
EXECUTING_TOOL("executing_tool"),
|
EXECUTING_TOOL("executing_tool"),
|
||||||
PROCESSING("processing");
|
PROCESSING("processing"),
|
||||||
|
/**
|
||||||
|
* Generic engine-emitted progress event (e.g. chunked-reasoner slice progress). The original
|
||||||
|
* engine event JSON is carried in {@link AiWorkflowProgressEvent#getEngineDetail()}, including
|
||||||
|
* a specific {@code phase} string the frontend can switch on.
|
||||||
|
*/
|
||||||
|
ENGINE_PROGRESS("engine_progress");
|
||||||
|
|
||||||
private final String value;
|
private final String value;
|
||||||
|
|
||||||
|
|||||||
+22
-2
@@ -23,8 +23,17 @@ public class AiWorkflowProgressEvent {
|
|||||||
/** Total number of plan steps, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
|
/** Total number of plan steps, for {@link AiWorkflowPhase#EXECUTING_TOOL} events. */
|
||||||
private Integer stepCount;
|
private Integer stepCount;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Engine-emitted event payload, for {@link AiWorkflowPhase#ENGINE_PROGRESS} events. The payload
|
||||||
|
* is a typed subtype keyed on its {@code phase} string (e.g. {@code "whole_doc_slice_done"})
|
||||||
|
* carrying phase-specific fields (slice index, page range, durations, etc.) that the frontend
|
||||||
|
* can render as detailed progress.
|
||||||
|
*/
|
||||||
|
private AiEngineProgressDetail engineDetail;
|
||||||
|
|
||||||
public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
|
public static AiWorkflowProgressEvent of(AiWorkflowPhase phase) {
|
||||||
return new AiWorkflowProgressEvent(phase, System.currentTimeMillis(), null, null, null);
|
return new AiWorkflowProgressEvent(
|
||||||
|
phase, System.currentTimeMillis(), null, null, null, null);
|
||||||
}
|
}
|
||||||
|
|
||||||
public static AiWorkflowProgressEvent executingTool(String tool, int stepIndex, int stepCount) {
|
public static AiWorkflowProgressEvent executingTool(String tool, int stepIndex, int stepCount) {
|
||||||
@@ -33,6 +42,17 @@ public class AiWorkflowProgressEvent {
|
|||||||
System.currentTimeMillis(),
|
System.currentTimeMillis(),
|
||||||
tool,
|
tool,
|
||||||
stepIndex,
|
stepIndex,
|
||||||
stepCount);
|
stepCount,
|
||||||
|
null);
|
||||||
|
}
|
||||||
|
|
||||||
|
public static AiWorkflowProgressEvent engineProgress(AiEngineProgressDetail detail) {
|
||||||
|
return new AiWorkflowProgressEvent(
|
||||||
|
AiWorkflowPhase.ENGINE_PROGRESS,
|
||||||
|
System.currentTimeMillis(),
|
||||||
|
null,
|
||||||
|
null,
|
||||||
|
null,
|
||||||
|
detail);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+66
@@ -7,6 +7,8 @@ import java.net.http.HttpRequest;
|
|||||||
import java.net.http.HttpResponse;
|
import java.net.http.HttpResponse;
|
||||||
import java.net.http.HttpTimeoutException;
|
import java.net.http.HttpTimeoutException;
|
||||||
import java.time.Duration;
|
import java.time.Duration;
|
||||||
|
import java.util.function.Consumer;
|
||||||
|
import java.util.stream.Stream;
|
||||||
|
|
||||||
import org.springframework.beans.factory.annotation.Autowired;
|
import org.springframework.beans.factory.annotation.Autowired;
|
||||||
import org.springframework.http.HttpStatus;
|
import org.springframework.http.HttpStatus;
|
||||||
@@ -83,6 +85,70 @@ public class AiEngineClient {
|
|||||||
return response.body();
|
return response.body();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* POST a JSON body and consume the response as a stream of NDJSON lines. Each line is passed to
|
||||||
|
* {@code lineConsumer} in arrival order; the call returns when the engine closes the stream.
|
||||||
|
*
|
||||||
|
* <p>This is the right shape for long-running orchestrator calls that emit incremental
|
||||||
|
* progress. The total HTTP timeout is the long-running timeout (typically 600s+), but in
|
||||||
|
* practice line arrival keeps the connection logically alive: as long as the engine emits
|
||||||
|
* events, the work is progressing. Genuine engine hangs still hit the total timeout.
|
||||||
|
*/
|
||||||
|
public void streamPost(String path, String jsonBody, Consumer<String> lineConsumer)
|
||||||
|
throws IOException {
|
||||||
|
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||||
|
if (!config.isEnabled()) {
|
||||||
|
throw new ResponseStatusException(
|
||||||
|
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
|
||||||
|
}
|
||||||
|
|
||||||
|
String url = config.getUrl().stripTrailing() + path;
|
||||||
|
Duration timeout = Duration.ofSeconds(config.getLongRunningTimeoutSeconds());
|
||||||
|
log.debug(
|
||||||
|
"Proxying AI engine streaming request to {} (timeout {}s)",
|
||||||
|
url,
|
||||||
|
timeout.toSeconds());
|
||||||
|
|
||||||
|
HttpRequest request =
|
||||||
|
HttpRequest.newBuilder()
|
||||||
|
.uri(URI.create(url))
|
||||||
|
.header("Content-Type", "application/json")
|
||||||
|
.header("Accept", "application/x-ndjson")
|
||||||
|
.timeout(timeout)
|
||||||
|
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||||
|
.build();
|
||||||
|
|
||||||
|
HttpResponse<Stream<String>> response;
|
||||||
|
try {
|
||||||
|
response = httpClient.send(request, HttpResponse.BodyHandlers.ofLines());
|
||||||
|
} catch (HttpTimeoutException e) {
|
||||||
|
throw new ResponseStatusException(HttpStatus.GATEWAY_TIMEOUT, "AI engine timed out", e);
|
||||||
|
} catch (IOException e) {
|
||||||
|
throw new ResponseStatusException(
|
||||||
|
HttpStatus.SERVICE_UNAVAILABLE, "AI engine unreachable: " + e.getMessage(), e);
|
||||||
|
} catch (InterruptedException e) {
|
||||||
|
Thread.currentThread().interrupt();
|
||||||
|
throw new ResponseStatusException(
|
||||||
|
HttpStatus.SERVICE_UNAVAILABLE, "AI engine request was interrupted");
|
||||||
|
}
|
||||||
|
|
||||||
|
int status = response.statusCode();
|
||||||
|
if (status >= 400) {
|
||||||
|
throw new ResponseStatusException(
|
||||||
|
HttpStatus.valueOf(status >= 500 ? 502 : status),
|
||||||
|
"AI engine returned error: " + status);
|
||||||
|
}
|
||||||
|
|
||||||
|
try (Stream<String> lines = response.body()) {
|
||||||
|
lines.forEach(
|
||||||
|
line -> {
|
||||||
|
if (!line.isEmpty()) {
|
||||||
|
lineConsumer.accept(line);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
public String get(String path) throws IOException {
|
public String get(String path) throws IOException {
|
||||||
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
|
||||||
if (!config.isEnabled()) {
|
if (!config.isEnabled()) {
|
||||||
|
|||||||
+69
-13
@@ -37,9 +37,10 @@ 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.AiDocumentIngestRequest;
|
||||||
|
import stirling.software.proprietary.model.api.ai.AiEngineProgressDetail;
|
||||||
import stirling.software.proprietary.model.api.ai.AiFile;
|
import stirling.software.proprietary.model.api.ai.AiFile;
|
||||||
import stirling.software.proprietary.model.api.ai.AiRagIngestRequest;
|
import stirling.software.proprietary.model.api.ai.AiPageText;
|
||||||
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;
|
||||||
@@ -61,7 +62,7 @@ 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 static final String DOCUMENTS_ENDPOINT = "/api/v1/documents";
|
||||||
|
|
||||||
private final CustomPDFDocumentFactory pdfDocumentFactory;
|
private final CustomPDFDocumentFactory pdfDocumentFactory;
|
||||||
private final AiEngineClient aiEngineClient;
|
private final AiEngineClient aiEngineClient;
|
||||||
@@ -77,6 +78,14 @@ public class AiWorkflowService {
|
|||||||
@FunctionalInterface
|
@FunctionalInterface
|
||||||
public interface ProgressListener {
|
public interface ProgressListener {
|
||||||
void onProgress(AiWorkflowProgressEvent event);
|
void onProgress(AiWorkflowProgressEvent event);
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Called when the engine emits a keep-alive heartbeat. Default is a no-op; consumers that
|
||||||
|
* forward to a downstream connection (e.g. an SSE emitter) override this to push a
|
||||||
|
* heartbeat through, so the next downstream-disconnect surfaces immediately rather than
|
||||||
|
* waiting for the next real progress event.
|
||||||
|
*/
|
||||||
|
default void onHeartbeat() {}
|
||||||
}
|
}
|
||||||
|
|
||||||
private static final ProgressListener NOOP_LISTENER = event -> {};
|
private static final ProgressListener NOOP_LISTENER = event -> {};
|
||||||
@@ -144,7 +153,7 @@ public class AiWorkflowService {
|
|||||||
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, listener);
|
||||||
return switch (response.getOutcome()) {
|
return switch (response.getOutcome()) {
|
||||||
case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
|
case NEED_CONTENT -> onNeedContent(response, filesById, request, listener);
|
||||||
case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
|
case NEED_INGEST -> onNeedIngest(response, filesById, request, listener);
|
||||||
@@ -278,22 +287,22 @@ public class AiWorkflowService {
|
|||||||
}
|
}
|
||||||
|
|
||||||
private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException {
|
private void ingestFile(AiFile file, MultipartFile multipartFile) throws IOException {
|
||||||
List<AiRagPageText> pages = new ArrayList<>();
|
List<AiPageText> pages = new ArrayList<>();
|
||||||
try (PDDocument document = pdfDocumentFactory.load(multipartFile, true)) {
|
try (PDDocument document = pdfDocumentFactory.load(multipartFile, true)) {
|
||||||
int pageCount = document.getNumberOfPages();
|
int pageCount = document.getNumberOfPages();
|
||||||
for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
|
for (int pageNumber = 1; pageNumber <= pageCount; pageNumber++) {
|
||||||
String pageText = pdfContentExtractor.extractPageTextRaw(document, pageNumber);
|
String pageText = pdfContentExtractor.extractPageTextRaw(document, pageNumber);
|
||||||
if (pageText != null && !pageText.isBlank()) {
|
if (pageText != null && !pageText.isBlank()) {
|
||||||
pages.add(new AiRagPageText(pageNumber, pageText));
|
pages.add(new AiPageText(pageNumber, pageText));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
AiRagIngestRequest ingestRequest =
|
AiDocumentIngestRequest ingestRequest =
|
||||||
new AiRagIngestRequest(file.getId(), file.getName(), pages);
|
new AiDocumentIngestRequest(file.getId(), file.getName(), pages);
|
||||||
String body = objectMapper.writeValueAsString(ingestRequest);
|
String body = objectMapper.writeValueAsString(ingestRequest);
|
||||||
aiEngineClient.postLongRunning(RAG_DOCUMENTS_ENDPOINT, body);
|
aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body);
|
||||||
log.debug(
|
log.debug(
|
||||||
"Ingested file into RAG: id={}, name={}, pages={}",
|
"Ingested document: id={}, name={}, pages={}",
|
||||||
file.getId(),
|
file.getId(),
|
||||||
file.getName(),
|
file.getName(),
|
||||||
pages.size());
|
pages.size());
|
||||||
@@ -651,10 +660,57 @@ public class AiWorkflowService {
|
|||||||
return response;
|
return response;
|
||||||
}
|
}
|
||||||
|
|
||||||
private AiWorkflowResponse invokeOrchestrator(WorkflowTurnRequest request) throws IOException {
|
/**
|
||||||
|
* Drive the engine's streaming orchestrator endpoint. Progress events are forwarded to {@code
|
||||||
|
* listener} as they arrive (each one keeps the SSE connection to the frontend alive too). The
|
||||||
|
* final {@code result} event carries the full {@link AiWorkflowResponse}; an {@code error}
|
||||||
|
* event surfaces engine-side failures.
|
||||||
|
*/
|
||||||
|
private AiWorkflowResponse invokeOrchestrator(
|
||||||
|
WorkflowTurnRequest request, ProgressListener listener) throws IOException {
|
||||||
String requestBody = objectMapper.writeValueAsString(request);
|
String requestBody = objectMapper.writeValueAsString(request);
|
||||||
String responseBody = aiEngineClient.post("/api/v1/orchestrator", requestBody);
|
AiWorkflowResponse[] resultHolder = new AiWorkflowResponse[1];
|
||||||
return objectMapper.readValue(responseBody, AiWorkflowResponse.class);
|
String[] errorHolder = new String[1];
|
||||||
|
|
||||||
|
aiEngineClient.streamPost(
|
||||||
|
"/api/v1/orchestrator",
|
||||||
|
requestBody,
|
||||||
|
line -> handleStreamLine(line, listener, resultHolder, errorHolder));
|
||||||
|
|
||||||
|
if (errorHolder[0] != null) {
|
||||||
|
throw new IOException("AI engine returned error: " + errorHolder[0]);
|
||||||
|
}
|
||||||
|
if (resultHolder[0] == null) {
|
||||||
|
throw new IOException("AI engine stream ended without a result");
|
||||||
|
}
|
||||||
|
return resultHolder[0];
|
||||||
|
}
|
||||||
|
|
||||||
|
private void handleStreamLine(
|
||||||
|
String line,
|
||||||
|
ProgressListener listener,
|
||||||
|
AiWorkflowResponse[] resultHolder,
|
||||||
|
String[] errorHolder) {
|
||||||
|
try {
|
||||||
|
JsonNode node = objectMapper.readTree(line);
|
||||||
|
String event = node.path("event").asText();
|
||||||
|
switch (event) {
|
||||||
|
case "progress" -> {
|
||||||
|
AiEngineProgressDetail detail =
|
||||||
|
objectMapper.treeToValue(node, AiEngineProgressDetail.class);
|
||||||
|
listener.onProgress(AiWorkflowProgressEvent.engineProgress(detail));
|
||||||
|
}
|
||||||
|
case "result" -> {
|
||||||
|
JsonNode response = node.path("response");
|
||||||
|
resultHolder[0] = objectMapper.treeToValue(response, AiWorkflowResponse.class);
|
||||||
|
}
|
||||||
|
case "error" -> errorHolder[0] = node.path("message").asText("unknown error");
|
||||||
|
case "heartbeat" -> listener.onHeartbeat();
|
||||||
|
default -> log.warn("Ignoring unknown engine stream event: {}", event);
|
||||||
|
}
|
||||||
|
} catch (JacksonException e) {
|
||||||
|
log.warn("Failed to parse engine stream line: {}", line, e);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@Data
|
@Data
|
||||||
|
|||||||
+48
-19
@@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
|
|||||||
import static org.mockito.ArgumentMatchers.anyInt;
|
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.doAnswer;
|
||||||
import static org.mockito.Mockito.lenient;
|
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;
|
||||||
@@ -22,6 +23,7 @@ import java.nio.file.Path;
|
|||||||
import java.util.ArrayList;
|
import java.util.ArrayList;
|
||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.concurrent.atomic.AtomicInteger;
|
import java.util.concurrent.atomic.AtomicInteger;
|
||||||
|
import java.util.function.Consumer;
|
||||||
import java.util.zip.ZipEntry;
|
import java.util.zip.ZipEntry;
|
||||||
import java.util.zip.ZipOutputStream;
|
import java.util.zip.ZipOutputStream;
|
||||||
|
|
||||||
@@ -419,7 +421,7 @@ class AiWorkflowServiceTest {
|
|||||||
}
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
void needIngestExtractsPageTextAndPostsToRagThenRetries() throws IOException {
|
void needIngestExtractsPageTextAndPostsThenRetries() throws IOException {
|
||||||
MockMultipartFile input = pdf("report.pdf", "bytes");
|
MockMultipartFile input = pdf("report.pdf", "bytes");
|
||||||
when(fileIdStrategy.idFor(any())).thenReturn("report-id");
|
when(fileIdStrategy.idFor(any())).thenReturn("report-id");
|
||||||
|
|
||||||
@@ -431,37 +433,64 @@ class AiWorkflowServiceTest {
|
|||||||
.thenReturn("page content");
|
.thenReturn("page content");
|
||||||
|
|
||||||
int[] orchestratorCalls = {0};
|
int[] orchestratorCalls = {0};
|
||||||
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString()))
|
doAnswer(
|
||||||
.thenAnswer(
|
|
||||||
inv -> {
|
inv -> {
|
||||||
orchestratorCalls[0]++;
|
orchestratorCalls[0]++;
|
||||||
|
String responseJson;
|
||||||
if (orchestratorCalls[0] == 1) {
|
if (orchestratorCalls[0] == 1) {
|
||||||
return """
|
responseJson =
|
||||||
{
|
"""
|
||||||
"outcome":"need_ingest",
|
{
|
||||||
"resumeWith":"pdf_question",
|
"outcome":"need_ingest",
|
||||||
"reason":"ingest first",
|
"resumeWith":"pdf_question",
|
||||||
"filesToIngest":[{"id":"report-id","name":"report.pdf"}],
|
"reason":"ingest first",
|
||||||
"contentTypes":["page_text"]
|
"filesToIngest":[{"id":"report-id","name":"report.pdf"}],
|
||||||
}
|
"contentTypes":["page_text"]
|
||||||
""";
|
}
|
||||||
|
""";
|
||||||
|
} else {
|
||||||
|
responseJson =
|
||||||
|
"""
|
||||||
|
{"outcome":"answer","answer":"done","evidence":[]}
|
||||||
|
""";
|
||||||
}
|
}
|
||||||
return """
|
Consumer<String> consumer = inv.getArgument(2);
|
||||||
{"outcome":"answer","answer":"done","evidence":[]}
|
consumer.accept(wrapAsResultEvent(responseJson));
|
||||||
""";
|
return null;
|
||||||
});
|
})
|
||||||
|
.when(aiEngineClient)
|
||||||
|
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
|
||||||
|
|
||||||
AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
|
AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
|
||||||
|
|
||||||
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
|
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
|
||||||
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/rag/documents"), anyString());
|
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString());
|
||||||
verify(aiEngineClient, times(2)).post(eq("/api/v1/orchestrator"), anyString());
|
verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any());
|
||||||
}
|
}
|
||||||
|
|
||||||
// --- helpers ---
|
// --- helpers ---
|
||||||
|
|
||||||
private void stubOrchestrator(String responseJson) throws IOException {
|
private void stubOrchestrator(String responseJson) throws IOException {
|
||||||
when(aiEngineClient.post(eq("/api/v1/orchestrator"), anyString())).thenReturn(responseJson);
|
doAnswer(
|
||||||
|
inv -> {
|
||||||
|
Consumer<String> consumer = inv.getArgument(2);
|
||||||
|
consumer.accept(wrapAsResultEvent(responseJson));
|
||||||
|
return null;
|
||||||
|
})
|
||||||
|
.when(aiEngineClient)
|
||||||
|
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Wrap a unary orchestrator response JSON as the NDJSON "result" event the streaming endpoint
|
||||||
|
* emits, so existing tests can keep their per-outcome JSON literals.
|
||||||
|
*/
|
||||||
|
private String wrapAsResultEvent(String responseJson) throws IOException {
|
||||||
|
return objectMapper
|
||||||
|
.createObjectNode()
|
||||||
|
.put("event", "result")
|
||||||
|
.set("response", objectMapper.readTree(responseJson))
|
||||||
|
.toString();
|
||||||
}
|
}
|
||||||
|
|
||||||
private void stubEndpoint(String endpoint, Resource body) {
|
private void stubEndpoint(String endpoint, Resource body) {
|
||||||
|
|||||||
+21
@@ -37,6 +37,22 @@ STIRLING_RAG_TOP_K=20
|
|||||||
# rather than chain more searches.
|
# rather than chain more searches.
|
||||||
STIRLING_RAG_MAX_SEARCHES=5
|
STIRLING_RAG_MAX_SEARCHES=5
|
||||||
|
|
||||||
|
# Chunked reasoner settings: how big each per-worker slice is (in characters),
|
||||||
|
# how many workers may run in parallel against the fast model, and how long
|
||||||
|
# any single worker is allowed to wait for a response before being abandoned.
|
||||||
|
# Worker timeouts protect gather_notes from upstream model stalls (which
|
||||||
|
# otherwise hang at the provider's ~10 minute HTTP default); the affected
|
||||||
|
# slice is dropped and the rest of the document still answers.
|
||||||
|
STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE=16000
|
||||||
|
STIRLING_CHUNKED_REASONER_CONCURRENCY=10
|
||||||
|
STIRLING_CHUNKED_REASONER_WORKER_TIMEOUT_SECONDS=60
|
||||||
|
|
||||||
|
# When the rendered slice notes would exceed this many characters, the
|
||||||
|
# reasoner folds them hierarchically with fast-model calls until they fit.
|
||||||
|
# This keeps the synthesis prompt under the model's context limit on long
|
||||||
|
# documents (a 3000-page novel produces ~900k chars of raw notes).
|
||||||
|
STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET=250000
|
||||||
|
|
||||||
# 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
|
||||||
STIRLING_MAX_CHARACTERS=200000
|
STIRLING_MAX_CHARACTERS=200000
|
||||||
@@ -51,3 +67,8 @@ STIRLING_LOG_LEVEL=INFO
|
|||||||
|
|
||||||
# Path to log file. Rolls daily, keeps 1 backup. Leave empty for console only.
|
# Path to log file. Rolls daily, keeps 1 backup. Leave empty for console only.
|
||||||
STIRLING_LOG_FILE=
|
STIRLING_LOG_FILE=
|
||||||
|
|
||||||
|
# Set true to log every outgoing httpx / Anthropic SDK request with timing.
|
||||||
|
# Use when diagnosing worker stalls: a hung call shows a "Request" line with
|
||||||
|
# no matching "Response" line. Noisy; leave off in normal use.
|
||||||
|
STIRLING_HTTP_DEBUG=false
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ from pydantic_ai import Agent
|
|||||||
from pydantic_ai.output import NativeOutput
|
from pydantic_ai.output import NativeOutput
|
||||||
|
|
||||||
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
|
from stirling.agents.math_presentation import MathIntentClassifier, extract_math_verdict
|
||||||
|
from stirling.agents.shared import ChunkedReasoner, WholeDocReaderCapability
|
||||||
from stirling.contracts import (
|
from stirling.contracts import (
|
||||||
AiFile,
|
AiFile,
|
||||||
EditPlanResponse,
|
EditPlanResponse,
|
||||||
@@ -24,44 +25,46 @@ from stirling.contracts import (
|
|||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
format_file_names,
|
format_file_names,
|
||||||
)
|
)
|
||||||
|
from stirling.documents import RagCapability
|
||||||
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
PDF_QUESTION_SYSTEM_PROMPT = (
|
PDF_QUESTION_SYSTEM_PROMPT = (
|
||||||
"You answer questions about PDF documents by retrieving relevant content with the "
|
"You answer questions about PDF documents using two retrieval tools:\n"
|
||||||
"search_knowledge tool. Use it before answering. Do not guess or use outside knowledge.\n"
|
|
||||||
"\n"
|
"\n"
|
||||||
"The search_knowledge tool has a finite call budget per run. When it is no longer "
|
"1. search_knowledge(query) - returns the passages most semantically similar "
|
||||||
"available, answer from what you have already retrieved.\n"
|
"to the query. Use it for targeted lookups: a specific fact, a named section, "
|
||||||
|
"a particular passage. Typically one or two calls is enough.\n"
|
||||||
|
"\n"
|
||||||
|
"2. read_full_document(query) - reads every page of the attached documents in "
|
||||||
|
"parallel and returns notes relevant to the query. Use it when answering "
|
||||||
|
"requires seeing the whole document end-to-end: summaries, aggregations "
|
||||||
|
"(largest, shortest, count), comparisons across sections. It is more "
|
||||||
|
"expensive than search_knowledge, so prefer search_knowledge when one or two "
|
||||||
|
"passages would suffice.\n"
|
||||||
|
"\n"
|
||||||
|
"Pick the right tool, call it, then answer from what you got back. Do not "
|
||||||
|
"guess or use outside knowledge.\n"
|
||||||
"\n"
|
"\n"
|
||||||
"Guidelines:\n"
|
"Guidelines:\n"
|
||||||
"- Make targeted search_knowledge calls. Typically one or two is enough.\n"
|
"- If the retrieved content does not support a confident answer, return not_found.\n"
|
||||||
"- Answer from the retrieved text. If the retrieved content doesn't support a confident "
|
"- Include a short list of evidence snippets (with page numbers where available) "
|
||||||
"answer, return not_found.\n"
|
"drawn from what the tools returned.\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"
|
"\n"
|
||||||
"Writing the not_found reason:\n"
|
"Writing the not_found reason:\n"
|
||||||
"- The reason is shown directly to the end user, so write it in plain, friendly "
|
"- The reason is shown directly to the end user, so write it in plain, friendly "
|
||||||
"language. One or two short sentences.\n"
|
"language. One or two short sentences.\n"
|
||||||
"- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
|
"- NEVER mention 'RAG', 'retrieval', 'chunks', 'search results', 'targeted search', "
|
||||||
"'search_knowledge', or other implementation details.\n"
|
"'search_knowledge', 'read_full_document', 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: "
|
"- For questions where the answer just isn't in the document, say so directly: "
|
||||||
"'I couldn't find that information in the document.'\n"
|
"'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 "
|
"- Do not make it sound like you're choosing not to answer."
|
||||||
"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 "
|
||||||
@@ -83,6 +86,9 @@ class PdfQuestionAgent:
|
|||||||
model_settings=runtime.fast_model_settings,
|
model_settings=runtime.fast_model_settings,
|
||||||
)
|
)
|
||||||
self._math_intent_classifier = MathIntentClassifier(runtime)
|
self._math_intent_classifier = MathIntentClassifier(runtime)
|
||||||
|
# Shared across whole-doc-reader instances so the worker agent and
|
||||||
|
# semaphore are constructed once and reused per request.
|
||||||
|
self._chunked_reasoner = ChunkedReasoner(runtime)
|
||||||
|
|
||||||
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
async def handle(self, request: PdfQuestionRequest) -> PdfQuestionResponse:
|
||||||
logger.info(
|
logger.info(
|
||||||
@@ -95,7 +101,7 @@ class PdfQuestionAgent:
|
|||||||
logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
|
logger.info("[pdf-question] missing ingestions: %s", [file.name for file in missing])
|
||||||
return NeedIngestResponse(
|
return NeedIngestResponse(
|
||||||
resume_with=SupportedCapability.PDF_QUESTION,
|
resume_with=SupportedCapability.PDF_QUESTION,
|
||||||
reason="Some files have not been ingested into RAG yet.",
|
reason="Some files have not been ingested yet.",
|
||||||
files_to_ingest=missing,
|
files_to_ingest=missing,
|
||||||
content_types=[PdfContentType.PAGE_TEXT],
|
content_types=[PdfContentType.PAGE_TEXT],
|
||||||
)
|
)
|
||||||
@@ -145,23 +151,37 @@ class PdfQuestionAgent:
|
|||||||
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
|
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
|
||||||
missing: list[AiFile] = []
|
missing: list[AiFile] = []
|
||||||
for file in files:
|
for file in files:
|
||||||
if not await self.runtime.rag_service.has_collection(file.id):
|
if not await self.runtime.documents.has_collection(file.id):
|
||||||
missing.append(file)
|
missing.append(file)
|
||||||
return missing
|
return missing
|
||||||
|
|
||||||
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
|
async def _run_answer_agent(self, request: PdfQuestionRequest) -> PdfQuestionTerminalResponse:
|
||||||
|
"""Drive a single smart-model agent with both retrieval tools.
|
||||||
|
|
||||||
|
The agent picks ``search_knowledge`` for targeted lookups and
|
||||||
|
``read_full_document`` for whole-document questions. Removing the
|
||||||
|
upstream classifier keeps that judgement in the same call that writes
|
||||||
|
the answer, and lets the agent mix tools when the question warrants it.
|
||||||
|
"""
|
||||||
rag = RagCapability(
|
rag = RagCapability(
|
||||||
rag_service=self.runtime.rag_service,
|
documents=self.runtime.documents,
|
||||||
collections=[file.id for file in request.files],
|
collections=[file.id for file in request.files],
|
||||||
top_k=self.runtime.settings.rag_default_top_k,
|
top_k=self.runtime.settings.rag_default_top_k,
|
||||||
max_searches=self.runtime.settings.rag_max_searches,
|
max_searches=self.runtime.settings.rag_max_searches,
|
||||||
)
|
)
|
||||||
|
whole_doc = WholeDocReaderCapability(
|
||||||
|
runtime=self.runtime,
|
||||||
|
files=request.files,
|
||||||
|
reasoner=self._chunked_reasoner,
|
||||||
|
)
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
model=self.runtime.smart_model,
|
model=self.runtime.smart_model,
|
||||||
output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
|
output_type=NativeOutput([PdfQuestionAnswerResponse, PdfQuestionNotFoundResponse]),
|
||||||
system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
|
system_prompt=PDF_QUESTION_SYSTEM_PROMPT,
|
||||||
instructions=rag.instructions,
|
# pydantic-ai accepts a list of (string-or-callable) instruction sources;
|
||||||
toolsets=[rag.toolset],
|
# it resolves each at run time and concatenates them for the model.
|
||||||
|
instructions=[rag.instructions, whole_doc.instructions],
|
||||||
|
toolsets=[rag.toolset, whole_doc.toolset],
|
||||||
model_settings=self.runtime.smart_model_settings,
|
model_settings=self.runtime.smart_model_settings,
|
||||||
)
|
)
|
||||||
prompt = self._build_prompt(request)
|
prompt = self._build_prompt(request)
|
||||||
@@ -184,5 +204,5 @@ class PdfQuestionAgent:
|
|||||||
f"Conversation history:\n{history}\n"
|
f"Conversation history:\n{history}\n"
|
||||||
f"Files: {format_file_names(request.files)}\n"
|
f"Files: {format_file_names(request.files)}\n"
|
||||||
f"Question: {request.question}\n"
|
f"Question: {request.question}\n"
|
||||||
"Use search_knowledge to retrieve the relevant content, then answer."
|
"Pick the right retrieval tool for this question, then answer from what it returns."
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -0,0 +1,10 @@
|
|||||||
|
"""Reasoning utilities shared across agents."""
|
||||||
|
|
||||||
|
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
|
||||||
|
from stirling.agents.shared.whole_doc_reader import WholeDocReaderCapability
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ChunkNotes",
|
||||||
|
"ChunkedReasoner",
|
||||||
|
"WholeDocReaderCapability",
|
||||||
|
]
|
||||||
@@ -0,0 +1,616 @@
|
|||||||
|
"""Chunked reasoning over long documents.
|
||||||
|
|
||||||
|
A reusable primitive for any agent that needs to answer a question that
|
||||||
|
requires reading a whole document end-to-end. The document is split into
|
||||||
|
character-budgeted chunks; each chunk is read by a parallel worker that
|
||||||
|
extracts question-relevant notes; if the gathered notes overflow the
|
||||||
|
synthesis context budget, the resulting notes are regrouped into fresh
|
||||||
|
chunks and run through the same extractor again, until they fit.
|
||||||
|
|
||||||
|
Pages are tracked by the wrapper, never asked of the model: keeps the model
|
||||||
|
output schema small and the page list authoritative.
|
||||||
|
|
||||||
|
Used wherever pure RAG retrieval is the wrong tool: aggregations ("largest
|
||||||
|
number"), comparisons ("shortest chapter"), and full summaries.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from pydantic_ai import Agent
|
||||||
|
from pydantic_ai.output import NativeOutput
|
||||||
|
|
||||||
|
from stirling.contracts import (
|
||||||
|
WholeDocCompressionRound,
|
||||||
|
WholeDocReadDone,
|
||||||
|
WholeDocReadStarted,
|
||||||
|
WholeDocSliceDone,
|
||||||
|
)
|
||||||
|
from stirling.contracts.documents import Page
|
||||||
|
from stirling.models import ApiModel
|
||||||
|
from stirling.services import AppRuntime, emit_progress
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ChunkNotes(ApiModel):
|
||||||
|
"""Public-facing notes for a span of pages.
|
||||||
|
|
||||||
|
Returned to callers of :meth:`ChunkedReasoner.gather_notes` and to the
|
||||||
|
inside of :meth:`ChunkedReasoner.reason`. The wrapper builds these from
|
||||||
|
the model's :class:`_ExtractedNotes` output and a deterministic page list.
|
||||||
|
"""
|
||||||
|
|
||||||
|
pages: list[int] = Field(description="Page numbers covered by these notes (1-indexed).")
|
||||||
|
summary: str = Field(description="One- to three-sentence summary of the covered range.")
|
||||||
|
relevant_excerpts: list[str] = Field(
|
||||||
|
default_factory=list,
|
||||||
|
description="Short verbatim quotes from the source content that bear on the user's question.",
|
||||||
|
)
|
||||||
|
facts: list[str] = Field(
|
||||||
|
default_factory=list,
|
||||||
|
description=(
|
||||||
|
"Concrete facts (numbers, names, dates, claims) the synthesiser may need. "
|
||||||
|
"Includes candidate values for aggregation questions."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _ExtractedNotes(BaseModel):
|
||||||
|
"""Model output for one extractor call.
|
||||||
|
|
||||||
|
No ``pages`` field: page numbers are mechanical aggregation the wrapper
|
||||||
|
computes deterministically. Keeping them out of the schema saves output
|
||||||
|
tokens for the bulkier excerpts/facts payload and prevents the model
|
||||||
|
from misreporting page coverage.
|
||||||
|
"""
|
||||||
|
|
||||||
|
summary: str = Field(description="One- to three-sentence summary of the supplied content.")
|
||||||
|
relevant_excerpts: list[str] = Field(
|
||||||
|
default_factory=list,
|
||||||
|
description=(
|
||||||
|
"Short verbatim quotes drawn from the supplied content that bear on the question. "
|
||||||
|
"Deduplicate; drop ones that don't bear on the question."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
facts: list[str] = Field(
|
||||||
|
default_factory=list,
|
||||||
|
description=(
|
||||||
|
"Distinct, deduplicated facts (numbers, names, dates, claims) needed to answer "
|
||||||
|
"the question. For aggregation questions retain ALL candidate values across the "
|
||||||
|
"supplied content so a later round can still pick the global winner."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class _Chunk:
|
||||||
|
"""A unit of work for the extractor: content + the pages it covers + a fallback.
|
||||||
|
|
||||||
|
``content`` is the formatted text fed to the model: raw page text with
|
||||||
|
``[Page N]`` markers in the first round, formatted prior-pass notes with
|
||||||
|
``[Notes from pages A-B]`` markers in subsequent rounds. ``pages`` is
|
||||||
|
attached to the resulting :class:`ChunkNotes` deterministically.
|
||||||
|
|
||||||
|
``fallback`` is the list of notes to keep if the extractor call fails. For
|
||||||
|
raw page chunks it's empty (a failed slice has no pre-extracted notes to
|
||||||
|
preserve). For chunks built from existing notes it's the input notes
|
||||||
|
themselves, so a failure doesn't lose page coverage.
|
||||||
|
"""
|
||||||
|
|
||||||
|
content: str
|
||||||
|
pages: list[int]
|
||||||
|
fallback: list[ChunkNotes]
|
||||||
|
label: str
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class _RoundResult:
|
||||||
|
"""Outcome of one extraction round.
|
||||||
|
|
||||||
|
``successes`` lets the loop detect rounds that made no forward progress
|
||||||
|
(every chunk failed) and bail rather than spinning. ``slowest`` is the
|
||||||
|
chunk with the longest successful extractor call this round, used for
|
||||||
|
diagnostic log lines on the first round.
|
||||||
|
"""
|
||||||
|
|
||||||
|
notes: list[ChunkNotes]
|
||||||
|
successes: int
|
||||||
|
slowest: tuple[str, float] | None
|
||||||
|
|
||||||
|
|
||||||
|
def _page_range_label(pages: list[Page]) -> str:
|
||||||
|
if not pages:
|
||||||
|
return "pages=?"
|
||||||
|
elif len(pages) == 1:
|
||||||
|
return f"pages={pages[0].page_number}"
|
||||||
|
else:
|
||||||
|
return f"pages={pages[0].page_number}-{pages[-1].page_number}"
|
||||||
|
|
||||||
|
|
||||||
|
def _note_range_label(notes: list[ChunkNotes]) -> str:
|
||||||
|
"""Render a "pages=A-B" label for a group of already-extracted notes."""
|
||||||
|
page_numbers = sorted({p for note in notes for p in note.pages})
|
||||||
|
if not page_numbers:
|
||||||
|
return "pages=?"
|
||||||
|
if len(page_numbers) == 1:
|
||||||
|
return f"pages={page_numbers[0]}"
|
||||||
|
return f"pages={page_numbers[0]}-{page_numbers[-1]}"
|
||||||
|
|
||||||
|
|
||||||
|
_EXTRACTOR_SYSTEM_PROMPT = (
|
||||||
|
"You are reading content from a document - either raw page text or "
|
||||||
|
"condensed notes from an earlier extraction pass - and your job is to "
|
||||||
|
"produce a tight set of notes that captures everything relevant to the "
|
||||||
|
"user's question. The same job runs many times in parallel across the "
|
||||||
|
"document and may run again to consolidate notes into smaller batches, "
|
||||||
|
"so be thorough: anything you skip cannot be recovered later.\n"
|
||||||
|
"\n"
|
||||||
|
"Output:\n"
|
||||||
|
"- summary: 1-3 sentences covering the supplied content.\n"
|
||||||
|
"- relevant_excerpts: short verbatim quotes from the supplied content "
|
||||||
|
"that bear on the question. Deduplicate; drop quotes that don't help.\n"
|
||||||
|
"- facts: concrete facts (numbers, names, dates, claims). Deduplicate; "
|
||||||
|
"drop irrelevant ones. For aggregation questions (largest, smallest, "
|
||||||
|
"count, total) retain ALL candidate values across the content so a "
|
||||||
|
"later step can still pick the global winner.\n"
|
||||||
|
"\n"
|
||||||
|
"Stay grounded in the supplied content. Do not infer or fabricate "
|
||||||
|
"anything that isn't already present. If nothing in the content is "
|
||||||
|
"relevant to the question, return empty excerpts and facts and a short "
|
||||||
|
"neutral summary."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class ChunkedReasoner:
|
||||||
|
"""Run a question against a long document by chunking, mapping, and looping.
|
||||||
|
|
||||||
|
Two consumption styles:
|
||||||
|
|
||||||
|
* Tools that already have a synthesising LLM call upstream call
|
||||||
|
:meth:`gather_notes` to get the structured notes and format them
|
||||||
|
themselves with :meth:`format_notes`.
|
||||||
|
* Callers that just want an answer call :meth:`reason`, which runs
|
||||||
|
:meth:`gather_notes` and then a single synthesis call governed by the
|
||||||
|
caller's ``answer_prompt`` and ``answer_type``.
|
||||||
|
|
||||||
|
Lifetime:
|
||||||
|
Construct once per agent that uses it. The extractor agent is built
|
||||||
|
at construction time and reused; the synthesis agent in :meth:`reason`
|
||||||
|
is built per call because its output type is generic.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
runtime: AppRuntime,
|
||||||
|
*,
|
||||||
|
chars_per_slice: int | None = None,
|
||||||
|
concurrency: int | None = None,
|
||||||
|
worker_timeout_seconds: float | None = None,
|
||||||
|
notes_char_budget: int | None = None,
|
||||||
|
) -> None:
|
||||||
|
chars = chars_per_slice if chars_per_slice is not None else runtime.settings.chunked_reasoner_chars_per_slice
|
||||||
|
conc = concurrency if concurrency is not None else runtime.settings.chunked_reasoner_concurrency
|
||||||
|
timeout = (
|
||||||
|
worker_timeout_seconds
|
||||||
|
if worker_timeout_seconds is not None
|
||||||
|
else runtime.settings.chunked_reasoner_worker_timeout_seconds
|
||||||
|
)
|
||||||
|
budget = (
|
||||||
|
notes_char_budget if notes_char_budget is not None else runtime.settings.chunked_reasoner_notes_char_budget
|
||||||
|
)
|
||||||
|
if chars <= 0:
|
||||||
|
raise ValueError("chars_per_slice must be positive")
|
||||||
|
if conc <= 0:
|
||||||
|
raise ValueError("concurrency must be positive")
|
||||||
|
if timeout <= 0:
|
||||||
|
raise ValueError("worker_timeout_seconds must be positive")
|
||||||
|
if budget <= 0:
|
||||||
|
raise ValueError("notes_char_budget must be positive")
|
||||||
|
self._runtime = runtime
|
||||||
|
self._chars_per_slice = chars
|
||||||
|
self._worker_timeout_seconds = timeout
|
||||||
|
self._notes_char_budget = budget
|
||||||
|
self._semaphore = asyncio.Semaphore(conc)
|
||||||
|
self._extractor: Agent[None, _ExtractedNotes] = Agent(
|
||||||
|
model=runtime.fast_model,
|
||||||
|
output_type=NativeOutput(_ExtractedNotes),
|
||||||
|
system_prompt=_EXTRACTOR_SYSTEM_PROMPT,
|
||||||
|
model_settings=runtime.fast_model_settings,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def gather_notes(self, pages: list[Page], question: str) -> list[ChunkNotes]:
|
||||||
|
"""Return notes covering every page that fit the synthesis budget.
|
||||||
|
|
||||||
|
Worker failures are tolerated: surviving notes are returned. Returns
|
||||||
|
an empty list only when every first-round chunk raises, which the
|
||||||
|
caller can treat as a hard failure.
|
||||||
|
|
||||||
|
Progress events fire as each first-round chunk finishes (in completion
|
||||||
|
order, not chunk order) carrying a monotonic ``completed`` counter so
|
||||||
|
consumers can render "Read X of Y" with X advancing by exactly one
|
||||||
|
per event. Subsequent compression rounds emit a single round-start
|
||||||
|
event each.
|
||||||
|
"""
|
||||||
|
if not pages:
|
||||||
|
raise ValueError("ChunkedReasoner.gather_notes requires at least one page")
|
||||||
|
|
||||||
|
chunks = [self._chunk_from_pages(slice_pages) for slice_pages in self._slice_pages(pages)]
|
||||||
|
slice_total = len(chunks)
|
||||||
|
logger.info(
|
||||||
|
"[chunked-reasoner] question=%r pages=%d slices=%d",
|
||||||
|
question,
|
||||||
|
len(pages),
|
||||||
|
slice_total,
|
||||||
|
)
|
||||||
|
await emit_progress(WholeDocReadStarted(question=question, pages=len(pages), slices=slice_total))
|
||||||
|
|
||||||
|
gather_start = time.perf_counter()
|
||||||
|
notes = await self._run_chunks(chunks, question)
|
||||||
|
|
||||||
|
await emit_progress(
|
||||||
|
WholeDocReadDone(
|
||||||
|
completed=len(notes),
|
||||||
|
slices=slice_total,
|
||||||
|
duration_seconds=round(time.perf_counter() - gather_start, 2),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return notes
|
||||||
|
|
||||||
|
async def _run_chunks(self, chunks: list[_Chunk], question: str) -> list[ChunkNotes]:
|
||||||
|
"""Run chunks through the extractor, regrouping and looping until under budget.
|
||||||
|
|
||||||
|
The first round emits per-chunk progress events for streaming UIs;
|
||||||
|
later rounds emit a single round-start event. Each round may produce
|
||||||
|
fewer notes than chunks (every chunk maps to at most one consolidated
|
||||||
|
note); when the rendered notes still exceed the budget, the survivors
|
||||||
|
are regrouped into fresh chunks and the loop runs again.
|
||||||
|
"""
|
||||||
|
round_number = 0
|
||||||
|
while True:
|
||||||
|
chunks_in = len(chunks)
|
||||||
|
result = await self._extract_chunks(chunks, question, round_number)
|
||||||
|
|
||||||
|
if result.slowest is not None:
|
||||||
|
slow_label, slow_duration = result.slowest
|
||||||
|
logger.info(
|
||||||
|
"[chunked-reasoner] round %d: %d/%d chunks succeeded; slowest %s (%.1fs)",
|
||||||
|
round_number,
|
||||||
|
result.successes,
|
||||||
|
chunks_in,
|
||||||
|
slow_label,
|
||||||
|
slow_duration,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.info(
|
||||||
|
"[chunked-reasoner] round %d: 0/%d chunks succeeded",
|
||||||
|
round_number,
|
||||||
|
chunks_in,
|
||||||
|
)
|
||||||
|
|
||||||
|
rendered_size = self._rendered_notes_size(result.notes)
|
||||||
|
if rendered_size <= self._notes_char_budget or len(result.notes) <= 1:
|
||||||
|
if round_number > 0:
|
||||||
|
logger.info(
|
||||||
|
"[chunked-reasoner] compression done after %d round(s): %d notes, %d chars",
|
||||||
|
round_number,
|
||||||
|
len(result.notes),
|
||||||
|
rendered_size,
|
||||||
|
)
|
||||||
|
return result.notes
|
||||||
|
|
||||||
|
if result.successes == 0:
|
||||||
|
# No forward progress this round; further rounds would
|
||||||
|
# reproduce the same shape. Return what we have.
|
||||||
|
logger.warning(
|
||||||
|
"[chunked-reasoner] round %d produced no successful extractions; bailing with %d notes",
|
||||||
|
round_number,
|
||||||
|
len(result.notes),
|
||||||
|
)
|
||||||
|
return result.notes
|
||||||
|
|
||||||
|
round_number += 1
|
||||||
|
groups = self._group_notes_for_compression(result.notes)
|
||||||
|
chunks = [self._chunk_from_notes(group) for group in groups]
|
||||||
|
logger.info(
|
||||||
|
"[chunked-reasoner] compression round %d: %d notes (%d chars) -> %d groups",
|
||||||
|
round_number,
|
||||||
|
len(result.notes),
|
||||||
|
rendered_size,
|
||||||
|
len(groups),
|
||||||
|
)
|
||||||
|
await emit_progress(
|
||||||
|
WholeDocCompressionRound(
|
||||||
|
round_number=round_number,
|
||||||
|
notes_in=len(result.notes),
|
||||||
|
groups=len(groups),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _extract_chunks(
|
||||||
|
self,
|
||||||
|
chunks: list[_Chunk],
|
||||||
|
question: str,
|
||||||
|
round_number: int,
|
||||||
|
) -> _RoundResult:
|
||||||
|
"""Run all chunks through the extractor in parallel; collect surviving notes.
|
||||||
|
|
||||||
|
Failures fall back to ``chunk.fallback`` (empty in the first round, so
|
||||||
|
failures drop; populated in compression rounds, so failures preserve
|
||||||
|
their input notes). The first round emits a
|
||||||
|
:class:`WholeDocSliceDone` per successful completion in completion
|
||||||
|
order, with a monotonic ``completed`` counter.
|
||||||
|
|
||||||
|
Returned notes are sorted by first page so downstream grouping packs
|
||||||
|
document-adjacent content together regardless of which task happened
|
||||||
|
to finish first.
|
||||||
|
"""
|
||||||
|
total = len(chunks)
|
||||||
|
pending: dict[asyncio.Task[tuple[ChunkNotes, float]], _Chunk] = {
|
||||||
|
asyncio.create_task(self._extract_chunk(chunk, question)): chunk for chunk in chunks
|
||||||
|
}
|
||||||
|
|
||||||
|
notes: list[ChunkNotes] = []
|
||||||
|
successes = 0
|
||||||
|
slowest: tuple[str, float] | None = None
|
||||||
|
completed = 0
|
||||||
|
|
||||||
|
try:
|
||||||
|
while pending:
|
||||||
|
done, _ = await asyncio.wait(pending.keys(), return_when=asyncio.FIRST_COMPLETED)
|
||||||
|
for task in done:
|
||||||
|
chunk = pending.pop(task)
|
||||||
|
exc = task.exception()
|
||||||
|
if exc is not None:
|
||||||
|
if chunk.fallback:
|
||||||
|
logger.warning(
|
||||||
|
"[chunked-reasoner] chunk %s failed: %s; preserving %d input note(s)",
|
||||||
|
chunk.label,
|
||||||
|
exc,
|
||||||
|
len(chunk.fallback),
|
||||||
|
)
|
||||||
|
notes.extend(chunk.fallback)
|
||||||
|
else:
|
||||||
|
logger.warning("[chunked-reasoner] chunk %s failed: %s", chunk.label, exc)
|
||||||
|
continue
|
||||||
|
extracted, duration = task.result()
|
||||||
|
notes.append(extracted)
|
||||||
|
successes += 1
|
||||||
|
completed += 1
|
||||||
|
if slowest is None or duration > slowest[1]:
|
||||||
|
slowest = (chunk.label, duration)
|
||||||
|
if round_number == 0:
|
||||||
|
await emit_progress(
|
||||||
|
WholeDocSliceDone(
|
||||||
|
completed=completed,
|
||||||
|
total=total,
|
||||||
|
pages=chunk.label,
|
||||||
|
duration_ms=int(duration * 1000),
|
||||||
|
excerpts=len(extracted.relevant_excerpts),
|
||||||
|
facts=len(extracted.facts),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
# On cancellation (typically a frontend disconnect propagating up
|
||||||
|
# through the streaming orchestrator) the per-chunk model calls
|
||||||
|
# would otherwise keep running to completion, billing tokens whose
|
||||||
|
# results nobody is reading. Cancel and drain so the upstream
|
||||||
|
# cancellation is the cancellation that matters.
|
||||||
|
if pending:
|
||||||
|
for task in pending:
|
||||||
|
task.cancel()
|
||||||
|
await asyncio.gather(*pending.keys(), return_exceptions=True)
|
||||||
|
|
||||||
|
notes.sort(key=lambda n: n.pages[0] if n.pages else 0)
|
||||||
|
return _RoundResult(notes=notes, successes=successes, slowest=slowest)
|
||||||
|
|
||||||
|
async def _extract_chunk(self, chunk: _Chunk, question: str) -> tuple[ChunkNotes, float]:
|
||||||
|
"""Run the extractor on one chunk and attach the chunk's pages to the output."""
|
||||||
|
try:
|
||||||
|
extracted, duration = await self._run_extractor(chunk.content, question, chunk.label)
|
||||||
|
except TimeoutError:
|
||||||
|
logger.warning(
|
||||||
|
"[chunked-reasoner] chunk %s timed out (limit %.1fs)",
|
||||||
|
chunk.label,
|
||||||
|
self._worker_timeout_seconds,
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
logger.debug(
|
||||||
|
"[chunked-reasoner] chunk %s: %d excerpt(s), %d fact(s) in %dms",
|
||||||
|
chunk.label,
|
||||||
|
len(extracted.relevant_excerpts),
|
||||||
|
len(extracted.facts),
|
||||||
|
int(duration * 1000),
|
||||||
|
)
|
||||||
|
return self._build_chunk_notes(extracted, chunk.pages), duration
|
||||||
|
|
||||||
|
async def _run_extractor(
|
||||||
|
self,
|
||||||
|
content: str,
|
||||||
|
question: str,
|
||||||
|
page_label: str,
|
||||||
|
) -> tuple[_ExtractedNotes, float]:
|
||||||
|
"""Inner primitive: run the extractor agent under semaphore + timeout."""
|
||||||
|
prompt = self._build_extraction_prompt(content, question)
|
||||||
|
async with self._semaphore:
|
||||||
|
start = time.perf_counter()
|
||||||
|
try:
|
||||||
|
result = await asyncio.wait_for(self._extractor.run(prompt), timeout=self._worker_timeout_seconds)
|
||||||
|
except TimeoutError:
|
||||||
|
duration = time.perf_counter() - start
|
||||||
|
logger.debug(
|
||||||
|
"[chunked-reasoner] extractor %s timed out after %dms",
|
||||||
|
page_label,
|
||||||
|
int(duration * 1000),
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
duration = time.perf_counter() - start
|
||||||
|
return result.output, duration
|
||||||
|
|
||||||
|
def _chunk_from_pages(self, pages: list[Page]) -> _Chunk:
|
||||||
|
"""Build a first-round chunk from a slice of raw pages."""
|
||||||
|
return _Chunk(
|
||||||
|
content="\n\n".join(f"[Page {p.page_number}]\n{p.text}" for p in pages),
|
||||||
|
pages=[p.page_number for p in pages],
|
||||||
|
fallback=[],
|
||||||
|
label=_page_range_label(pages),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _chunk_from_notes(self, group: list[ChunkNotes]) -> _Chunk:
|
||||||
|
"""Build a compression-round chunk from a group of prior-pass notes.
|
||||||
|
|
||||||
|
``fallback`` is the input group itself: if the extractor call fails,
|
||||||
|
the originals stay in the working set so page coverage isn't lost.
|
||||||
|
"""
|
||||||
|
return _Chunk(
|
||||||
|
content=self.format_notes(group),
|
||||||
|
pages=sorted({p for note in group for p in note.pages}),
|
||||||
|
fallback=group,
|
||||||
|
label=_note_range_label(group),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _group_notes_for_compression(self, notes: list[ChunkNotes]) -> list[list[ChunkNotes]]:
|
||||||
|
"""Pack consecutive notes into groups whose rendered size fits ``chars_per_slice``.
|
||||||
|
|
||||||
|
Each group becomes one compression-round chunk. Sized to match the
|
||||||
|
first-round slice budget so the extractor sees roughly the same input
|
||||||
|
footprint regardless of which round is running. Single notes that
|
||||||
|
exceed the budget on their own become their own group.
|
||||||
|
"""
|
||||||
|
groups: list[list[ChunkNotes]] = []
|
||||||
|
current: list[ChunkNotes] = []
|
||||||
|
current_chars = 0
|
||||||
|
for note in notes:
|
||||||
|
note_chars = self._rendered_notes_size([note])
|
||||||
|
if current and current_chars + note_chars > self._chars_per_slice:
|
||||||
|
groups.append(current)
|
||||||
|
current = []
|
||||||
|
current_chars = 0
|
||||||
|
current.append(note)
|
||||||
|
current_chars += note_chars
|
||||||
|
if current:
|
||||||
|
groups.append(current)
|
||||||
|
return groups
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _build_chunk_notes(extracted: _ExtractedNotes, pages: list[int]) -> ChunkNotes:
|
||||||
|
"""Build a public ChunkNotes from the model's output and the wrapper's pages."""
|
||||||
|
return ChunkNotes(
|
||||||
|
pages=pages,
|
||||||
|
summary=extracted.summary,
|
||||||
|
relevant_excerpts=extracted.relevant_excerpts,
|
||||||
|
facts=extracted.facts,
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _build_extraction_prompt(content: str, question: str) -> str:
|
||||||
|
"""Single prompt shape used for every round.
|
||||||
|
|
||||||
|
The system prompt explains the role; the user prompt just hands over
|
||||||
|
the question and the content. Whether ``content`` is raw page text
|
||||||
|
with ``[Page N]`` markers or formatted notes with ``[Notes from
|
||||||
|
pages A-B]`` markers, the same instructions apply.
|
||||||
|
"""
|
||||||
|
return f"User question:\n{question}\n\nContent:\n{content}"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _rendered_notes_size(notes: list[ChunkNotes]) -> int:
|
||||||
|
"""Length in characters of what :meth:`format_notes` would produce."""
|
||||||
|
return len(ChunkedReasoner.format_notes(notes))
|
||||||
|
|
||||||
|
async def reason[T: BaseModel](
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
pages: list[Page],
|
||||||
|
question: str,
|
||||||
|
answer_prompt: str,
|
||||||
|
answer_type: type[T],
|
||||||
|
) -> T:
|
||||||
|
"""Map over pages, then synthesise a structured answer of type ``T``.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pages: Document pages in order.
|
||||||
|
question: The user's question, passed to both workers and the
|
||||||
|
synthesiser. Workers use it to decide what's relevant.
|
||||||
|
answer_prompt: System prompt for the synthesis stage. Should
|
||||||
|
instruct the model to answer ``question`` from the notes
|
||||||
|
supplied. Owned by the caller because the answer's tone,
|
||||||
|
format, and grounding rules are domain-specific.
|
||||||
|
answer_type: Pydantic model describing the structured answer.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
An instance of ``answer_type`` produced by the synthesis stage.
|
||||||
|
"""
|
||||||
|
notes = await self.gather_notes(pages, question)
|
||||||
|
if not notes:
|
||||||
|
raise RuntimeError("All chunked-reasoning workers failed; no notes to synthesise from")
|
||||||
|
return await self._synthesise(question, notes, answer_prompt, answer_type)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def format_notes(notes: list[ChunkNotes]) -> str:
|
||||||
|
"""Render notes as readable text for inclusion in another agent's tool result.
|
||||||
|
|
||||||
|
Order is preserved. Page numbers, summary, excerpts and facts are all
|
||||||
|
emitted; empty sections are omitted.
|
||||||
|
"""
|
||||||
|
sections: list[str] = []
|
||||||
|
for n in notes:
|
||||||
|
page_label = (
|
||||||
|
f"pages {n.pages[0]}-{n.pages[-1]}"
|
||||||
|
if len(n.pages) > 1
|
||||||
|
else f"page {n.pages[0]}"
|
||||||
|
if n.pages
|
||||||
|
else "unknown pages"
|
||||||
|
)
|
||||||
|
block = [f"[Notes from {page_label}]", f"Summary: {n.summary}"]
|
||||||
|
if n.relevant_excerpts:
|
||||||
|
block.append("Relevant excerpts:")
|
||||||
|
block.extend(f"- {e}" for e in n.relevant_excerpts)
|
||||||
|
if n.facts:
|
||||||
|
block.append("Facts:")
|
||||||
|
block.extend(f"- {f}" for f in n.facts)
|
||||||
|
sections.append("\n".join(block))
|
||||||
|
return "\n\n".join(sections)
|
||||||
|
|
||||||
|
def _slice_pages(self, pages: list[Page]) -> list[list[Page]]:
|
||||||
|
"""Group consecutive pages into character-budgeted slices.
|
||||||
|
|
||||||
|
Page boundaries are preserved: a single page is never split across
|
||||||
|
slices. If one page exceeds the budget on its own, it becomes its
|
||||||
|
own slice.
|
||||||
|
"""
|
||||||
|
slices: list[list[Page]] = []
|
||||||
|
current: list[Page] = []
|
||||||
|
current_chars = 0
|
||||||
|
for page in pages:
|
||||||
|
if current and current_chars + page.char_count > self._chars_per_slice:
|
||||||
|
slices.append(current)
|
||||||
|
current = []
|
||||||
|
current_chars = 0
|
||||||
|
current.append(page)
|
||||||
|
current_chars += page.char_count
|
||||||
|
if current:
|
||||||
|
slices.append(current)
|
||||||
|
return slices
|
||||||
|
|
||||||
|
async def _synthesise[T: BaseModel](
|
||||||
|
self,
|
||||||
|
question: str,
|
||||||
|
notes: list[ChunkNotes],
|
||||||
|
answer_prompt: str,
|
||||||
|
answer_type: type[T],
|
||||||
|
) -> T:
|
||||||
|
agent: Agent[None, T] = Agent(
|
||||||
|
model=self._runtime.smart_model,
|
||||||
|
output_type=NativeOutput(answer_type),
|
||||||
|
system_prompt=answer_prompt,
|
||||||
|
model_settings=self._runtime.smart_model_settings,
|
||||||
|
)
|
||||||
|
prompt = f"User question:\n{question}\n\nNotes from across the document:\n\n{self.format_notes(notes)}"
|
||||||
|
result = await agent.run(prompt)
|
||||||
|
return result.output
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
"""Tool capability that lets an agent read whole documents end-to-end.
|
||||||
|
|
||||||
|
Companion to :class:`stirling.documents.RagCapability`. Where ``RagCapability``
|
||||||
|
gives an agent targeted vector retrieval, this gives it map-style whole-document
|
||||||
|
reading: every page is read in parallel by fast-model workers, and the
|
||||||
|
question-relevant notes are returned for the agent to synthesise.
|
||||||
|
|
||||||
|
Use both capabilities together when the agent should pick its strategy:
|
||||||
|
``search_knowledge`` for specific lookups, ``read_full_document`` for
|
||||||
|
aggregations, comparisons, and summaries.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
||||||
|
from pydantic_ai.toolsets import AbstractToolset
|
||||||
|
|
||||||
|
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner
|
||||||
|
from stirling.contracts import AiFile
|
||||||
|
from stirling.services import AppRuntime
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
# Cap on per-run calls. One pass already reads every page of every attached
|
||||||
|
# document, so a second call is almost always the smart model second-guessing
|
||||||
|
# itself on a near-identical query (and doubles wall-clock time for a sizeable
|
||||||
|
# document). If a follow-up genuinely needs more, ``search_knowledge`` is the
|
||||||
|
# right escape hatch. Configurable per-construction in case a future caller
|
||||||
|
# can prove a real two-read use case; the default stays at 1.
|
||||||
|
DEFAULT_MAX_READS = 1
|
||||||
|
|
||||||
|
|
||||||
|
class WholeDocReaderCapability:
|
||||||
|
"""Bundles instructions and the ``read_full_document`` toolset for agent injection.
|
||||||
|
|
||||||
|
Lifecycle: a ``WholeDocReaderCapability`` instance is intended to live for
|
||||||
|
the duration of a single agent run.
|
||||||
|
|
||||||
|
The agent picks between this and :class:`RagCapability` per the tool
|
||||||
|
descriptions: targeted retrieval vs whole-document reading.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
runtime: AppRuntime,
|
||||||
|
files: list[AiFile],
|
||||||
|
*,
|
||||||
|
reasoner: ChunkedReasoner | None = None,
|
||||||
|
max_reads: int = DEFAULT_MAX_READS,
|
||||||
|
) -> None:
|
||||||
|
self._runtime = runtime
|
||||||
|
self._files = files
|
||||||
|
self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime)
|
||||||
|
self._max_reads = max_reads
|
||||||
|
self._read_count = 0
|
||||||
|
toolset: FunctionToolset[None] = FunctionToolset()
|
||||||
|
toolset.add_function(
|
||||||
|
self._read_full_document,
|
||||||
|
name="read_full_document",
|
||||||
|
prepare=self._prepare_read_full_document,
|
||||||
|
)
|
||||||
|
self._toolset = toolset
|
||||||
|
|
||||||
|
@property
|
||||||
|
def instructions(self) -> str:
|
||||||
|
names = ", ".join(f.name for f in self._files) if self._files else "the attached documents"
|
||||||
|
return (
|
||||||
|
"You have a 'read_full_document' tool that reads every page of "
|
||||||
|
f"{names} in parallel and returns notes relevant to a query. "
|
||||||
|
"Use it when answering requires seeing the whole document end-to-end "
|
||||||
|
"(summaries, aggregations, comparisons across sections). One call "
|
||||||
|
"already reads everything; phrase the query to cover all the angles "
|
||||||
|
"you need in a single pass. For follow-ups or specific lookups use "
|
||||||
|
"'search_knowledge', which is cheaper and targeted."
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def toolset(self) -> AbstractToolset[None]:
|
||||||
|
return self._toolset
|
||||||
|
|
||||||
|
async def _prepare_read_full_document(
|
||||||
|
self,
|
||||||
|
ctx: RunContext[None],
|
||||||
|
tool_def: ToolDefinition,
|
||||||
|
) -> ToolDefinition | None:
|
||||||
|
"""Hide the tool from the agent's toolset once the per-run budget is spent.
|
||||||
|
Mirrors the search_knowledge prepare callback."""
|
||||||
|
if self._read_count >= self._max_reads:
|
||||||
|
return None
|
||||||
|
return tool_def
|
||||||
|
|
||||||
|
async def _read_full_document(self, query: str) -> str:
|
||||||
|
"""Read every page of the attached documents and return notes relevant to the query.
|
||||||
|
|
||||||
|
Use this when answering needs the whole document end-to-end - summaries,
|
||||||
|
aggregations like 'largest number' or 'shortest chapter', or comparisons
|
||||||
|
across sections. Slow and expensive (one fast-model call per slice per
|
||||||
|
document); prefer search_knowledge for targeted lookups.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: A focused description of what to extract from the documents,
|
||||||
|
phrased so a worker reading just one slice can decide what's
|
||||||
|
relevant to the user's question.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Per-document sections of structured notes (page numbers, summary,
|
||||||
|
relevant excerpts, extracted facts), already ordered by page.
|
||||||
|
"""
|
||||||
|
self._read_count += 1
|
||||||
|
if not self._files:
|
||||||
|
return "No documents attached to read."
|
||||||
|
|
||||||
|
sections: list[str] = []
|
||||||
|
for file in self._files:
|
||||||
|
pages = await self._runtime.documents.read_pages(file.id)
|
||||||
|
if not pages:
|
||||||
|
logger.info(
|
||||||
|
"[whole-doc-reader] no stored pages for %s (id=%s); skipping",
|
||||||
|
file.name,
|
||||||
|
file.id,
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
notes = await self._reasoner.gather_notes(pages, query)
|
||||||
|
if not notes:
|
||||||
|
continue
|
||||||
|
sections.append(f"=== {file.name} ===\n{ChunkedReasoner.format_notes(notes)}")
|
||||||
|
|
||||||
|
if not sections:
|
||||||
|
return "Could not read any document content."
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
"[whole-doc-reader] read query=%r files=%d -> %d chars",
|
||||||
|
query,
|
||||||
|
len(self._files),
|
||||||
|
sum(len(s) for s in sections),
|
||||||
|
)
|
||||||
|
return "\n\n".join(sections)
|
||||||
@@ -19,13 +19,13 @@ from stirling.agents.pdf_comment import PdfCommentAgent
|
|||||||
from stirling.api.middleware import UserIdMiddleware
|
from stirling.api.middleware import UserIdMiddleware
|
||||||
from stirling.api.routes import (
|
from stirling.api.routes import (
|
||||||
agent_draft_router,
|
agent_draft_router,
|
||||||
|
document_router,
|
||||||
execution_router,
|
execution_router,
|
||||||
ledger_router,
|
ledger_router,
|
||||||
orchestrator_router,
|
orchestrator_router,
|
||||||
pdf_comments_router,
|
pdf_comments_router,
|
||||||
pdf_edit_router,
|
pdf_edit_router,
|
||||||
pdf_question_router,
|
pdf_question_router,
|
||||||
rag_router,
|
|
||||||
)
|
)
|
||||||
from stirling.config import AppSettings, load_settings
|
from stirling.config import AppSettings, load_settings
|
||||||
from stirling.contracts import HealthResponse
|
from stirling.contracts import HealthResponse
|
||||||
@@ -57,7 +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()
|
await runtime.documents.close()
|
||||||
if tracer_provider:
|
if tracer_provider:
|
||||||
tracer_provider.shutdown()
|
tracer_provider.shutdown()
|
||||||
|
|
||||||
@@ -69,7 +69,7 @@ app.include_router(pdf_edit_router)
|
|||||||
app.include_router(pdf_question_router)
|
app.include_router(pdf_question_router)
|
||||||
app.include_router(agent_draft_router)
|
app.include_router(agent_draft_router)
|
||||||
app.include_router(execution_router)
|
app.include_router(execution_router)
|
||||||
app.include_router(rag_router)
|
app.include_router(document_router)
|
||||||
app.include_router(ledger_router)
|
app.include_router(ledger_router)
|
||||||
app.include_router(pdf_comments_router)
|
app.include_router(pdf_comments_router)
|
||||||
|
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ from stirling.agents import (
|
|||||||
)
|
)
|
||||||
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.documents import DocumentService
|
||||||
from stirling.services import AppRuntime
|
from stirling.services import AppRuntime
|
||||||
|
|
||||||
|
|
||||||
@@ -39,8 +39,8 @@ def get_execution_planning_agent(request: Request) -> ExecutionPlanningAgent:
|
|||||||
return request.app.state.execution_planning_agent
|
return request.app.state.execution_planning_agent
|
||||||
|
|
||||||
|
|
||||||
def get_rag_service(request: Request) -> RagService:
|
def get_document_service(request: Request) -> DocumentService:
|
||||||
return request.app.state.runtime.rag_service
|
return request.app.state.runtime.documents
|
||||||
|
|
||||||
|
|
||||||
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
|
def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
|
||||||
|
|||||||
@@ -1,19 +1,19 @@
|
|||||||
from .agent_drafts import router as agent_draft_router
|
from .agent_drafts import router as agent_draft_router
|
||||||
|
from .documents import router as document_router
|
||||||
from .execution import router as execution_router
|
from .execution import router as execution_router
|
||||||
from .ledger import router as ledger_router
|
from .ledger import router as ledger_router
|
||||||
from .orchestrator import router as orchestrator_router
|
from .orchestrator import router as orchestrator_router
|
||||||
from .pdf_comments import router as pdf_comments_router
|
from .pdf_comments import router as pdf_comments_router
|
||||||
from .pdf_edit import router as pdf_edit_router
|
from .pdf_edit import router as pdf_edit_router
|
||||||
from .pdf_questions import router as pdf_question_router
|
from .pdf_questions import router as pdf_question_router
|
||||||
from .rag import router as rag_router
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"agent_draft_router",
|
"agent_draft_router",
|
||||||
|
"document_router",
|
||||||
"execution_router",
|
"execution_router",
|
||||||
"ledger_router",
|
"ledger_router",
|
||||||
"orchestrator_router",
|
"orchestrator_router",
|
||||||
"pdf_comments_router",
|
"pdf_comments_router",
|
||||||
"pdf_edit_router",
|
"pdf_edit_router",
|
||||||
"pdf_question_router",
|
"pdf_question_router",
|
||||||
"rag_router",
|
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -0,0 +1,49 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Annotated
|
||||||
|
|
||||||
|
from fastapi import APIRouter, Depends
|
||||||
|
|
||||||
|
from stirling.api.dependencies import get_document_service
|
||||||
|
from stirling.contracts import (
|
||||||
|
DeleteDocumentResponse,
|
||||||
|
IngestDocumentRequest,
|
||||||
|
IngestDocumentResponse,
|
||||||
|
)
|
||||||
|
from stirling.documents import DocumentService
|
||||||
|
from stirling.models import FileId
|
||||||
|
|
||||||
|
router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("", response_model=IngestDocumentResponse)
|
||||||
|
async def ingest_document(
|
||||||
|
request: IngestDocumentRequest,
|
||||||
|
documents: Annotated[DocumentService, Depends(get_document_service)],
|
||||||
|
) -> IngestDocumentResponse:
|
||||||
|
"""Replace-ingest a document's content under ``document_id``.
|
||||||
|
|
||||||
|
Stores both representations in one shot:
|
||||||
|
* embedded chunks for RAG search,
|
||||||
|
* ordered page text for whole-document reading.
|
||||||
|
Any previously-stored content for this document is removed first.
|
||||||
|
"""
|
||||||
|
pages = request.page_text or []
|
||||||
|
chunks_indexed = await documents.ingest(
|
||||||
|
collection=request.document_id,
|
||||||
|
pages=pages,
|
||||||
|
source=request.source,
|
||||||
|
)
|
||||||
|
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed)
|
||||||
|
|
||||||
|
|
||||||
|
@router.delete("/{document_id}", response_model=DeleteDocumentResponse)
|
||||||
|
async def delete_document(
|
||||||
|
document_id: FileId,
|
||||||
|
documents: Annotated[DocumentService, Depends(get_document_service)],
|
||||||
|
) -> DeleteDocumentResponse:
|
||||||
|
"""Remove a document's content. Idempotent."""
|
||||||
|
existed = await documents.has_collection(document_id)
|
||||||
|
if existed:
|
||||||
|
await documents.delete_collection(document_id)
|
||||||
|
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
|
||||||
@@ -1,19 +1,170 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from typing import Annotated
|
import asyncio
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Annotated, assert_never
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends
|
from fastapi import APIRouter, Depends
|
||||||
|
from fastapi.responses import StreamingResponse
|
||||||
|
|
||||||
from stirling.agents import OrchestratorAgent
|
from stirling.agents import OrchestratorAgent
|
||||||
from stirling.api.dependencies import get_orchestrator_agent
|
from stirling.api.dependencies import get_orchestrator_agent
|
||||||
from stirling.contracts import OrchestratorRequest, OrchestratorResponse
|
from stirling.contracts import OrchestratorRequest, OrchestratorResponse, ProgressEvent
|
||||||
|
from stirling.services import reset_progress_emitter, set_progress_emitter
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Cadence for keep-alive heartbeats on the streaming endpoint. Java forwards
|
||||||
|
# them to the frontend as SSE comments; their job is to make every layer of
|
||||||
|
# the connection visibly alive at this rhythm so disconnects surface within a
|
||||||
|
# bounded window instead of waiting for the next progress event.
|
||||||
|
HEARTBEAT_INTERVAL_SECONDS = 10.0
|
||||||
|
|
||||||
router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
|
router = APIRouter(prefix="/api/v1/orchestrator", tags=["orchestrator"])
|
||||||
|
|
||||||
|
|
||||||
@router.post("", response_model=OrchestratorResponse)
|
@router.post("")
|
||||||
async def orchestrate(
|
async def orchestrate(
|
||||||
request: OrchestratorRequest,
|
request: OrchestratorRequest,
|
||||||
agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
|
agent: Annotated[OrchestratorAgent, Depends(get_orchestrator_agent)],
|
||||||
) -> OrchestratorResponse:
|
) -> StreamingResponse:
|
||||||
return await agent.handle(request)
|
"""Run the orchestrator and stream NDJSON events.
|
||||||
|
|
||||||
|
Each output line is a JSON object with an ``event`` field. ``progress``
|
||||||
|
events arrive whenever an inner agent reports work (e.g. each
|
||||||
|
chunked-reasoner slice completing); the final ``result`` event carries the
|
||||||
|
typed orchestrator response. ``error`` events surface failures without
|
||||||
|
breaking the connection. ``heartbeat`` events fire on a fixed cadence to
|
||||||
|
keep idle connections visibly alive so disconnects propagate.
|
||||||
|
|
||||||
|
The stream itself is the liveness signal: as long as events flow, work is
|
||||||
|
alive. Java consumes this with a long total timeout and treats line
|
||||||
|
arrival as forward progress.
|
||||||
|
"""
|
||||||
|
return StreamingResponse(
|
||||||
|
_OrchestratorStream(
|
||||||
|
agent=agent,
|
||||||
|
request=request,
|
||||||
|
heartbeat_interval_seconds=HEARTBEAT_INTERVAL_SECONDS,
|
||||||
|
).iterate(),
|
||||||
|
media_type="application/x-ndjson",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class _ProgressFrame:
|
||||||
|
event: ProgressEvent
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class _ResultFrame:
|
||||||
|
response: OrchestratorResponse
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class _ErrorFrame:
|
||||||
|
message: str
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True, slots=True)
|
||||||
|
class _HeartbeatFrame:
|
||||||
|
"""No payload: a heartbeat exists only to push bytes through the pipe.
|
||||||
|
|
||||||
|
Without periodic traffic, a slow workflow phase (e.g. all extractor
|
||||||
|
workers busy on long calls) leaves the engine writer, Java's SSE
|
||||||
|
forwarder, and the frontend's fetch all silently waiting. A closed
|
||||||
|
connection at any layer wouldn't surface until the next real event,
|
||||||
|
which could be many tens of seconds away. Heartbeats bound that window
|
||||||
|
to :data:`HEARTBEAT_INTERVAL_SECONDS`.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
type _StreamFrame = _ProgressFrame | _ResultFrame | _ErrorFrame | _HeartbeatFrame
|
||||||
|
|
||||||
|
|
||||||
|
def _serialize_frame(frame: _StreamFrame) -> bytes:
|
||||||
|
"""Render a frame as one NDJSON line."""
|
||||||
|
match frame:
|
||||||
|
case _ProgressFrame(event=event):
|
||||||
|
body = {"event": "progress", **event.model_dump(mode="json")}
|
||||||
|
case _ResultFrame(response=response):
|
||||||
|
body = {"event": "result", "response": response.model_dump(mode="json")}
|
||||||
|
case _ErrorFrame(message=message):
|
||||||
|
body = {"event": "error", "message": message}
|
||||||
|
case _HeartbeatFrame():
|
||||||
|
body = {"event": "heartbeat"}
|
||||||
|
case _:
|
||||||
|
assert_never(frame)
|
||||||
|
return (json.dumps(body) + "\n").encode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
class _OrchestratorStream:
|
||||||
|
"""Drives one streaming orchestrator request.
|
||||||
|
|
||||||
|
Owns the per-request queue and pumps progress events through it; the agent
|
||||||
|
runs as a child task so its emissions and the streaming response interleave.
|
||||||
|
A heartbeat task pushes keep-alive messages onto the same queue at a fixed
|
||||||
|
cadence so the connection stays visibly alive between progress events.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
agent: OrchestratorAgent,
|
||||||
|
request: OrchestratorRequest,
|
||||||
|
heartbeat_interval_seconds: float,
|
||||||
|
) -> None:
|
||||||
|
self._agent = agent
|
||||||
|
self._request = request
|
||||||
|
self._heartbeat_interval_seconds = heartbeat_interval_seconds
|
||||||
|
self._queue: asyncio.Queue[_StreamFrame | None] = asyncio.Queue()
|
||||||
|
|
||||||
|
async def iterate(self) -> AsyncIterator[bytes]:
|
||||||
|
token = set_progress_emitter(self._emit_progress)
|
||||||
|
agent_task = asyncio.create_task(self._run_agent())
|
||||||
|
heartbeat_task = asyncio.create_task(self._emit_heartbeats())
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
frame = await self._queue.get()
|
||||||
|
if frame is None:
|
||||||
|
break
|
||||||
|
yield _serialize_frame(frame)
|
||||||
|
finally:
|
||||||
|
reset_progress_emitter(token)
|
||||||
|
await self._cancel_task(heartbeat_task)
|
||||||
|
await self._cancel_task(agent_task)
|
||||||
|
|
||||||
|
async def _emit_progress(self, event: ProgressEvent) -> None:
|
||||||
|
await self._queue.put(_ProgressFrame(event=event))
|
||||||
|
|
||||||
|
async def _emit_heartbeats(self) -> None:
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(self._heartbeat_interval_seconds)
|
||||||
|
await self._queue.put(_HeartbeatFrame())
|
||||||
|
|
||||||
|
async def _run_agent(self) -> None:
|
||||||
|
try:
|
||||||
|
response = await self._agent.handle(self._request)
|
||||||
|
await self._queue.put(_ResultFrame(response=response))
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception as exc:
|
||||||
|
logger.exception("orchestrator stream failed")
|
||||||
|
await self._queue.put(_ErrorFrame(message=str(exc)))
|
||||||
|
finally:
|
||||||
|
await self._queue.put(None)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _cancel_task(task: asyncio.Task[None]) -> None:
|
||||||
|
if task.done():
|
||||||
|
return
|
||||||
|
task.cancel()
|
||||||
|
try:
|
||||||
|
await task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
except Exception:
|
||||||
|
logger.exception("background task failed during cancellation", exc_info=True)
|
||||||
|
|||||||
@@ -1,63 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from typing import Annotated
|
|
||||||
|
|
||||||
from fastapi import APIRouter, Depends
|
|
||||||
|
|
||||||
from stirling.api.dependencies import get_rag_service
|
|
||||||
from stirling.contracts import (
|
|
||||||
DeleteDocumentResponse,
|
|
||||||
IngestDocumentRequest,
|
|
||||||
IngestDocumentResponse,
|
|
||||||
PdfContentType,
|
|
||||||
)
|
|
||||||
from stirling.models import FileId
|
|
||||||
from stirling.rag import Document, RagService
|
|
||||||
|
|
||||||
router = APIRouter(prefix="/api/v1/rag", tags=["rag"])
|
|
||||||
|
|
||||||
|
|
||||||
@router.post("/documents", response_model=IngestDocumentResponse)
|
|
||||||
async def ingest_document(
|
|
||||||
request: IngestDocumentRequest,
|
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
|
||||||
) -> IngestDocumentResponse:
|
|
||||||
"""Replace-ingest a document's content under ``document_id``.
|
|
||||||
|
|
||||||
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.delete("/documents/{document_id}", response_model=DeleteDocumentResponse)
|
|
||||||
async def delete_document(
|
|
||||||
document_id: FileId,
|
|
||||||
rag: Annotated[RagService, Depends(get_rag_service)],
|
|
||||||
) -> DeleteDocumentResponse:
|
|
||||||
"""Remove a document's content from RAG. Idempotent."""
|
|
||||||
existed = await rag.has_collection(document_id)
|
|
||||||
if existed:
|
|
||||||
await rag.delete_collection(document_id)
|
|
||||||
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
|
|
||||||
@@ -38,18 +38,38 @@ class AppSettings(BaseSettings):
|
|||||||
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")
|
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
|
||||||
|
|
||||||
|
# Chunked reasoner settings (whole-document map-reduce).
|
||||||
|
chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE")
|
||||||
|
chunked_reasoner_concurrency: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CONCURRENCY")
|
||||||
|
chunked_reasoner_worker_timeout_seconds: float = Field(
|
||||||
|
validation_alias="STIRLING_CHUNKED_REASONER_WORKER_TIMEOUT_SECONDS"
|
||||||
|
)
|
||||||
|
# Maximum size, in characters, of the rendered notes block before the
|
||||||
|
# reasoner folds slice notes hierarchically. The Anthropic context limit
|
||||||
|
# is 200k tokens (~880k chars); we leave a generous margin for the
|
||||||
|
# downstream agent's system prompt, history, tool definitions, and
|
||||||
|
# response budget.
|
||||||
|
chunked_reasoner_notes_char_budget: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_NOTES_CHAR_BUDGET")
|
||||||
|
|
||||||
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")
|
||||||
|
|
||||||
log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
|
log_level: str = Field(default="INFO", validation_alias="STIRLING_LOG_LEVEL")
|
||||||
log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
|
log_file: str = Field(default="", validation_alias="STIRLING_LOG_FILE")
|
||||||
|
# When true, raises httpx + httpcore logger levels so every outgoing
|
||||||
|
# model SDK call is logged with timing. Use to diagnose worker stalls:
|
||||||
|
# a hung request shows the "Request: POST ..." line with no matching
|
||||||
|
# response line, confirming the hang is transport-layer (not in our
|
||||||
|
# code or the Anthropic SDK itself). Off by default — DEBUG-level
|
||||||
|
# output is high-volume.
|
||||||
|
http_debug: bool = Field(default=False, validation_alias="STIRLING_HTTP_DEBUG")
|
||||||
|
|
||||||
posthog_enabled: bool = Field(validation_alias="STIRLING_POSTHOG_ENABLED")
|
posthog_enabled: bool = Field(validation_alias="STIRLING_POSTHOG_ENABLED")
|
||||||
posthog_api_key: str = Field(validation_alias="STIRLING_POSTHOG_API_KEY")
|
posthog_api_key: str = Field(validation_alias="STIRLING_POSTHOG_API_KEY")
|
||||||
posthog_host: str = Field(validation_alias="STIRLING_POSTHOG_HOST")
|
posthog_host: str = Field(validation_alias="STIRLING_POSTHOG_HOST")
|
||||||
|
|
||||||
|
|
||||||
def _configure_logging(level_name: str, log_file: str) -> None:
|
def _configure_logging(level_name: str, log_file: str, http_debug: bool) -> None:
|
||||||
"""Configure the ``stirling`` logger hierarchy."""
|
"""Configure the ``stirling`` logger hierarchy."""
|
||||||
level = logging.getLevelNamesMapping().get(level_name.upper())
|
level = logging.getLevelNamesMapping().get(level_name.upper())
|
||||||
if level is None:
|
if level is None:
|
||||||
@@ -83,11 +103,43 @@ def _configure_logging(level_name: str, log_file: str) -> None:
|
|||||||
fh.setLevel(level)
|
fh.setLevel(level)
|
||||||
root.addHandler(fh)
|
root.addHandler(fh)
|
||||||
|
|
||||||
|
if http_debug:
|
||||||
|
_enable_http_debug(formatter)
|
||||||
|
|
||||||
|
|
||||||
|
def _enable_http_debug(formatter: logging.Formatter) -> None:
|
||||||
|
"""Surface every httpx/httpcore call against the Anthropic API.
|
||||||
|
|
||||||
|
httpx emits one INFO line per request with the URL and final status,
|
||||||
|
which is the most useful signal for diagnosing hung worker calls: a
|
||||||
|
successful call shows "Request" then "Response" within a second or two;
|
||||||
|
a hung one shows "Request" with no matching response until it's
|
||||||
|
cancelled. httpcore at DEBUG drills down to TCP / HTTP/2 stream events
|
||||||
|
if the user wants to see exactly where bytes stop flowing.
|
||||||
|
|
||||||
|
The ``stirling`` console handler is scoped to its own logger tree, so
|
||||||
|
we attach a dedicated stream handler here. Without it, httpx records
|
||||||
|
propagate to the root logger which has no handler in our setup and the
|
||||||
|
output is silently dropped.
|
||||||
|
"""
|
||||||
|
handler = logging.StreamHandler()
|
||||||
|
handler.setFormatter(formatter)
|
||||||
|
handler.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
|
for name, level in (("httpx", logging.INFO), ("httpcore", logging.DEBUG)):
|
||||||
|
lg = logging.getLogger(name)
|
||||||
|
lg.setLevel(level)
|
||||||
|
# Idempotent: avoid stacking handlers on settings reload.
|
||||||
|
if not any(getattr(h, "_stirling_http_debug", False) for h in lg.handlers):
|
||||||
|
handler._stirling_http_debug = True # type: ignore[attr-defined]
|
||||||
|
lg.addHandler(handler)
|
||||||
|
lg.propagate = False
|
||||||
|
|
||||||
|
|
||||||
@lru_cache(maxsize=1)
|
@lru_cache(maxsize=1)
|
||||||
def load_settings() -> AppSettings:
|
def load_settings() -> AppSettings:
|
||||||
load_dotenv(ENV_FILE)
|
load_dotenv(ENV_FILE)
|
||||||
load_dotenv(ENV_LOCAL_FILE, override=True)
|
load_dotenv(ENV_LOCAL_FILE, override=True)
|
||||||
settings = AppSettings.model_validate({})
|
settings = AppSettings.model_validate({})
|
||||||
_configure_logging(settings.log_level, settings.log_file)
|
_configure_logging(settings.log_level, settings.log_file, settings.http_debug)
|
||||||
return settings
|
return settings
|
||||||
|
|||||||
@@ -28,6 +28,14 @@ from .common import (
|
|||||||
format_conversation_history,
|
format_conversation_history,
|
||||||
format_file_names,
|
format_file_names,
|
||||||
)
|
)
|
||||||
|
from .documents import (
|
||||||
|
DeleteDocumentResponse,
|
||||||
|
IngestDocumentRequest,
|
||||||
|
IngestDocumentResponse,
|
||||||
|
Page,
|
||||||
|
PageRange,
|
||||||
|
PageText,
|
||||||
|
)
|
||||||
from .execution import (
|
from .execution import (
|
||||||
AgentExecutionRequest,
|
AgentExecutionRequest,
|
||||||
CannotContinueExecutionAction,
|
CannotContinueExecutionAction,
|
||||||
@@ -79,11 +87,12 @@ from .pdf_questions import (
|
|||||||
PdfQuestionResponse,
|
PdfQuestionResponse,
|
||||||
PdfQuestionTerminalResponse,
|
PdfQuestionTerminalResponse,
|
||||||
)
|
)
|
||||||
from .rag import (
|
from .progress import (
|
||||||
DeleteDocumentResponse,
|
ProgressEvent,
|
||||||
IngestDocumentRequest,
|
WholeDocCompressionRound,
|
||||||
IngestDocumentResponse,
|
WholeDocReadDone,
|
||||||
IngestedPageText,
|
WholeDocReadStarted,
|
||||||
|
WholeDocSliceDone,
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
@@ -123,7 +132,6 @@ __all__ = [
|
|||||||
"HealthResponse",
|
"HealthResponse",
|
||||||
"IngestDocumentRequest",
|
"IngestDocumentRequest",
|
||||||
"IngestDocumentResponse",
|
"IngestDocumentResponse",
|
||||||
"IngestedPageText",
|
|
||||||
"MathAuditorToolReportArtifact",
|
"MathAuditorToolReportArtifact",
|
||||||
"NeedContentFileRequest",
|
"NeedContentFileRequest",
|
||||||
"NeedContentResponse",
|
"NeedContentResponse",
|
||||||
@@ -131,6 +139,9 @@ __all__ = [
|
|||||||
"NextExecutionAction",
|
"NextExecutionAction",
|
||||||
"OrchestratorRequest",
|
"OrchestratorRequest",
|
||||||
"OrchestratorResponse",
|
"OrchestratorResponse",
|
||||||
|
"Page",
|
||||||
|
"PageRange",
|
||||||
|
"PageText",
|
||||||
"PdfCommentInstruction",
|
"PdfCommentInstruction",
|
||||||
"PdfCommentReport",
|
"PdfCommentReport",
|
||||||
"PdfCommentRequest",
|
"PdfCommentRequest",
|
||||||
@@ -146,6 +157,7 @@ __all__ = [
|
|||||||
"PdfQuestionResponse",
|
"PdfQuestionResponse",
|
||||||
"PdfQuestionTerminalResponse",
|
"PdfQuestionTerminalResponse",
|
||||||
"PdfTextSelection",
|
"PdfTextSelection",
|
||||||
|
"ProgressEvent",
|
||||||
"Requisition",
|
"Requisition",
|
||||||
"Severity",
|
"Severity",
|
||||||
"StepKind",
|
"StepKind",
|
||||||
@@ -156,6 +168,10 @@ __all__ = [
|
|||||||
"ToolReportArtifact",
|
"ToolReportArtifact",
|
||||||
"UnsupportedCapabilityResponse",
|
"UnsupportedCapabilityResponse",
|
||||||
"Verdict",
|
"Verdict",
|
||||||
|
"WholeDocCompressionRound",
|
||||||
|
"WholeDocReadDone",
|
||||||
|
"WholeDocReadStarted",
|
||||||
|
"WholeDocSliceDone",
|
||||||
"WorkflowArtifact",
|
"WorkflowArtifact",
|
||||||
"WorkflowOutcome",
|
"WorkflowOutcome",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -167,10 +167,10 @@ ToolReportArtifact = MathAuditorToolReportArtifact
|
|||||||
|
|
||||||
|
|
||||||
class NeedIngestResponse(ApiModel):
|
class NeedIngestResponse(ApiModel):
|
||||||
"""Signal that the listed files must be ingested into RAG before the agent can continue.
|
"""Signal that the listed files must be ingested before the agent can continue.
|
||||||
|
|
||||||
Java's handling: for each file, extract the requested content types, POST to
|
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.
|
``/api/v1/documents`` keyed by ``file.id``, then retry the original request.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST
|
outcome: Literal[WorkflowOutcome.NEED_INGEST] = WorkflowOutcome.NEED_INGEST
|
||||||
|
|||||||
@@ -0,0 +1,61 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pydantic import Field
|
||||||
|
|
||||||
|
from stirling.models import ApiModel
|
||||||
|
|
||||||
|
from .common import FileId
|
||||||
|
|
||||||
|
|
||||||
|
class PageText(ApiModel):
|
||||||
|
"""A single page of extracted text on the ingest wire."""
|
||||||
|
|
||||||
|
page_number: int = Field(ge=1)
|
||||||
|
text: str
|
||||||
|
|
||||||
|
|
||||||
|
class Page(ApiModel):
|
||||||
|
"""A single page of a document, retrieved from storage.
|
||||||
|
|
||||||
|
``char_count`` is precomputed at ingest time and reported here so callers
|
||||||
|
can budget how much content they want to read without first concatenating
|
||||||
|
the text of every page.
|
||||||
|
"""
|
||||||
|
|
||||||
|
page_number: int = Field(ge=1)
|
||||||
|
text: str
|
||||||
|
char_count: int = Field(ge=0)
|
||||||
|
|
||||||
|
|
||||||
|
class PageRange(ApiModel):
|
||||||
|
"""Inclusive page range for partial reads. Both bounds are 1-indexed."""
|
||||||
|
|
||||||
|
start: int = Field(ge=1)
|
||||||
|
end: int = Field(ge=1)
|
||||||
|
|
||||||
|
|
||||||
|
class IngestDocumentRequest(ApiModel):
|
||||||
|
"""Replace-ingest a document's content under the given ``document_id``.
|
||||||
|
|
||||||
|
Each call wipes any previously-stored content for the document and writes
|
||||||
|
both the vector-chunk and ordered-page representations from the supplied
|
||||||
|
pages.
|
||||||
|
|
||||||
|
``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[PageText] | None = None
|
||||||
|
|
||||||
|
|
||||||
|
class IngestDocumentResponse(ApiModel):
|
||||||
|
document_id: FileId
|
||||||
|
chunks_indexed: int
|
||||||
|
|
||||||
|
|
||||||
|
class DeleteDocumentResponse(ApiModel):
|
||||||
|
document_id: FileId
|
||||||
|
deleted: bool
|
||||||
@@ -0,0 +1,70 @@
|
|||||||
|
"""Progress events emitted by deep callees during a streaming orchestrator run.
|
||||||
|
|
||||||
|
Each subclass models one engine-side phase. The Java side forwards the JSON
|
||||||
|
verbatim into ``AiWorkflowProgressEvent.engineDetail``; the frontend switches
|
||||||
|
on ``phase`` and renders the typed fields.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import Annotated, Literal
|
||||||
|
|
||||||
|
from pydantic import Field
|
||||||
|
|
||||||
|
from stirling.models import ApiModel
|
||||||
|
|
||||||
|
|
||||||
|
class WholeDocReadStarted(ApiModel):
|
||||||
|
phase: Literal["whole_doc_read_started"] = "whole_doc_read_started"
|
||||||
|
question: str
|
||||||
|
pages: int
|
||||||
|
slices: int
|
||||||
|
|
||||||
|
|
||||||
|
class WholeDocSliceDone(ApiModel):
|
||||||
|
"""Emitted as each chunked-reasoner worker completes.
|
||||||
|
|
||||||
|
``completed`` is a monotonically increasing counter (1..total) reflecting
|
||||||
|
the order in which workers finished, NOT the slice's position in the
|
||||||
|
document. Callers showing "Read X of Y" should use this directly so X
|
||||||
|
increments by one with each event.
|
||||||
|
"""
|
||||||
|
|
||||||
|
phase: Literal["whole_doc_slice_done"] = "whole_doc_slice_done"
|
||||||
|
completed: int
|
||||||
|
total: int
|
||||||
|
pages: str
|
||||||
|
duration_ms: int
|
||||||
|
excerpts: int
|
||||||
|
facts: int
|
||||||
|
|
||||||
|
|
||||||
|
class WholeDocCompressionRound(ApiModel):
|
||||||
|
"""Emitted when the gathered slice notes exceed the synthesis context
|
||||||
|
budget and the reasoner consolidates them with a fast-model fold pass.
|
||||||
|
|
||||||
|
Long documents (a 3000-page novel produces ~900k chars of raw notes)
|
||||||
|
would otherwise overflow the smart-model's prompt. ``notes_in`` is the
|
||||||
|
count entering the round; ``groups`` is the number of fold calls fired
|
||||||
|
(each producing one consolidated note). One or two rounds usually fit;
|
||||||
|
the event fires per round so callers can render "Consolidating notes
|
||||||
|
(round N)..." rather than going silent through the fold.
|
||||||
|
"""
|
||||||
|
|
||||||
|
phase: Literal["whole_doc_compression_round"] = "whole_doc_compression_round"
|
||||||
|
round_number: int
|
||||||
|
notes_in: int
|
||||||
|
groups: int
|
||||||
|
|
||||||
|
|
||||||
|
class WholeDocReadDone(ApiModel):
|
||||||
|
phase: Literal["whole_doc_read_done"] = "whole_doc_read_done"
|
||||||
|
completed: int
|
||||||
|
slices: int
|
||||||
|
duration_seconds: float
|
||||||
|
|
||||||
|
|
||||||
|
type ProgressEvent = Annotated[
|
||||||
|
WholeDocReadStarted | WholeDocSliceDone | WholeDocCompressionRound | WholeDocReadDone,
|
||||||
|
Field(discriminator="phase"),
|
||||||
|
]
|
||||||
@@ -1,39 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from pydantic import Field
|
|
||||||
|
|
||||||
from stirling.models import ApiModel
|
|
||||||
|
|
||||||
from .common import FileId
|
|
||||||
|
|
||||||
|
|
||||||
class IngestedPageText(ApiModel):
|
|
||||||
page_number: int = Field(ge=1)
|
|
||||||
text: str
|
|
||||||
|
|
||||||
|
|
||||||
class IngestDocumentRequest(ApiModel):
|
|
||||||
"""Replace-ingest a document's content into RAG under the given document_id.
|
|
||||||
|
|
||||||
Each content-type field is optional; the endpoint replaces the document's entire
|
|
||||||
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 IngestDocumentResponse(ApiModel):
|
|
||||||
document_id: FileId
|
|
||||||
chunks_indexed: int
|
|
||||||
|
|
||||||
|
|
||||||
class DeleteDocumentResponse(ApiModel):
|
|
||||||
document_id: FileId
|
|
||||||
deleted: bool
|
|
||||||
@@ -1,4 +1,14 @@
|
|||||||
# RAG Integration Guide
|
# Document Storage
|
||||||
|
|
||||||
|
The `documents` package owns all stored content for a document under a single
|
||||||
|
`collection` (file id):
|
||||||
|
|
||||||
|
* **Vector chunks** — small, embedded chunks for RAG-style retrieval.
|
||||||
|
* **Ordered pages** — the original page text retained in document order, used
|
||||||
|
for whole-document reading.
|
||||||
|
|
||||||
|
Both representations are populated by a single `ingest()` call and removed
|
||||||
|
together by `delete_collection()`.
|
||||||
|
|
||||||
## Adding RAG to an Agent
|
## Adding RAG to an Agent
|
||||||
|
|
||||||
@@ -22,19 +32,21 @@ That's it. The agent gets a `search_knowledge` tool it can call autonomously.
|
|||||||
|
|
||||||
## Scoping to Specific Collections
|
## Scoping to Specific Collections
|
||||||
|
|
||||||
Collections are named buckets of indexed documents — think folders. By default an agent searches everything in the store. Pass `collections=` to restrict it to only the docs indexed under those names.
|
Collections are named buckets of indexed documents - think folders. By default
|
||||||
|
an agent searches everything in the store. Pass `collections=` to restrict it
|
||||||
|
to only the docs indexed under those names.
|
||||||
|
|
||||||
```python
|
```python
|
||||||
from stirling.rag import RagCapability
|
from stirling.documents import RagCapability
|
||||||
|
|
||||||
# Only searches docs indexed under "company-docs" — ignores everything else
|
# Only searches docs indexed under "company-docs"
|
||||||
scoped = RagCapability(runtime.rag_service, collections=["company-docs"], top_k=3)
|
scoped = RagCapability(runtime.documents, collections=["company-docs"], top_k=3)
|
||||||
|
|
||||||
# Searches multiple collections
|
# Searches multiple collections
|
||||||
multi = RagCapability(runtime.rag_service, collections=["company-docs", "product-specs"])
|
multi = RagCapability(runtime.documents, collections=["company-docs", "product-specs"])
|
||||||
|
|
||||||
# No collections arg = searches all collections in the store
|
# No collections arg = searches all collections in the store
|
||||||
everything = RagCapability(runtime.rag_service)
|
everything = RagCapability(runtime.documents)
|
||||||
```
|
```
|
||||||
|
|
||||||
## Config
|
## Config
|
||||||
@@ -51,7 +63,8 @@ STIRLING_RAG_CHUNK_OVERLAP=64
|
|||||||
STIRLING_RAG_TOP_K=5
|
STIRLING_RAG_TOP_K=5
|
||||||
```
|
```
|
||||||
|
|
||||||
Provider credentials (and any local overrides) go in the uncommitted `engine/.env.local`:
|
Provider credentials (and any local overrides) go in the uncommitted
|
||||||
|
`engine/.env.local`:
|
||||||
|
|
||||||
```
|
```
|
||||||
VOYAGE_API_KEY=your-key
|
VOYAGE_API_KEY=your-key
|
||||||
@@ -59,13 +72,17 @@ VOYAGE_API_KEY=your-key
|
|||||||
|
|
||||||
## Backends
|
## Backends
|
||||||
|
|
||||||
**`sqlite`** — Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev and self-hosted deployments.
|
**`sqlite`** - Embedded sqlite-vec. Single `.db` file, zero ops. Ideal for dev
|
||||||
|
and self-hosted deployments.
|
||||||
|
|
||||||
**`pgvector`** — External PostgreSQL with the `vector` extension. Point `STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
|
**`pgvector`** - External PostgreSQL with the `vector` extension. Point
|
||||||
|
`STIRLING_RAG_PGVECTOR_DSN` at your Postgres instance.
|
||||||
|
|
||||||
Both backends implement the same `VectorStore` interface, so agents and the RAG service work identically regardless of which you pick.
|
Both backends implement the same `DocumentStore` interface, so agents and the
|
||||||
|
service work identically regardless of which you pick.
|
||||||
|
|
||||||
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model string accordingly and point at the server via its native env var:
|
For a self-hosted embedding server (e.g. Ollama, TEI, vLLM) set the model
|
||||||
|
string accordingly and point at the server via its native env var:
|
||||||
|
|
||||||
```
|
```
|
||||||
# Ollama running on another machine
|
# Ollama running on another machine
|
||||||
@@ -81,8 +98,5 @@ OPENAI_BASE_URL=http://192.168.1.50:8080/v1
|
|||||||
|
|
||||||
| Method | Endpoint | Purpose |
|
| Method | Endpoint | Purpose |
|
||||||
|--------|----------|---------|
|
|--------|----------|---------|
|
||||||
| GET | `/api/v1/rag/status` | Report embedding model and existing collections |
|
| POST | `/api/v1/documents` | Replace-ingest a document's pages |
|
||||||
| POST | `/api/v1/rag/index` | Index text into a collection |
|
| DELETE | `/api/v1/documents/{document_id}` | Delete a document's stored content |
|
||||||
| POST | `/api/v1/rag/search` | Search a collection |
|
|
||||||
| GET | `/api/v1/rag/collections` | List collections |
|
|
||||||
| DELETE | `/api/v1/rag/collections/{name}` | Delete a collection |
|
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from stirling.documents.embedder import EmbeddingService
|
||||||
|
from stirling.documents.pgvector_store import PgVectorStore
|
||||||
|
from stirling.documents.rag_capability import RagCapability
|
||||||
|
from stirling.documents.service import DocumentService
|
||||||
|
from stirling.documents.sqlite_vec_store import SqliteVecStore
|
||||||
|
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"Document",
|
||||||
|
"DocumentService",
|
||||||
|
"DocumentStore",
|
||||||
|
"EmbeddingService",
|
||||||
|
"PgVectorStore",
|
||||||
|
"RagCapability",
|
||||||
|
"SearchResult",
|
||||||
|
"SqliteVecStore",
|
||||||
|
"StoredPage",
|
||||||
|
]
|
||||||
@@ -2,8 +2,8 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from pydantic_ai import Embedder
|
from pydantic_ai import Embedder
|
||||||
|
|
||||||
from stirling.rag.chunker import chunk_text
|
from stirling.documents.chunker import chunk_text
|
||||||
from stirling.rag.store import Document
|
from stirling.documents.store import Document
|
||||||
|
|
||||||
# Keep each upstream embed request under every major provider's per-call limit while
|
# 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
|
# still batching large enough that a book-sized document ingests in a reasonable number
|
||||||
+88
-6
@@ -5,14 +5,20 @@ import json
|
|||||||
import psycopg
|
import psycopg
|
||||||
from pgvector.psycopg import register_vector_async
|
from pgvector.psycopg import register_vector_async
|
||||||
|
|
||||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
from stirling.contracts.documents import Page, PageRange
|
||||||
|
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||||
|
|
||||||
|
|
||||||
class PgVectorStore(VectorStore):
|
class PgVectorStore(DocumentStore):
|
||||||
"""PostgreSQL + pgvector backed store.
|
"""PostgreSQL + pgvector backed store.
|
||||||
|
|
||||||
Connects to an external Postgres instance (DSN provided via config) and uses the
|
Connects to an external Postgres instance (DSN provided via config) and uses the
|
||||||
`vector` extension for similarity search. The schema is created on first use.
|
`vector` extension for similarity search. The schema is created on first use.
|
||||||
|
|
||||||
|
Holds two tables under the same connection:
|
||||||
|
|
||||||
|
* ``rag_documents`` - vector chunks for RAG search.
|
||||||
|
* ``document_pages`` - ordered page text for whole-document reading.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, dsn: str) -> None:
|
def __init__(self, dsn: str) -> None:
|
||||||
@@ -32,11 +38,20 @@ class PgVectorStore(VectorStore):
|
|||||||
async with await self._connect() as conn:
|
async with await self._connect() as conn:
|
||||||
async with conn.cursor() as cur:
|
async with conn.cursor() as cur:
|
||||||
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
||||||
|
await cur.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS documents_meta (
|
||||||
|
collection TEXT PRIMARY KEY,
|
||||||
|
source TEXT NOT NULL
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
await cur.execute(
|
await cur.execute(
|
||||||
"""
|
"""
|
||||||
CREATE TABLE IF NOT EXISTS rag_documents (
|
CREATE TABLE IF NOT EXISTS rag_documents (
|
||||||
id TEXT NOT NULL,
|
id TEXT NOT NULL,
|
||||||
collection TEXT NOT NULL,
|
collection TEXT NOT NULL
|
||||||
|
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||||
text TEXT NOT NULL,
|
text TEXT NOT NULL,
|
||||||
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
|
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||||||
embedding vector NOT NULL,
|
embedding vector NOT NULL,
|
||||||
@@ -45,9 +60,36 @@ class PgVectorStore(VectorStore):
|
|||||||
"""
|
"""
|
||||||
)
|
)
|
||||||
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
|
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
|
||||||
|
await cur.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS document_pages (
|
||||||
|
collection TEXT NOT NULL
|
||||||
|
REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||||
|
page_number INTEGER NOT NULL,
|
||||||
|
text TEXT NOT NULL,
|
||||||
|
char_count INTEGER NOT NULL,
|
||||||
|
PRIMARY KEY (collection, page_number)
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||||
await conn.commit()
|
await conn.commit()
|
||||||
self._initialized = True
|
self._initialized = True
|
||||||
|
|
||||||
|
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||||
|
await self._ensure_schema()
|
||||||
|
async with await self._connect() as conn:
|
||||||
|
async with conn.cursor() as cur:
|
||||||
|
await cur.execute(
|
||||||
|
"""
|
||||||
|
INSERT INTO documents_meta (collection, source)
|
||||||
|
VALUES (%s, %s)
|
||||||
|
ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source
|
||||||
|
""",
|
||||||
|
(collection, source),
|
||||||
|
)
|
||||||
|
await conn.commit()
|
||||||
|
|
||||||
async def add_documents(
|
async def add_documents(
|
||||||
self,
|
self,
|
||||||
collection: str,
|
collection: str,
|
||||||
@@ -106,18 +148,58 @@ class PgVectorStore(VectorStore):
|
|||||||
for r in rows
|
for r in rows
|
||||||
]
|
]
|
||||||
|
|
||||||
|
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||||
|
await self._ensure_schema()
|
||||||
|
async with await self._connect() as conn:
|
||||||
|
async with conn.cursor() as cur:
|
||||||
|
await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,))
|
||||||
|
if pages:
|
||||||
|
await cur.executemany(
|
||||||
|
"""
|
||||||
|
INSERT INTO document_pages (collection, page_number, text, char_count)
|
||||||
|
VALUES (%s, %s, %s, %s)
|
||||||
|
""",
|
||||||
|
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||||
|
)
|
||||||
|
await conn.commit()
|
||||||
|
|
||||||
|
async def read_pages(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
page_range: PageRange | None = None,
|
||||||
|
) -> list[Page]:
|
||||||
|
await self._ensure_schema()
|
||||||
|
async with await self._connect() as conn:
|
||||||
|
async with conn.cursor() as cur:
|
||||||
|
if page_range is None:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT page_number, text, char_count FROM document_pages "
|
||||||
|
"WHERE collection = %s ORDER BY page_number",
|
||||||
|
(collection,),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
await cur.execute(
|
||||||
|
"SELECT page_number, text, char_count FROM document_pages "
|
||||||
|
"WHERE collection = %s AND page_number BETWEEN %s AND %s "
|
||||||
|
"ORDER BY page_number",
|
||||||
|
(collection, page_range.start, page_range.end),
|
||||||
|
)
|
||||||
|
rows = await cur.fetchall()
|
||||||
|
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
|
||||||
|
|
||||||
async def delete_collection(self, collection: str) -> None:
|
async def delete_collection(self, collection: str) -> None:
|
||||||
await self._ensure_schema()
|
await self._ensure_schema()
|
||||||
async with await self._connect() as conn:
|
async with await self._connect() as conn:
|
||||||
async with conn.cursor() as cur:
|
async with conn.cursor() as cur:
|
||||||
await cur.execute("DELETE FROM rag_documents WHERE collection = %s", (collection,))
|
# Cascade FKs handle rag_documents and document_pages.
|
||||||
|
await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,))
|
||||||
await conn.commit()
|
await conn.commit()
|
||||||
|
|
||||||
async def list_collections(self) -> list[str]:
|
async def list_collections(self) -> list[str]:
|
||||||
await self._ensure_schema()
|
await self._ensure_schema()
|
||||||
async with await self._connect() as conn:
|
async with await self._connect() as conn:
|
||||||
async with conn.cursor() as cur:
|
async with conn.cursor() as cur:
|
||||||
await cur.execute("SELECT DISTINCT collection FROM rag_documents ORDER BY collection")
|
await cur.execute("SELECT collection FROM documents_meta ORDER BY collection")
|
||||||
rows = await cur.fetchall()
|
rows = await cur.fetchall()
|
||||||
return [r[0] for r in rows]
|
return [r[0] for r in rows]
|
||||||
|
|
||||||
@@ -126,7 +208,7 @@ class PgVectorStore(VectorStore):
|
|||||||
async with await self._connect() as conn:
|
async with await self._connect() as conn:
|
||||||
async with conn.cursor() as cur:
|
async with conn.cursor() as cur:
|
||||||
await cur.execute(
|
await cur.execute(
|
||||||
"SELECT 1 FROM rag_documents WHERE collection = %s LIMIT 1",
|
"SELECT 1 FROM documents_meta WHERE collection = %s",
|
||||||
(collection,),
|
(collection,),
|
||||||
)
|
)
|
||||||
row = await cur.fetchone()
|
row = await cur.fetchone()
|
||||||
+7
-7
@@ -6,9 +6,9 @@ from collections.abc import Awaitable, Callable
|
|||||||
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
from pydantic_ai import FunctionToolset, RunContext, ToolDefinition
|
||||||
from pydantic_ai.toolsets import AbstractToolset
|
from pydantic_ai.toolsets import AbstractToolset
|
||||||
|
|
||||||
|
from stirling.documents.service import DocumentService
|
||||||
|
from stirling.documents.store import SearchResult
|
||||||
from stirling.models import FileId
|
from stirling.models import FileId
|
||||||
from stirling.rag.service import RagService
|
|
||||||
from stirling.rag.store import SearchResult
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -34,12 +34,12 @@ class RagCapability:
|
|||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
rag_service: RagService,
|
documents: DocumentService,
|
||||||
collections: list[FileId] | None = None,
|
collections: list[FileId] | None = None,
|
||||||
top_k: int = 5,
|
top_k: int = 5,
|
||||||
max_searches: int = 5,
|
max_searches: int = 5,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._rag_service = rag_service
|
self._documents = documents
|
||||||
self._collections = collections
|
self._collections = collections
|
||||||
self._top_k = top_k
|
self._top_k = top_k
|
||||||
self._max_searches = max_searches
|
self._max_searches = max_searches
|
||||||
@@ -74,7 +74,7 @@ class RagCapability:
|
|||||||
)
|
)
|
||||||
|
|
||||||
async def _dynamic_instructions(self) -> str:
|
async def _dynamic_instructions(self) -> str:
|
||||||
collections = await self._rag_service.list_collections()
|
collections = await self._documents.list_collections()
|
||||||
if collections:
|
if collections:
|
||||||
names = ", ".join(collections)
|
names = ", ".join(collections)
|
||||||
collection_desc = f"the following knowledge base collections: {names}"
|
collection_desc = f"the following knowledge base collections: {names}"
|
||||||
@@ -115,12 +115,12 @@ class RagCapability:
|
|||||||
if self._collections:
|
if self._collections:
|
||||||
all_results = []
|
all_results = []
|
||||||
for col in self._collections:
|
for col in self._collections:
|
||||||
col_results = await self._rag_service.search(query, collection=col, top_k=k)
|
col_results = await self._documents.search(query, collection=col, top_k=k)
|
||||||
all_results.extend(col_results)
|
all_results.extend(col_results)
|
||||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||||
results = all_results[:k]
|
results = all_results[:k]
|
||||||
else:
|
else:
|
||||||
results = await self._rag_service.search(query, top_k=k)
|
results = await self._documents.search(query, top_k=k)
|
||||||
|
|
||||||
if not results:
|
if not results:
|
||||||
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
|
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
|
||||||
@@ -0,0 +1,133 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from stirling.contracts.documents import Page, PageRange, PageText
|
||||||
|
from stirling.documents.embedder import EmbeddingService
|
||||||
|
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||||
|
from stirling.models import FileId
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
PAGE_NUMBER_METADATA_KEY = "page_number"
|
||||||
|
CONTENT_TYPE_METADATA_KEY = "content_type"
|
||||||
|
PAGE_TEXT_CONTENT_TYPE = "page_text"
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentService:
|
||||||
|
"""Top-level facade for stored document content.
|
||||||
|
|
||||||
|
Holds two representations of every document under a single ``collection``:
|
||||||
|
|
||||||
|
* **Vector chunks** for RAG-style semantic retrieval (``search``).
|
||||||
|
* **Ordered pages** for whole-document reading (``read_pages``).
|
||||||
|
|
||||||
|
Both are populated by :meth:`ingest` from a single ``pages`` payload. Agents
|
||||||
|
pick the strategy that fits the question; they don't need to know which
|
||||||
|
storage they're hitting.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
|
||||||
|
self._embedder = embedder
|
||||||
|
self._store = store
|
||||||
|
self._default_top_k = default_top_k
|
||||||
|
|
||||||
|
async def ingest(
|
||||||
|
self,
|
||||||
|
collection: FileId,
|
||||||
|
pages: list[PageText],
|
||||||
|
source: str,
|
||||||
|
) -> int:
|
||||||
|
"""Replace-ingest a document. Returns the number of vector chunks indexed.
|
||||||
|
|
||||||
|
This wipes any previously-stored content for ``collection`` and writes
|
||||||
|
both the vector-chunk and page-text representations from the same
|
||||||
|
``pages`` payload. Pages with empty/whitespace-only text are skipped
|
||||||
|
for chunking but still written to the page store so page numbering is
|
||||||
|
preserved end-to-end.
|
||||||
|
"""
|
||||||
|
await self._store.delete_collection(collection)
|
||||||
|
await self._store.ensure_collection(collection, source)
|
||||||
|
|
||||||
|
stored_pages = [StoredPage(page_number=p.page_number, text=p.text, char_count=len(p.text)) for p in pages]
|
||||||
|
await self._store.add_pages(collection, stored_pages)
|
||||||
|
|
||||||
|
chunks: list[Document] = []
|
||||||
|
for page in pages:
|
||||||
|
if not page.text.strip():
|
||||||
|
continue
|
||||||
|
chunks.extend(
|
||||||
|
self._embedder.chunk_and_prepare(
|
||||||
|
text=page.text,
|
||||||
|
source=f"{source}:page:{page.page_number}",
|
||||||
|
base_metadata={
|
||||||
|
PAGE_NUMBER_METADATA_KEY: str(page.page_number),
|
||||||
|
CONTENT_TYPE_METADATA_KEY: PAGE_TEXT_CONTENT_TYPE,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if not chunks:
|
||||||
|
return 0
|
||||||
|
embeddings = await self._embedder.embed_documents([doc.text for doc in chunks])
|
||||||
|
await self._store.add_documents(collection, chunks, embeddings)
|
||||||
|
return len(chunks)
|
||||||
|
|
||||||
|
async def search(
|
||||||
|
self,
|
||||||
|
query: str,
|
||||||
|
collection: FileId | None = None,
|
||||||
|
top_k: int | None = None,
|
||||||
|
) -> list[SearchResult]:
|
||||||
|
"""Embed query and search across one or all collections.
|
||||||
|
|
||||||
|
If collection is None, searches all available collections and merges results.
|
||||||
|
"""
|
||||||
|
k = top_k if top_k is not None else self._default_top_k
|
||||||
|
query_embedding = await self._embedder.embed_query(query)
|
||||||
|
|
||||||
|
if collection is not None:
|
||||||
|
if not await self._store.has_collection(collection):
|
||||||
|
return []
|
||||||
|
return await self._store.search(collection, query_embedding, k)
|
||||||
|
|
||||||
|
# Search all collections, skipping any that error (e.g. dimension mismatch)
|
||||||
|
collections = await self._store.list_collections()
|
||||||
|
all_results: list[SearchResult] = []
|
||||||
|
for col_name in collections:
|
||||||
|
try:
|
||||||
|
results = await self._store.search(col_name, query_embedding, k)
|
||||||
|
all_results.extend(results)
|
||||||
|
except Exception: # noqa: BLE001 - any backend error on one collection should not stop the others
|
||||||
|
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
|
||||||
|
|
||||||
|
# Sort by score descending, return top_k across all collections
|
||||||
|
all_results.sort(key=lambda r: r.score, reverse=True)
|
||||||
|
return all_results[:k]
|
||||||
|
|
||||||
|
async def read_pages(
|
||||||
|
self,
|
||||||
|
collection: FileId,
|
||||||
|
page_range: PageRange | None = None,
|
||||||
|
) -> list[Page]:
|
||||||
|
"""Return ordered page text for ``collection``.
|
||||||
|
|
||||||
|
Empty list if the collection has no stored pages.
|
||||||
|
"""
|
||||||
|
return await self._store.read_pages(collection, page_range)
|
||||||
|
|
||||||
|
async def delete_collection(self, collection: FileId) -> None:
|
||||||
|
"""Remove a collection's chunks and pages."""
|
||||||
|
await self._store.delete_collection(collection)
|
||||||
|
|
||||||
|
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."""
|
||||||
|
return [FileId(name) for name in await self._store.list_collections()]
|
||||||
|
|
||||||
|
async def close(self) -> None:
|
||||||
|
"""Release the underlying store's resources."""
|
||||||
|
await self._store.close()
|
||||||
+90
-12
@@ -9,10 +9,11 @@ from pathlib import Path
|
|||||||
|
|
||||||
import sqlite_vec
|
import sqlite_vec
|
||||||
|
|
||||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
from stirling.contracts.documents import Page, PageRange
|
||||||
|
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
|
||||||
|
|
||||||
|
|
||||||
class SqliteVecStore(VectorStore):
|
class SqliteVecStore(DocumentStore):
|
||||||
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
|
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
|
||||||
|
|
||||||
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
|
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
|
||||||
@@ -32,6 +33,8 @@ class SqliteVecStore(VectorStore):
|
|||||||
conn.enable_load_extension(True)
|
conn.enable_load_extension(True)
|
||||||
sqlite_vec.load(conn)
|
sqlite_vec.load(conn)
|
||||||
conn.enable_load_extension(False)
|
conn.enable_load_extension(False)
|
||||||
|
# Required so cascade deletes from documents_meta clean up child tables.
|
||||||
|
conn.execute("PRAGMA foreign_keys=ON")
|
||||||
if self._db_path is not None:
|
if self._db_path is not None:
|
||||||
conn.execute("PRAGMA journal_mode=WAL")
|
conn.execute("PRAGMA journal_mode=WAL")
|
||||||
|
|
||||||
@@ -45,10 +48,18 @@ class SqliteVecStore(VectorStore):
|
|||||||
return cls(":memory:")
|
return cls(":memory:")
|
||||||
|
|
||||||
def _init_schema(self) -> None:
|
def _init_schema(self) -> None:
|
||||||
|
self._conn.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS documents_meta (
|
||||||
|
collection TEXT PRIMARY KEY,
|
||||||
|
source TEXT NOT NULL
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
self._conn.execute(
|
self._conn.execute(
|
||||||
"""
|
"""
|
||||||
CREATE TABLE IF NOT EXISTS collections (
|
CREATE TABLE IF NOT EXISTS collections (
|
||||||
name TEXT PRIMARY KEY,
|
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||||
dim INTEGER NOT NULL,
|
dim INTEGER NOT NULL,
|
||||||
table_name TEXT NOT NULL
|
table_name TEXT NOT NULL
|
||||||
)
|
)
|
||||||
@@ -58,7 +69,7 @@ class SqliteVecStore(VectorStore):
|
|||||||
"""
|
"""
|
||||||
CREATE TABLE IF NOT EXISTS documents (
|
CREATE TABLE IF NOT EXISTS documents (
|
||||||
id TEXT NOT NULL,
|
id TEXT NOT NULL,
|
||||||
collection TEXT NOT NULL,
|
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||||
text TEXT NOT NULL,
|
text TEXT NOT NULL,
|
||||||
metadata TEXT NOT NULL DEFAULT '{}',
|
metadata TEXT NOT NULL DEFAULT '{}',
|
||||||
vec_rowid INTEGER NOT NULL,
|
vec_rowid INTEGER NOT NULL,
|
||||||
@@ -67,6 +78,32 @@ class SqliteVecStore(VectorStore):
|
|||||||
"""
|
"""
|
||||||
)
|
)
|
||||||
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
|
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
|
||||||
|
self._conn.execute(
|
||||||
|
"""
|
||||||
|
CREATE TABLE IF NOT EXISTS document_pages (
|
||||||
|
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||||
|
page_number INTEGER NOT NULL,
|
||||||
|
text TEXT NOT NULL,
|
||||||
|
char_count INTEGER NOT NULL,
|
||||||
|
PRIMARY KEY (collection, page_number)
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
|
||||||
|
self._conn.commit()
|
||||||
|
|
||||||
|
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||||
|
async with self._lock:
|
||||||
|
await asyncio.to_thread(self._sync_ensure_collection, collection, source)
|
||||||
|
|
||||||
|
def _sync_ensure_collection(self, collection: str, source: str) -> None:
|
||||||
|
self._conn.execute(
|
||||||
|
"""
|
||||||
|
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
|
||||||
|
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
|
||||||
|
""",
|
||||||
|
(collection, source),
|
||||||
|
)
|
||||||
self._conn.commit()
|
self._conn.commit()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@@ -196,18 +233,56 @@ class SqliteVecStore(VectorStore):
|
|||||||
for r in results
|
for r in results
|
||||||
]
|
]
|
||||||
|
|
||||||
|
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||||
|
async with self._lock:
|
||||||
|
await asyncio.to_thread(self._sync_add_pages, collection, pages)
|
||||||
|
|
||||||
|
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||||
|
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,))
|
||||||
|
if pages:
|
||||||
|
self._conn.executemany(
|
||||||
|
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)",
|
||||||
|
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||||
|
)
|
||||||
|
self._conn.commit()
|
||||||
|
|
||||||
|
async def read_pages(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
page_range: PageRange | None = None,
|
||||||
|
) -> list[Page]:
|
||||||
|
async with self._lock:
|
||||||
|
return await asyncio.to_thread(self._sync_read_pages, collection, page_range)
|
||||||
|
|
||||||
|
def _sync_read_pages(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
page_range: PageRange | None,
|
||||||
|
) -> list[Page]:
|
||||||
|
if page_range is None:
|
||||||
|
rows = self._conn.execute(
|
||||||
|
"SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number",
|
||||||
|
(collection,),
|
||||||
|
).fetchall()
|
||||||
|
else:
|
||||||
|
rows = self._conn.execute(
|
||||||
|
"SELECT page_number, text, char_count FROM document_pages "
|
||||||
|
"WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number",
|
||||||
|
(collection, page_range.start, page_range.end),
|
||||||
|
).fetchall()
|
||||||
|
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
|
||||||
|
|
||||||
async def delete_collection(self, collection: str) -> None:
|
async def delete_collection(self, collection: str) -> None:
|
||||||
async with self._lock:
|
async with self._lock:
|
||||||
await asyncio.to_thread(self._sync_delete_collection, collection)
|
await asyncio.to_thread(self._sync_delete_collection, collection)
|
||||||
|
|
||||||
def _sync_delete_collection(self, collection: str) -> None:
|
def _sync_delete_collection(self, collection: str) -> None:
|
||||||
|
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
|
||||||
|
# (collections, documents, document_pages) when documents_meta is deleted.
|
||||||
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||||
if row is None:
|
if row is not None:
|
||||||
return
|
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
|
||||||
table_name = row[0]
|
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,))
|
||||||
self._conn.execute(f"DROP TABLE IF EXISTS {table_name}")
|
|
||||||
self._conn.execute("DELETE FROM documents WHERE collection = ?", (collection,))
|
|
||||||
self._conn.execute("DELETE FROM collections WHERE name = ?", (collection,))
|
|
||||||
self._conn.commit()
|
self._conn.commit()
|
||||||
|
|
||||||
async def list_collections(self) -> list[str]:
|
async def list_collections(self) -> list[str]:
|
||||||
@@ -215,7 +290,7 @@ class SqliteVecStore(VectorStore):
|
|||||||
return await asyncio.to_thread(self._sync_list_collections)
|
return await asyncio.to_thread(self._sync_list_collections)
|
||||||
|
|
||||||
def _sync_list_collections(self) -> list[str]:
|
def _sync_list_collections(self) -> list[str]:
|
||||||
rows = self._conn.execute("SELECT name FROM collections ORDER BY name").fetchall()
|
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall()
|
||||||
return [r[0] for r in rows]
|
return [r[0] for r in rows]
|
||||||
|
|
||||||
async def has_collection(self, collection: str) -> bool:
|
async def has_collection(self, collection: str) -> bool:
|
||||||
@@ -223,7 +298,10 @@ class SqliteVecStore(VectorStore):
|
|||||||
return await asyncio.to_thread(self._sync_has_collection, collection)
|
return await asyncio.to_thread(self._sync_has_collection, collection)
|
||||||
|
|
||||||
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 documents_meta WHERE collection = ?",
|
||||||
|
(collection,),
|
||||||
|
).fetchone()
|
||||||
return row is not None
|
return row is not None
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
@@ -0,0 +1,113 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from abc import ABC, abstractmethod
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
|
||||||
|
from stirling.contracts.documents import Page, PageRange
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class Document:
|
||||||
|
"""A chunk of text with metadata, ready for embedding and storage."""
|
||||||
|
|
||||||
|
id: str
|
||||||
|
text: str
|
||||||
|
metadata: dict[str, str] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SearchResult:
|
||||||
|
"""A document returned from a vector search with its relevance score."""
|
||||||
|
|
||||||
|
document: Document
|
||||||
|
score: float
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class StoredPage:
|
||||||
|
"""A page as written to the store. ``char_count`` is precomputed at ingest."""
|
||||||
|
|
||||||
|
page_number: int
|
||||||
|
text: str
|
||||||
|
char_count: int
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentStore(ABC):
|
||||||
|
"""Abstract interface for document storage backends.
|
||||||
|
|
||||||
|
Backends hold two representations of every document:
|
||||||
|
|
||||||
|
* **Vector chunks** - small, embedded chunks used for RAG search.
|
||||||
|
* **Ordered pages** - the original page text retained in document order,
|
||||||
|
used for whole-document reading.
|
||||||
|
|
||||||
|
Both representations live under the same ``collection`` (file id) and are
|
||||||
|
rooted at a single parent row in ``documents_meta``. Removing that parent
|
||||||
|
row cascades to both child representations, so :meth:`delete_collection`
|
||||||
|
is one logical delete.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||||
|
"""Upsert the top-level ``documents_meta`` row for this collection.
|
||||||
|
|
||||||
|
Must be called before :meth:`add_pages` or :meth:`add_documents`. Both
|
||||||
|
of those write into child tables that hold a foreign key to the parent
|
||||||
|
row, so it must exist first.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def add_documents(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
documents: list[Document],
|
||||||
|
embeddings: list[list[float]],
|
||||||
|
) -> None:
|
||||||
|
"""Store vector chunks with their embeddings in the named collection."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def search(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
query_embedding: list[float],
|
||||||
|
top_k: int = 5,
|
||||||
|
) -> list[SearchResult]:
|
||||||
|
"""Return the top_k most similar vector chunks from the collection."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||||
|
"""Replace the stored pages for ``collection`` with the supplied pages.
|
||||||
|
|
||||||
|
Implementations must remove any previously-stored pages for the
|
||||||
|
collection before writing, so callers can re-ingest by calling this
|
||||||
|
method again.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def read_pages(
|
||||||
|
self,
|
||||||
|
collection: str,
|
||||||
|
page_range: PageRange | None = None,
|
||||||
|
) -> list[Page]:
|
||||||
|
"""Return ordered pages for ``collection``.
|
||||||
|
|
||||||
|
If ``page_range`` is ``None`` all pages are returned. Otherwise only
|
||||||
|
pages whose ``page_number`` falls within the inclusive range are
|
||||||
|
returned. Pages are always ordered by ``page_number`` ascending.
|
||||||
|
"""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def delete_collection(self, collection: str) -> None:
|
||||||
|
"""Remove a collection's chunks and pages."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def list_collections(self) -> list[str]:
|
||||||
|
"""Return names of all existing collections."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def has_collection(self, collection: str) -> bool:
|
||||||
|
"""Check whether a collection exists."""
|
||||||
|
|
||||||
|
@abstractmethod
|
||||||
|
async def close(self) -> None:
|
||||||
|
"""Release any resources held by the store (connections, handles, etc.)."""
|
||||||
@@ -12,9 +12,12 @@ FileId = NewType("FileId", str)
|
|||||||
|
|
||||||
|
|
||||||
class ApiModel(BaseModel):
|
class ApiModel(BaseModel):
|
||||||
|
"""Base for every contract model crossing a service boundary."""
|
||||||
|
|
||||||
model_config = ConfigDict(
|
model_config = ConfigDict(
|
||||||
alias_generator=to_camel,
|
alias_generator=to_camel,
|
||||||
extra="forbid",
|
extra="forbid",
|
||||||
validate_by_name=True,
|
validate_by_name=True,
|
||||||
validate_by_alias=True,
|
validate_by_alias=True,
|
||||||
|
serialize_by_alias=True,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,19 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from stirling.rag.capability import RagCapability
|
|
||||||
from stirling.rag.embedder import EmbeddingService
|
|
||||||
from stirling.rag.pgvector_store import PgVectorStore
|
|
||||||
from stirling.rag.service import RagService
|
|
||||||
from stirling.rag.sqlite_vec_store import SqliteVecStore
|
|
||||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
|
||||||
|
|
||||||
__all__ = [
|
|
||||||
"Document",
|
|
||||||
"EmbeddingService",
|
|
||||||
"PgVectorStore",
|
|
||||||
"RagCapability",
|
|
||||||
"RagService",
|
|
||||||
"SearchResult",
|
|
||||||
"SqliteVecStore",
|
|
||||||
"VectorStore",
|
|
||||||
]
|
|
||||||
@@ -1,102 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import logging
|
|
||||||
|
|
||||||
from stirling.models import FileId
|
|
||||||
from stirling.rag.embedder import EmbeddingService
|
|
||||||
from stirling.rag.store import Document, SearchResult, VectorStore
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
class RagService:
|
|
||||||
"""Orchestrates embedding and vector storage for RAG workflows."""
|
|
||||||
|
|
||||||
def __init__(self, embedder: EmbeddingService, store: VectorStore, default_top_k: int = 5) -> None:
|
|
||||||
self._embedder = embedder
|
|
||||||
self._store = store
|
|
||||||
self._default_top_k = default_top_k
|
|
||||||
|
|
||||||
async def index_text(
|
|
||||||
self,
|
|
||||||
collection: FileId,
|
|
||||||
text: str,
|
|
||||||
source: str = "",
|
|
||||||
metadata: dict[str, str] | None = None,
|
|
||||||
) -> int:
|
|
||||||
"""Chunk, embed, and store text. Returns the number of chunks indexed."""
|
|
||||||
documents = self._embedder.chunk_and_prepare(text, source=source, base_metadata=metadata)
|
|
||||||
if not documents:
|
|
||||||
return 0
|
|
||||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
|
||||||
await self._store.add_documents(collection, documents, embeddings)
|
|
||||||
return len(documents)
|
|
||||||
|
|
||||||
async def index_documents(self, collection: FileId, documents: list[Document]) -> int:
|
|
||||||
"""Embed and store pre-chunked documents. Returns the number stored."""
|
|
||||||
if not documents:
|
|
||||||
return 0
|
|
||||||
embeddings = await self._embedder.embed_documents([doc.text for doc in documents])
|
|
||||||
await self._store.add_documents(collection, documents, embeddings)
|
|
||||||
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(
|
|
||||||
self,
|
|
||||||
query: str,
|
|
||||||
collection: FileId | None = None,
|
|
||||||
top_k: int | None = None,
|
|
||||||
) -> list[SearchResult]:
|
|
||||||
"""Embed query and search across one or all collections.
|
|
||||||
|
|
||||||
If collection is None, searches all available collections and merges results.
|
|
||||||
"""
|
|
||||||
k = top_k if top_k is not None else self._default_top_k
|
|
||||||
query_embedding = await self._embedder.embed_query(query)
|
|
||||||
|
|
||||||
if collection is not None:
|
|
||||||
if not await self._store.has_collection(collection):
|
|
||||||
return []
|
|
||||||
return await self._store.search(collection, query_embedding, k)
|
|
||||||
|
|
||||||
# Search all collections, skipping any that error (e.g. dimension mismatch)
|
|
||||||
collections = await self._store.list_collections()
|
|
||||||
all_results: list[SearchResult] = []
|
|
||||||
for col_name in collections:
|
|
||||||
try:
|
|
||||||
results = await self._store.search(col_name, query_embedding, k)
|
|
||||||
all_results.extend(results)
|
|
||||||
except Exception: # noqa: BLE001 — any backend error on one collection should not stop the others
|
|
||||||
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
|
|
||||||
|
|
||||||
# Sort by score descending, return top_k across all collections
|
|
||||||
all_results.sort(key=lambda r: r.score, reverse=True)
|
|
||||||
return all_results[:k]
|
|
||||||
|
|
||||||
async def delete_collection(self, collection: FileId) -> None:
|
|
||||||
"""Remove a collection and all its documents."""
|
|
||||||
await self._store.delete_collection(collection)
|
|
||||||
|
|
||||||
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."""
|
|
||||||
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()
|
|
||||||
@@ -1,63 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from abc import ABC, abstractmethod
|
|
||||||
from dataclasses import dataclass, field
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class Document:
|
|
||||||
"""A chunk of text with metadata, ready for embedding and storage."""
|
|
||||||
|
|
||||||
id: str
|
|
||||||
text: str
|
|
||||||
metadata: dict[str, str] = field(default_factory=dict)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class SearchResult:
|
|
||||||
"""A document returned from a vector search with its relevance score."""
|
|
||||||
|
|
||||||
document: Document
|
|
||||||
score: float
|
|
||||||
|
|
||||||
|
|
||||||
class VectorStore(ABC):
|
|
||||||
"""Abstract interface for vector storage backends.
|
|
||||||
|
|
||||||
Implementations must handle persistence, collection management,
|
|
||||||
and nearest-neighbor search over pre-computed embeddings.
|
|
||||||
"""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def add_documents(
|
|
||||||
self,
|
|
||||||
collection: str,
|
|
||||||
documents: list[Document],
|
|
||||||
embeddings: list[list[float]],
|
|
||||||
) -> None:
|
|
||||||
"""Store documents with their embeddings in the named collection."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def search(
|
|
||||||
self,
|
|
||||||
collection: str,
|
|
||||||
query_embedding: list[float],
|
|
||||||
top_k: int = 5,
|
|
||||||
) -> list[SearchResult]:
|
|
||||||
"""Return the top_k most similar documents from the collection."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def delete_collection(self, collection: str) -> None:
|
|
||||||
"""Remove a collection and all its documents."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def list_collections(self) -> list[str]:
|
|
||||||
"""Return names of all existing collections."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def has_collection(self, collection: str) -> bool:
|
|
||||||
"""Check whether a collection exists."""
|
|
||||||
|
|
||||||
@abstractmethod
|
|
||||||
async def close(self) -> None:
|
|
||||||
"""Release any resources held by the store (connections, handles, etc.)."""
|
|
||||||
@@ -1,11 +1,21 @@
|
|||||||
"""Shared services used by the Stirling AI runtime."""
|
"""Shared services used by the Stirling AI runtime."""
|
||||||
|
|
||||||
|
from .progress import (
|
||||||
|
ProgressEmitter,
|
||||||
|
emit_progress,
|
||||||
|
reset_progress_emitter,
|
||||||
|
set_progress_emitter,
|
||||||
|
)
|
||||||
from .runtime import AppRuntime, build_model_settings, build_runtime
|
from .runtime import AppRuntime, build_model_settings, build_runtime
|
||||||
from .tracking import setup_posthog_tracking
|
from .tracking import setup_posthog_tracking
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"AppRuntime",
|
"AppRuntime",
|
||||||
|
"ProgressEmitter",
|
||||||
"build_model_settings",
|
"build_model_settings",
|
||||||
"build_runtime",
|
"build_runtime",
|
||||||
|
"emit_progress",
|
||||||
|
"reset_progress_emitter",
|
||||||
|
"set_progress_emitter",
|
||||||
"setup_posthog_tracking",
|
"setup_posthog_tracking",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -0,0 +1,47 @@
|
|||||||
|
"""Per-request progress emission, plumbed via a ContextVar so deep call stacks
|
||||||
|
can publish typed events to the streaming orchestrator endpoint without every
|
||||||
|
intermediate layer knowing about it.
|
||||||
|
|
||||||
|
Outside a streaming request no emitter is bound and ``emit_progress`` is a
|
||||||
|
no-op, so callers in agents/services can emit unconditionally.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from collections.abc import Awaitable, Callable
|
||||||
|
from contextvars import ContextVar, Token
|
||||||
|
|
||||||
|
from stirling.contracts import ProgressEvent
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
type ProgressEmitter = Callable[[ProgressEvent], Awaitable[None]]
|
||||||
|
|
||||||
|
_emitter: ContextVar[ProgressEmitter | None] = ContextVar("stirling_progress_emitter", default=None)
|
||||||
|
|
||||||
|
|
||||||
|
def set_progress_emitter(emitter: ProgressEmitter | None) -> Token[ProgressEmitter | None]:
|
||||||
|
return _emitter.set(emitter)
|
||||||
|
|
||||||
|
|
||||||
|
def reset_progress_emitter(token: Token[ProgressEmitter | None]) -> None:
|
||||||
|
_emitter.reset(token)
|
||||||
|
|
||||||
|
|
||||||
|
async def emit_progress(event: ProgressEvent) -> None:
|
||||||
|
"""Publish ``event`` to the current request's emitter, if any.
|
||||||
|
|
||||||
|
Failures inside the emitter are logged and swallowed so progress emission
|
||||||
|
can never break the work it's reporting on.
|
||||||
|
"""
|
||||||
|
emitter = _emitter.get()
|
||||||
|
if emitter is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
await emitter(event)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception:
|
||||||
|
logger.exception("progress emitter raised; dropping event %r", event.phase)
|
||||||
@@ -4,28 +4,49 @@ import logging
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import assert_never
|
from typing import assert_never
|
||||||
|
|
||||||
|
import httpx
|
||||||
from pydantic_ai.models import Model, infer_model
|
from pydantic_ai.models import Model, infer_model
|
||||||
|
from pydantic_ai.models.anthropic import AnthropicModel
|
||||||
|
from pydantic_ai.providers.anthropic import AnthropicProvider
|
||||||
from pydantic_ai.settings import ModelSettings
|
from pydantic_ai.settings import ModelSettings
|
||||||
|
|
||||||
from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
|
from stirling.config import ENGINE_ROOT, AppSettings, RagBackend
|
||||||
from stirling.rag import (
|
from stirling.documents import (
|
||||||
|
DocumentService,
|
||||||
|
DocumentStore,
|
||||||
EmbeddingService,
|
EmbeddingService,
|
||||||
PgVectorStore,
|
PgVectorStore,
|
||||||
RagCapability,
|
RagCapability,
|
||||||
RagService,
|
|
||||||
SqliteVecStore,
|
SqliteVecStore,
|
||||||
VectorStore,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _build_anthropic_http_client() -> httpx.AsyncClient:
|
||||||
|
"""Build the httpx client used for Anthropic API calls.
|
||||||
|
|
||||||
|
We disable connection-pool keepalive so every request opens a fresh
|
||||||
|
TCP+TLS connection. The default HTTP/1.1 pool reuses connections that
|
||||||
|
Anthropic's front door (Cloudflare) sometimes closes silently between
|
||||||
|
requests; the next request that picks up a stale connection sends its
|
||||||
|
body into a black hole and never gets a response, hanging until our
|
||||||
|
chunked-reasoner timeout fires.
|
||||||
|
|
||||||
|
A fresh handshake costs ~150ms — rounding error against a 5-15s LLM
|
||||||
|
call. The trade is determinism: we never reuse a connection that might
|
||||||
|
have died in the pool. See ``STIRLING_HTTP_DEBUG`` traces of slice 6
|
||||||
|
on 2026-05-06 for the concrete failure mode this addresses.
|
||||||
|
"""
|
||||||
|
return httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=0))
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class AppRuntime:
|
class AppRuntime:
|
||||||
settings: AppSettings
|
settings: AppSettings
|
||||||
fast_model: Model
|
fast_model: Model
|
||||||
smart_model: Model
|
smart_model: Model
|
||||||
rag_service: RagService
|
documents: DocumentService
|
||||||
rag_capability: RagCapability
|
rag_capability: RagCapability
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -53,47 +74,62 @@ def validate_structured_output_support(model: Model, model_name: str) -> None:
|
|||||||
raise ValueError(f"Unsupported model {model_name}. This model does not support structured outputs.")
|
raise ValueError(f"Unsupported model {model_name}. This model does not support structured outputs.")
|
||||||
|
|
||||||
|
|
||||||
def _build_vector_store(settings: AppSettings) -> VectorStore:
|
def _build_document_store(settings: AppSettings) -> DocumentStore:
|
||||||
"""Build the configured vector store backend."""
|
"""Build the configured document store backend."""
|
||||||
if settings.rag_backend == RagBackend.SQLITE:
|
if settings.rag_backend == RagBackend.SQLITE:
|
||||||
store_path = settings.rag_store_path
|
store_path = settings.rag_store_path
|
||||||
# Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
|
# Treat ":memory:" as a special in-process token; otherwise resolve against the engine root.
|
||||||
if str(store_path) != ":memory:" and not store_path.is_absolute():
|
if str(store_path) != ":memory:" and not store_path.is_absolute():
|
||||||
store_path = ENGINE_ROOT / store_path
|
store_path = ENGINE_ROOT / store_path
|
||||||
logger.info("RAG backend=sqlite, db_path=%s", store_path)
|
logger.info("Document store backend=sqlite, db_path=%s", store_path)
|
||||||
return SqliteVecStore(db_path=store_path)
|
return SqliteVecStore(db_path=store_path)
|
||||||
if settings.rag_backend == RagBackend.PGVECTOR:
|
if settings.rag_backend == RagBackend.PGVECTOR:
|
||||||
logger.info("RAG backend=pgvector, dsn=<configured>")
|
logger.info("Document store backend=pgvector, dsn=<configured>")
|
||||||
return PgVectorStore(dsn=settings.rag_pgvector_dsn)
|
return PgVectorStore(dsn=settings.rag_pgvector_dsn)
|
||||||
assert_never(settings.rag_backend)
|
assert_never(settings.rag_backend)
|
||||||
|
|
||||||
|
|
||||||
def _build_rag(settings: AppSettings) -> tuple[RagService, RagCapability]:
|
def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]:
|
||||||
"""Build the RAG service and capability."""
|
"""Build the document service and the RAG-search capability that wraps it."""
|
||||||
logger.info("RAG: embedding_model=%s", settings.rag_embedding_model)
|
logger.info("Documents: embedding_model=%s", settings.rag_embedding_model)
|
||||||
embedder = EmbeddingService(
|
embedder = EmbeddingService(
|
||||||
model_name=settings.rag_embedding_model,
|
model_name=settings.rag_embedding_model,
|
||||||
chunk_size=settings.rag_chunk_size,
|
chunk_size=settings.rag_chunk_size,
|
||||||
chunk_overlap=settings.rag_chunk_overlap,
|
chunk_overlap=settings.rag_chunk_overlap,
|
||||||
)
|
)
|
||||||
store = _build_vector_store(settings)
|
store = _build_document_store(settings)
|
||||||
service = RagService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
|
service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
|
||||||
capability = RagCapability(rag_service=service, top_k=settings.rag_default_top_k)
|
capability = RagCapability(documents=service, top_k=settings.rag_default_top_k)
|
||||||
return service, capability
|
return service, capability
|
||||||
|
|
||||||
|
|
||||||
def build_runtime(settings: AppSettings) -> AppRuntime:
|
def build_runtime(settings: AppSettings) -> AppRuntime:
|
||||||
fast_model = infer_model(settings.fast_model_name)
|
fast_model = _build_model(settings.fast_model_name)
|
||||||
smart_model = infer_model(settings.smart_model_name)
|
smart_model = _build_model(settings.smart_model_name)
|
||||||
validate_structured_output_support(fast_model, settings.fast_model_name)
|
validate_structured_output_support(fast_model, settings.fast_model_name)
|
||||||
validate_structured_output_support(smart_model, settings.smart_model_name)
|
validate_structured_output_support(smart_model, settings.smart_model_name)
|
||||||
|
|
||||||
rag_service, rag_capability = _build_rag(settings)
|
documents, rag_capability = _build_documents(settings)
|
||||||
|
|
||||||
return AppRuntime(
|
return AppRuntime(
|
||||||
settings=settings,
|
settings=settings,
|
||||||
fast_model=fast_model,
|
fast_model=fast_model,
|
||||||
smart_model=smart_model,
|
smart_model=smart_model,
|
||||||
rag_service=rag_service,
|
documents=documents,
|
||||||
rag_capability=rag_capability,
|
rag_capability=rag_capability,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _build_model(model_name: str) -> Model:
|
||||||
|
"""Construct a model, injecting our keepalive-free httpx client for
|
||||||
|
Anthropic models so workers don't pick up stale pooled connections.
|
||||||
|
|
||||||
|
Other providers fall back to ``infer_model`` defaults; the stale-pool
|
||||||
|
issue is specific to the Cloudflare-fronted Anthropic API in our
|
||||||
|
observations and the fix doesn't necessarily apply elsewhere.
|
||||||
|
"""
|
||||||
|
if model_name.startswith("anthropic:"):
|
||||||
|
bare_name = model_name.removeprefix("anthropic:")
|
||||||
|
provider = AnthropicProvider(http_client=_build_anthropic_http_client())
|
||||||
|
return AnthropicModel(bare_name, provider=provider)
|
||||||
|
return infer_model(model_name)
|
||||||
|
|||||||
@@ -0,0 +1,511 @@
|
|||||||
|
"""Tests for ``ChunkedReasoner``: chunking, fan-out, and synthesis wiring.
|
||||||
|
|
||||||
|
LLM calls are stubbed at the agent boundary; the runtime fixture supplies a
|
||||||
|
``test`` model so construction succeeds without provider credentials.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from unittest.mock import AsyncMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner, ChunkNotes
|
||||||
|
from stirling.contracts import WholeDocSliceDone
|
||||||
|
from stirling.contracts.documents import Page
|
||||||
|
from stirling.services import reset_progress_emitter, set_progress_emitter
|
||||||
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class _StubAgentResult[T]:
|
||||||
|
output: T
|
||||||
|
|
||||||
|
|
||||||
|
class _Answer(BaseModel):
|
||||||
|
answer: str
|
||||||
|
|
||||||
|
|
||||||
|
def _page(n: int, text: str) -> Page:
|
||||||
|
return Page(page_number=n, text=text, char_count=len(text))
|
||||||
|
|
||||||
|
|
||||||
|
# Slicing logic
|
||||||
|
|
||||||
|
|
||||||
|
class TestSlicePages:
|
||||||
|
"""``_slice_pages`` is pure: no model calls, no I/O."""
|
||||||
|
|
||||||
|
def test_groups_consecutive_pages_under_budget(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=20)
|
||||||
|
pages = [_page(1, "abc"), _page(2, "defgh"), _page(3, "ij")]
|
||||||
|
slices = reasoner._slice_pages(pages)
|
||||||
|
|
||||||
|
assert len(slices) == 1
|
||||||
|
assert [p.page_number for p in slices[0]] == [1, 2, 3]
|
||||||
|
|
||||||
|
def test_starts_new_slice_when_budget_exceeded(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(1, "a" * 6), _page(2, "b" * 6), _page(3, "c" * 6)]
|
||||||
|
slices = reasoner._slice_pages(pages)
|
||||||
|
|
||||||
|
# Each page is 6 chars, budget 10 -> two pages would be 12 (over),
|
||||||
|
# so the slicer breaks after each page.
|
||||||
|
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
|
||||||
|
|
||||||
|
def test_oversized_page_becomes_its_own_slice(self, runtime: AppRuntime) -> None:
|
||||||
|
"""A single page larger than the budget is never split. The reasoner
|
||||||
|
accepts that this slice exceeds the budget rather than breaking page
|
||||||
|
boundaries."""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(1, "small"), _page(2, "x" * 100), _page(3, "tiny")]
|
||||||
|
slices = reasoner._slice_pages(pages)
|
||||||
|
|
||||||
|
assert [[p.page_number for p in s] for s in slices] == [[1], [2], [3]]
|
||||||
|
|
||||||
|
def test_preserves_input_order(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
|
||||||
|
pages = [_page(i, f"page-{i}") for i in range(1, 11)]
|
||||||
|
slices = reasoner._slice_pages(pages)
|
||||||
|
|
||||||
|
flattened = [p.page_number for s in slices for p in s]
|
||||||
|
assert flattened == list(range(1, 11))
|
||||||
|
|
||||||
|
|
||||||
|
# End-to-end orchestration
|
||||||
|
|
||||||
|
|
||||||
|
class TestReason:
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_runs_one_chunk_per_slice_and_synthesises(self, runtime: AppRuntime) -> None:
|
||||||
|
"""Three small pages with a generous budget produce one chunk and one extractor call;
|
||||||
|
the synthesis stage receives notes from all chunks and returns the final answer."""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=1000)
|
||||||
|
pages = [_page(1, "alpha"), _page(2, "beta"), _page(3, "gamma")]
|
||||||
|
|
||||||
|
canned_notes = ChunkNotes(pages=[1, 2, 3], summary="all three pages", facts=["fact-1"])
|
||||||
|
canned_answer = _Answer(answer="final answer")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
|
||||||
|
):
|
||||||
|
result = await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="summarise this",
|
||||||
|
answer_prompt="answer the question from the notes",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result == canned_answer
|
||||||
|
assert chunk_mock.await_count == 1
|
||||||
|
synth_mock.assert_awaited_once()
|
||||||
|
synth_args = synth_mock.await_args
|
||||||
|
assert synth_args is not None
|
||||||
|
# _synthesise(question, notes, answer_prompt, answer_type)
|
||||||
|
_, notes_arg, _, type_arg = synth_args.args
|
||||||
|
assert notes_arg == [canned_notes]
|
||||||
|
assert type_arg is _Answer
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_fans_out_when_pages_exceed_slice_budget(self, runtime: AppRuntime) -> None:
|
||||||
|
"""Pages that don't fit into a single slice produce one extractor call per slice."""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(i, "x" * 8) for i in range(1, 6)]
|
||||||
|
|
||||||
|
canned_notes = ChunkNotes(pages=[0], summary="placeholder")
|
||||||
|
canned_answer = _Answer(answer="ok")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner, "_extract_chunk", AsyncMock(return_value=(canned_notes, 0.0))) as chunk_mock,
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)),
|
||||||
|
):
|
||||||
|
await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="aggregate",
|
||||||
|
answer_prompt="answer",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 5 pages, budget 10, each page 8 chars -> 5 slices -> 5 chunk calls.
|
||||||
|
assert chunk_mock.await_count == 5
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_skips_first_round_chunks_that_raise_and_continues(self, runtime: AppRuntime) -> None:
|
||||||
|
"""First-round chunks have no fallback notes, so a failure is dropped
|
||||||
|
rather than preserving anything; the surviving notes still flow into
|
||||||
|
synthesis."""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(i, "x" * 8) for i in range(1, 4)]
|
||||||
|
|
||||||
|
good = ChunkNotes(pages=[1], summary="ok")
|
||||||
|
async_results = [good, RuntimeError("chunk boom"), good]
|
||||||
|
|
||||||
|
async def _chunk(*_args: object, **_kwargs: object) -> tuple[ChunkNotes, float]:
|
||||||
|
value = async_results.pop(0)
|
||||||
|
if isinstance(value, BaseException):
|
||||||
|
raise value
|
||||||
|
return value, 0.0
|
||||||
|
|
||||||
|
canned_answer = _Answer(answer="resilient")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=_chunk)),
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock(return_value=canned_answer)) as synth_mock,
|
||||||
|
):
|
||||||
|
result = await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="aggregate",
|
||||||
|
answer_prompt="answer",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result == canned_answer
|
||||||
|
synth_args = synth_mock.await_args
|
||||||
|
assert synth_args is not None
|
||||||
|
_, notes_arg, _, _ = synth_args.args
|
||||||
|
assert len(notes_arg) == 2
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_raises_when_every_first_round_chunk_fails(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(i, "x" * 8) for i in range(1, 3)]
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner, "_extract_chunk", AsyncMock(side_effect=RuntimeError("boom"))),
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
|
||||||
|
pytest.raises(RuntimeError, match="no notes"),
|
||||||
|
):
|
||||||
|
await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="anything",
|
||||||
|
answer_prompt="answer",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
synth_mock.assert_not_awaited()
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_rejects_empty_pages(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime)
|
||||||
|
with pytest.raises(ValueError, match="at least one page"):
|
||||||
|
await reasoner.reason(
|
||||||
|
pages=[],
|
||||||
|
question="x",
|
||||||
|
answer_prompt="y",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_progress_events_carry_monotonic_completion_counter(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10)
|
||||||
|
pages = [_page(i, "x" * 8) for i in range(1, 4)]
|
||||||
|
|
||||||
|
# Each chunk's worker awaits a different release event; we release them in
|
||||||
|
# reverse order so completion order is the inverse of slice order.
|
||||||
|
release_events = [asyncio.Event() for _ in pages]
|
||||||
|
next_call_index = 0
|
||||||
|
|
||||||
|
async def _gated_worker(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||||
|
nonlocal next_call_index
|
||||||
|
this_call = next_call_index
|
||||||
|
next_call_index += 1
|
||||||
|
await release_events[this_call].wait()
|
||||||
|
return _StubAgentResult(output=ChunkNotes(pages=[this_call + 1], summary=f"slice-{this_call}"))
|
||||||
|
|
||||||
|
emitted: list[WholeDocSliceDone] = []
|
||||||
|
|
||||||
|
async def _capture_emitter(event: object) -> None:
|
||||||
|
if isinstance(event, WholeDocSliceDone):
|
||||||
|
emitted.append(event)
|
||||||
|
|
||||||
|
async def _release_in_reverse() -> None:
|
||||||
|
# Wait briefly so all three worker tasks are blocked on their events
|
||||||
|
# before we start releasing them.
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
for ev in reversed(release_events):
|
||||||
|
ev.set()
|
||||||
|
# Yield so the just-released worker can run to completion before
|
||||||
|
# we release the next one — keeps ordering deterministic.
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
token = set_progress_emitter(_capture_emitter)
|
||||||
|
try:
|
||||||
|
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_gated_worker)):
|
||||||
|
gather_task = asyncio.create_task(reasoner.gather_notes(pages, "anything"))
|
||||||
|
await _release_in_reverse()
|
||||||
|
notes = await gather_task
|
||||||
|
finally:
|
||||||
|
reset_progress_emitter(token)
|
||||||
|
|
||||||
|
assert len(notes) == 3
|
||||||
|
assert [event.completed for event in emitted] == [1, 2, 3]
|
||||||
|
assert all(event.total == 3 for event in emitted)
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_worker_timeout_is_terminal_for_the_chunk(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
|
||||||
|
pages = [_page(1, "x" * 8), _page(2, "y" * 8)]
|
||||||
|
attempts = 0
|
||||||
|
|
||||||
|
async def _hang_forever(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||||
|
nonlocal attempts
|
||||||
|
attempts += 1
|
||||||
|
await asyncio.sleep(10)
|
||||||
|
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang_forever)),
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
|
||||||
|
pytest.raises(RuntimeError, match="no notes"),
|
||||||
|
):
|
||||||
|
await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="anything",
|
||||||
|
answer_prompt="answer",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
# One attempt per slice; no retry path.
|
||||||
|
assert attempts == len(pages)
|
||||||
|
synth_mock.assert_not_awaited()
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_worker_timeout_drops_stalled_chunks(self, runtime: AppRuntime) -> None:
|
||||||
|
"""A worker that exceeds ``worker_timeout_seconds`` is abandoned, not awaited.
|
||||||
|
|
||||||
|
Without this guard one stuck upstream call would pin gather_notes to its
|
||||||
|
provider HTTP timeout (~10 minutes), starving the orchestrator request.
|
||||||
|
"""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10, worker_timeout_seconds=0.05)
|
||||||
|
pages = [_page(i, "x" * 8) for i in range(1, 4)]
|
||||||
|
|
||||||
|
async def _hang(*_args: object, **_kwargs: object) -> _StubAgentResult[ChunkNotes]:
|
||||||
|
await asyncio.sleep(10)
|
||||||
|
return _StubAgentResult(output=ChunkNotes(pages=[0], summary="never"))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_hang)),
|
||||||
|
patch.object(reasoner, "_synthesise", AsyncMock()) as synth_mock,
|
||||||
|
pytest.raises(RuntimeError, match="no notes"),
|
||||||
|
):
|
||||||
|
await reasoner.reason(
|
||||||
|
pages=pages,
|
||||||
|
question="anything",
|
||||||
|
answer_prompt="answer",
|
||||||
|
answer_type=_Answer,
|
||||||
|
)
|
||||||
|
|
||||||
|
synth_mock.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
# Prompt construction
|
||||||
|
|
||||||
|
|
||||||
|
class TestPromptConstruction:
|
||||||
|
def test_extraction_prompt_includes_question_and_page_markers(self, runtime: AppRuntime) -> None:
|
||||||
|
"""A first-round chunk's content carries ``[Page N]`` markers; the
|
||||||
|
extraction prompt prepends the user question."""
|
||||||
|
reasoner = ChunkedReasoner(runtime)
|
||||||
|
chunk = reasoner._chunk_from_pages([_page(2, "page two body"), _page(3, "page three body")])
|
||||||
|
prompt = reasoner._build_extraction_prompt(chunk.content, "what is on page two?")
|
||||||
|
|
||||||
|
assert "what is on page two?" in prompt
|
||||||
|
assert "[Page 2]" in prompt
|
||||||
|
assert "[Page 3]" in prompt
|
||||||
|
assert "page two body" in prompt
|
||||||
|
|
||||||
|
def test_format_notes_groups_by_page_label(self) -> None:
|
||||||
|
notes = [
|
||||||
|
ChunkNotes(pages=[1], summary="single", facts=["f-1"]),
|
||||||
|
ChunkNotes(pages=[2, 3, 4], summary="range", relevant_excerpts=["quote-1"]),
|
||||||
|
]
|
||||||
|
rendered = ChunkedReasoner.format_notes(notes)
|
||||||
|
|
||||||
|
assert "[Notes from page 1]" in rendered
|
||||||
|
assert "[Notes from pages 2-4]" in rendered
|
||||||
|
assert "f-1" in rendered
|
||||||
|
assert "quote-1" in rendered
|
||||||
|
|
||||||
|
|
||||||
|
# Hierarchical compression
|
||||||
|
#
|
||||||
|
# The compression loop is part of ``_run_chunks`` and isn't exposed
|
||||||
|
# directly, so these tests drive it end-to-end via ``gather_notes`` with a
|
||||||
|
# stubbed extractor that controls per-call output (and per-call failure
|
||||||
|
# patterns) by counting calls.
|
||||||
|
|
||||||
|
|
||||||
|
class TestCompression:
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_no_compression_when_under_budget(self, runtime: AppRuntime) -> None:
|
||||||
|
"""First-round notes that already fit the budget result in zero
|
||||||
|
compression rounds: the only extractor calls are one per slice."""
|
||||||
|
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=10_000)
|
||||||
|
pages = [_page(i, "x" * 150) for i in range(1, 5)] # each page exceeds slice budget alone -> 4 slices
|
||||||
|
|
||||||
|
canned = _ExtractedNotes(summary="ok")
|
||||||
|
with patch.object(
|
||||||
|
reasoner._extractor,
|
||||||
|
"run",
|
||||||
|
AsyncMock(return_value=_StubAgentResult(output=canned)),
|
||||||
|
) as ext_mock:
|
||||||
|
notes = await reasoner.gather_notes(pages, "anything")
|
||||||
|
|
||||||
|
assert ext_mock.await_count == 4
|
||||||
|
assert len(notes) == 4
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_runs_compression_when_over_budget(self, runtime: AppRuntime) -> None:
|
||||||
|
"""When first-round notes overflow the budget, the loop regroups them
|
||||||
|
and runs the extractor again. Output is shorter than input; pages from
|
||||||
|
every input slice survive in the consolidated notes."""
|
||||||
|
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=200)
|
||||||
|
pages = [_page(i, "x" * 150) for i in range(1, 5)]
|
||||||
|
|
||||||
|
call_count = 0
|
||||||
|
|
||||||
|
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
|
||||||
|
nonlocal call_count
|
||||||
|
call_count += 1
|
||||||
|
if call_count <= 4:
|
||||||
|
# Round 1: each note ~80 chars rendered. 4 * 80 = 320 chars, over the 200 budget.
|
||||||
|
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
|
||||||
|
# Round 2: smaller note so the post-round set fits the budget.
|
||||||
|
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
|
||||||
|
|
||||||
|
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)) as ext_mock:
|
||||||
|
notes = await reasoner.gather_notes(pages, "anything")
|
||||||
|
|
||||||
|
# 4 first-round + 2 compression-round calls = 6 total.
|
||||||
|
assert ext_mock.await_count == 6
|
||||||
|
# Compressed from 4 notes to 2.
|
||||||
|
assert len(notes) == 2
|
||||||
|
# Pages from every original slice are preserved through compression.
|
||||||
|
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_compression_preserves_input_notes_when_a_group_fails(self, runtime: AppRuntime) -> None:
|
||||||
|
"""A compression chunk that raises has its input notes carried forward
|
||||||
|
rather than dropped, so page coverage isn't silently lost. The
|
||||||
|
succeeding chunk is replaced by its consolidated note.
|
||||||
|
|
||||||
|
Budget is sized so the post-round survivors (2 preserved + 1
|
||||||
|
consolidated) fit, leaving a single compression round as the
|
||||||
|
observable interaction."""
|
||||||
|
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=300)
|
||||||
|
pages = [_page(i, "x" * 150) for i in range(1, 5)]
|
||||||
|
|
||||||
|
call_count = 0
|
||||||
|
|
||||||
|
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
|
||||||
|
nonlocal call_count
|
||||||
|
call_count += 1
|
||||||
|
if call_count <= 4:
|
||||||
|
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
|
||||||
|
if call_count == 5:
|
||||||
|
# The first compression call (covering 2 of the round-1 notes) fails.
|
||||||
|
raise RuntimeError("compression group fails")
|
||||||
|
return _StubAgentResult(output=_ExtractedNotes(summary="ok"))
|
||||||
|
|
||||||
|
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
|
||||||
|
notes = await reasoner.gather_notes(pages, "anything")
|
||||||
|
|
||||||
|
# 2 preserved round-1 notes + 1 consolidated note = 3 notes total. Pages
|
||||||
|
# from every original slice are still covered (preservation worked).
|
||||||
|
assert len(notes) == 3
|
||||||
|
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||||
|
|
||||||
|
consolidated = [n for n in notes if n.summary == "ok"]
|
||||||
|
preserved = [n for n in notes if n.summary.startswith("x")]
|
||||||
|
assert len(consolidated) == 1
|
||||||
|
assert len(preserved) == 2
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_compression_bails_when_every_group_fails(self, runtime: AppRuntime) -> None:
|
||||||
|
"""If every chunk in a compression round fails, every input note is
|
||||||
|
preserved (none consolidated). The loop exits rather than retrying
|
||||||
|
the same shape forever."""
|
||||||
|
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=200, notes_char_budget=200)
|
||||||
|
pages = [_page(i, "x" * 150) for i in range(1, 5)]
|
||||||
|
|
||||||
|
call_count = 0
|
||||||
|
|
||||||
|
async def _stub(*_args: object, **_kwargs: object) -> _StubAgentResult[object]:
|
||||||
|
nonlocal call_count
|
||||||
|
call_count += 1
|
||||||
|
if call_count <= 4:
|
||||||
|
return _StubAgentResult(output=_ExtractedNotes(summary="x" * 60))
|
||||||
|
raise RuntimeError("compression always fails")
|
||||||
|
|
||||||
|
with patch.object(reasoner._extractor, "run", AsyncMock(side_effect=_stub)):
|
||||||
|
notes = await reasoner.gather_notes(pages, "anything")
|
||||||
|
|
||||||
|
# All 4 round-1 notes preserved through the bailed compression round.
|
||||||
|
assert len(notes) == 4
|
||||||
|
assert sorted({p for n in notes for p in n.pages}) == [1, 2, 3, 4]
|
||||||
|
|
||||||
|
|
||||||
|
class TestGroupNotesForCompression:
|
||||||
|
"""``_group_notes_for_compression`` is pure and packs by rendered char count."""
|
||||||
|
|
||||||
|
def test_packs_consecutive_notes_under_budget(self, runtime: AppRuntime) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=10_000)
|
||||||
|
notes = [ChunkNotes(pages=[i], summary=f"s-{i}") for i in range(1, 5)]
|
||||||
|
|
||||||
|
groups = reasoner._group_notes_for_compression(notes)
|
||||||
|
|
||||||
|
assert len(groups) == 1
|
||||||
|
assert [n.pages[0] for n in groups[0]] == [1, 2, 3, 4]
|
||||||
|
|
||||||
|
def test_starts_new_group_when_budget_exceeded(self, runtime: AppRuntime) -> None:
|
||||||
|
"""Each note already exceeds the per-group budget, so each becomes its
|
||||||
|
own group; this matches how slice-pages handles oversize pages."""
|
||||||
|
reasoner = ChunkedReasoner(runtime, chars_per_slice=5)
|
||||||
|
notes = [ChunkNotes(pages=[i], summary=f"slice-{i}-with-prose-to-fill-space") for i in range(1, 5)]
|
||||||
|
|
||||||
|
groups = reasoner._group_notes_for_compression(notes)
|
||||||
|
|
||||||
|
assert [[n.pages[0] for n in g] for g in groups] == [[1], [2], [3], [4]]
|
||||||
|
|
||||||
|
|
||||||
|
class TestExtractChunk:
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_pages_are_unioned_for_compression_chunks(self, runtime: AppRuntime) -> None:
|
||||||
|
"""A compression chunk's resulting note carries the union of input pages.
|
||||||
|
The model output schema doesn't include pages, so the wrapper is the
|
||||||
|
single source of truth."""
|
||||||
|
from stirling.agents.shared.chunked_reasoner import _ExtractedNotes
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime)
|
||||||
|
group = [
|
||||||
|
ChunkNotes(pages=[1, 2], summary="a"),
|
||||||
|
ChunkNotes(pages=[3], summary="b"),
|
||||||
|
ChunkNotes(pages=[4, 5], summary="c"),
|
||||||
|
]
|
||||||
|
chunk = reasoner._chunk_from_notes(group)
|
||||||
|
canned = _ExtractedNotes(summary="merged", facts=["x"], relevant_excerpts=["y"])
|
||||||
|
|
||||||
|
with patch.object(
|
||||||
|
reasoner._extractor,
|
||||||
|
"run",
|
||||||
|
AsyncMock(return_value=_StubAgentResult(output=canned)),
|
||||||
|
):
|
||||||
|
note, _ = await reasoner._extract_chunk(chunk, "anything")
|
||||||
|
|
||||||
|
assert note.pages == [1, 2, 3, 4, 5]
|
||||||
|
assert note.summary == "merged"
|
||||||
|
assert note.facts == ["x"]
|
||||||
|
assert note.relevant_excerpts == ["y"]
|
||||||
@@ -0,0 +1,220 @@
|
|||||||
|
"""Tests for ``WholeDocReaderCapability``: tool dispatch, multi-file iteration,
|
||||||
|
budget enforcement, and graceful handling of missing pages.
|
||||||
|
|
||||||
|
The map-phase LLM call is patched at the reasoner boundary so tests don't hit
|
||||||
|
any model.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import replace
|
||||||
|
from unittest.mock import AsyncMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability
|
||||||
|
from stirling.contracts import AiFile, PageText
|
||||||
|
from stirling.documents import Document, DocumentService, SqliteVecStore
|
||||||
|
from stirling.models import FileId
|
||||||
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
|
|
||||||
|
class StubEmbedder:
|
||||||
|
"""Deterministic embeddings so tests don't need a real provider."""
|
||||||
|
|
||||||
|
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,
|
||||||
|
text: str,
|
||||||
|
source: str = "",
|
||||||
|
base_metadata: dict[str, str] | None = None,
|
||||||
|
) -> list[Document]:
|
||||||
|
from stirling.documents.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
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def runtime_with_stub_docs(runtime: AppRuntime) -> AppRuntime:
|
||||||
|
stub = DocumentService(
|
||||||
|
embedder=StubEmbedder(), # type: ignore[arg-type]
|
||||||
|
store=SqliteVecStore.ephemeral(),
|
||||||
|
default_top_k=runtime.settings.rag_default_top_k,
|
||||||
|
)
|
||||||
|
return replace(runtime, documents=stub)
|
||||||
|
|
||||||
|
|
||||||
|
def _ai_file(file_id: str, name: str) -> AiFile:
|
||||||
|
return AiFile(id=FileId(file_id), name=name)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_full_document_returns_formatted_notes_for_single_file(
|
||||||
|
runtime_with_stub_docs: AppRuntime,
|
||||||
|
) -> None:
|
||||||
|
"""The tool reads the file's stored pages, calls the reasoner's map phase,
|
||||||
|
and returns the formatted notes prefixed by the file name."""
|
||||||
|
pages = [
|
||||||
|
PageText(page_number=1, text="Chapter one prose."),
|
||||||
|
PageText(page_number=2, text="Chapter two prose."),
|
||||||
|
]
|
||||||
|
await runtime_with_stub_docs.documents.ingest(FileId("doc-id"), pages, source="doc.pdf")
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||||
|
canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])]
|
||||||
|
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned_notes)) as gather_mock:
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("doc-id", "doc.pdf")],
|
||||||
|
reasoner=reasoner,
|
||||||
|
)
|
||||||
|
result = await capability._read_full_document("what is in the document?")
|
||||||
|
|
||||||
|
gather_mock.assert_awaited_once()
|
||||||
|
call = gather_mock.await_args
|
||||||
|
assert call is not None
|
||||||
|
pages_arg = call.args[0]
|
||||||
|
assert [p.page_number for p in pages_arg] == [1, 2]
|
||||||
|
assert "=== doc.pdf ===" in result
|
||||||
|
assert "fact-A" in result
|
||||||
|
assert "[Notes from pages 1-2]" in result
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs: AppRuntime) -> None:
|
||||||
|
"""Multi-file requests run the map phase per file and return one section
|
||||||
|
per file in the formatted output."""
|
||||||
|
for cid, source in (("doc-a", "a.pdf"), ("doc-b", "b.pdf")):
|
||||||
|
await runtime_with_stub_docs.documents.ingest(
|
||||||
|
FileId(cid),
|
||||||
|
[PageText(page_number=1, text=f"contents of {cid}")],
|
||||||
|
source=source,
|
||||||
|
)
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||||
|
notes_by_call = [
|
||||||
|
[ChunkNotes(pages=[1], summary="a-summary")],
|
||||||
|
[ChunkNotes(pages=[1], summary="b-summary")],
|
||||||
|
]
|
||||||
|
|
||||||
|
async def _gather(*_args: object, **_kwargs: object) -> list[ChunkNotes]:
|
||||||
|
return notes_by_call.pop(0)
|
||||||
|
|
||||||
|
with patch.object(reasoner, "gather_notes", AsyncMock(side_effect=_gather)) as gather_mock:
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")],
|
||||||
|
reasoner=reasoner,
|
||||||
|
)
|
||||||
|
result = await capability._read_full_document("compare them")
|
||||||
|
|
||||||
|
assert gather_mock.await_count == 2
|
||||||
|
assert "=== a.pdf ===" in result
|
||||||
|
assert "a-summary" in result
|
||||||
|
assert "=== b.pdf ===" in result
|
||||||
|
assert "b-summary" in result
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_full_document_skips_files_without_pages(runtime_with_stub_docs: AppRuntime) -> None:
|
||||||
|
"""Files with no stored pages are quietly skipped; the tool still runs
|
||||||
|
the map phase for files that do have pages."""
|
||||||
|
await runtime_with_stub_docs.documents.ingest(
|
||||||
|
FileId("present"),
|
||||||
|
[PageText(page_number=1, text="real content")],
|
||||||
|
source="present.pdf",
|
||||||
|
)
|
||||||
|
# 'missing' is never ingested -> read_pages returns [].
|
||||||
|
|
||||||
|
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||||
|
canned = [ChunkNotes(pages=[1], summary="present summary")]
|
||||||
|
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=canned)) as gather_mock:
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")],
|
||||||
|
reasoner=reasoner,
|
||||||
|
)
|
||||||
|
result = await capability._read_full_document("anything")
|
||||||
|
|
||||||
|
gather_mock.assert_awaited_once()
|
||||||
|
assert "=== present.pdf ===" in result
|
||||||
|
assert "missing.pdf" not in result
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_full_document_returns_empty_message_when_no_pages_anywhere(
|
||||||
|
runtime_with_stub_docs: AppRuntime,
|
||||||
|
) -> None:
|
||||||
|
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||||
|
with patch.object(reasoner, "gather_notes", AsyncMock()) as gather_mock:
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("nope", "nope.pdf")],
|
||||||
|
reasoner=reasoner,
|
||||||
|
)
|
||||||
|
result = await capability._read_full_document("anything")
|
||||||
|
|
||||||
|
gather_mock.assert_not_awaited()
|
||||||
|
assert result == "Could not read any document content."
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_full_document_budget_hides_tool_when_exhausted(
|
||||||
|
runtime_with_stub_docs: AppRuntime,
|
||||||
|
) -> None:
|
||||||
|
"""The prepare callback returns None once max_reads is reached so the
|
||||||
|
agent can no longer call the tool on subsequent turns. Mirrors
|
||||||
|
RagCapability's per-run budget."""
|
||||||
|
await runtime_with_stub_docs.documents.ingest(
|
||||||
|
FileId("doc-id"),
|
||||||
|
[PageText(page_number=1, text="content")],
|
||||||
|
source="doc.pdf",
|
||||||
|
)
|
||||||
|
reasoner = ChunkedReasoner(runtime_with_stub_docs)
|
||||||
|
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])):
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("doc-id", "doc.pdf")],
|
||||||
|
reasoner=reasoner,
|
||||||
|
max_reads=1,
|
||||||
|
)
|
||||||
|
sentinel: object = object()
|
||||||
|
|
||||||
|
# Budget intact -> prepare returns the tool.
|
||||||
|
assert await capability._prepare_read_full_document(None, sentinel) is sentinel # type: ignore[arg-type]
|
||||||
|
|
||||||
|
# Spend the budget.
|
||||||
|
await capability._read_full_document("anything")
|
||||||
|
|
||||||
|
# Budget spent -> prepare returns None.
|
||||||
|
assert await capability._prepare_read_full_document(None, sentinel) is None # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRuntime) -> None:
|
||||||
|
capability = WholeDocReaderCapability(
|
||||||
|
runtime=runtime_with_stub_docs,
|
||||||
|
files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")],
|
||||||
|
)
|
||||||
|
text = capability.instructions
|
||||||
|
|
||||||
|
assert "alpha.pdf" in text
|
||||||
|
assert "beta.pdf" in text
|
||||||
|
assert "read_full_document" in text
|
||||||
@@ -31,6 +31,10 @@ def build_app_settings() -> AppSettings:
|
|||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
rag_max_searches=5,
|
rag_max_searches=5,
|
||||||
|
chunked_reasoner_chars_per_slice=16_000,
|
||||||
|
chunked_reasoner_concurrency=10,
|
||||||
|
chunked_reasoner_notes_char_budget=250_000,
|
||||||
|
chunked_reasoner_worker_timeout_seconds=60.0,
|
||||||
max_pages=200,
|
max_pages=200,
|
||||||
max_characters=200_000,
|
max_characters=200_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
@@ -43,6 +43,10 @@ class StubSettingsProvider:
|
|||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
rag_max_searches=5,
|
rag_max_searches=5,
|
||||||
|
chunked_reasoner_chars_per_slice=16_000,
|
||||||
|
chunked_reasoner_concurrency=10,
|
||||||
|
chunked_reasoner_worker_timeout_seconds=60.0,
|
||||||
|
chunked_reasoner_notes_char_budget=250_000,
|
||||||
max_pages=100,
|
max_pages=100,
|
||||||
max_characters=100_000,
|
max_characters=100_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
@@ -2,14 +2,15 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from stirling.contracts import PageText
|
||||||
|
from stirling.documents.chunker import chunk_text
|
||||||
|
from stirling.documents.rag_capability import RagCapability
|
||||||
|
from stirling.documents.service import DocumentService
|
||||||
|
from stirling.documents.sqlite_vec_store import SqliteVecStore
|
||||||
|
from stirling.documents.store import Document, SearchResult
|
||||||
from stirling.models import FileId
|
from stirling.models import FileId
|
||||||
from stirling.rag.capability import RagCapability
|
|
||||||
from stirling.rag.chunker import chunk_text
|
|
||||||
from stirling.rag.service import RagService
|
|
||||||
from stirling.rag.sqlite_vec_store import SqliteVecStore
|
|
||||||
from stirling.rag.store import Document, SearchResult
|
|
||||||
|
|
||||||
# ── chunk_text ──────────────────────────────────────────────────────────
|
# chunk_text
|
||||||
|
|
||||||
|
|
||||||
class TestChunkText:
|
class TestChunkText:
|
||||||
@@ -26,7 +27,6 @@ class TestChunkText:
|
|||||||
def test_splits_on_paragraph_boundaries(self) -> None:
|
def test_splits_on_paragraph_boundaries(self) -> None:
|
||||||
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
|
text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph."
|
||||||
chunks = chunk_text(text, chunk_size=30, overlap=0)
|
chunks = chunk_text(text, chunk_size=30, overlap=0)
|
||||||
# Each paragraph fits in 30 chars, so they should be split
|
|
||||||
assert len(chunks) >= 2
|
assert len(chunks) >= 2
|
||||||
assert "First paragraph." in chunks[0]
|
assert "First paragraph." in chunks[0]
|
||||||
|
|
||||||
@@ -35,22 +35,19 @@ class TestChunkText:
|
|||||||
chunks = chunk_text(text, chunk_size=100, overlap=10)
|
chunks = chunk_text(text, chunk_size=100, overlap=10)
|
||||||
assert len(chunks) > 1
|
assert len(chunks) > 1
|
||||||
for chunk in chunks:
|
for chunk in chunks:
|
||||||
# Chunks may slightly exceed due to sentence boundary snapping
|
assert len(chunk) <= 200
|
||||||
assert len(chunk) <= 200 # generous upper bound
|
|
||||||
|
|
||||||
def test_overlap_produces_shared_content(self) -> None:
|
def test_overlap_produces_shared_content(self) -> None:
|
||||||
sentences = [f"Sentence number {i}." for i in range(20)]
|
sentences = [f"Sentence number {i}." for i in range(20)]
|
||||||
text = " ".join(sentences)
|
text = " ".join(sentences)
|
||||||
chunks = chunk_text(text, chunk_size=100, overlap=30)
|
chunks = chunk_text(text, chunk_size=100, overlap=30)
|
||||||
if len(chunks) >= 2:
|
if len(chunks) >= 2:
|
||||||
# After word-boundary snapping, the second chunk should share
|
words_in_first_tail = chunks[0].split()[-3:]
|
||||||
# some content with the tail of the first chunk
|
|
||||||
words_in_first_tail = chunks[0].split()[-3:] # last 3 words
|
|
||||||
overlap_text = " ".join(words_in_first_tail)
|
overlap_text = " ".join(words_in_first_tail)
|
||||||
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
|
assert overlap_text in chunks[1], f"Expected overlap '{overlap_text}' in chunk[1]: '{chunks[1][:80]}...'"
|
||||||
|
|
||||||
|
|
||||||
# ── SqliteVecStore ──────────────────────────────────────────────────────
|
# SqliteVecStore
|
||||||
|
|
||||||
|
|
||||||
class TestSqliteVecStore:
|
class TestSqliteVecStore:
|
||||||
@@ -59,12 +56,12 @@ class TestSqliteVecStore:
|
|||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_add_and_search(self) -> None:
|
async def test_add_and_search(self) -> None:
|
||||||
store = SqliteVecStore.ephemeral()
|
store = SqliteVecStore.ephemeral()
|
||||||
|
await store.ensure_collection("test-col", "test.pdf")
|
||||||
docs = [
|
docs = [
|
||||||
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
|
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
|
||||||
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
|
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
|
||||||
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
|
Document(id="3", text="The weather today is sunny", metadata={"source": "test"}),
|
||||||
]
|
]
|
||||||
# Simple 3-dimensional embeddings for testing
|
|
||||||
embeddings = [
|
embeddings = [
|
||||||
[1.0, 0.0, 0.0],
|
[1.0, 0.0, 0.0],
|
||||||
[0.9, 0.1, 0.0],
|
[0.9, 0.1, 0.0],
|
||||||
@@ -72,17 +69,16 @@ class TestSqliteVecStore:
|
|||||||
]
|
]
|
||||||
await store.add_documents("test-col", docs, embeddings)
|
await store.add_documents("test-col", docs, embeddings)
|
||||||
|
|
||||||
# Search with a query close to the programming-related docs
|
|
||||||
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
|
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
|
||||||
assert len(results) == 2
|
assert len(results) == 2
|
||||||
assert isinstance(results[0], SearchResult)
|
assert isinstance(results[0], SearchResult)
|
||||||
# The closest should be doc "1" (exact match on first dimension)
|
|
||||||
assert results[0].document.id == "1"
|
assert results[0].document.id == "1"
|
||||||
assert results[0].score > 0.5
|
assert results[0].score > 0.5
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_list_and_has_collection(self) -> None:
|
async def test_list_and_has_collection(self) -> None:
|
||||||
store = SqliteVecStore.ephemeral()
|
store = SqliteVecStore.ephemeral()
|
||||||
|
await store.ensure_collection("my-collection", "test.pdf")
|
||||||
docs = [Document(id="1", text="test", metadata={})]
|
docs = [Document(id="1", text="test", metadata={})]
|
||||||
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
|
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
|
||||||
|
|
||||||
@@ -94,6 +90,7 @@ class TestSqliteVecStore:
|
|||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_delete_collection(self) -> None:
|
async def test_delete_collection(self) -> None:
|
||||||
store = SqliteVecStore.ephemeral()
|
store = SqliteVecStore.ephemeral()
|
||||||
|
await store.ensure_collection("to-delete", "test.pdf")
|
||||||
docs = [Document(id="1", text="test", metadata={})]
|
docs = [Document(id="1", text="test", metadata={})]
|
||||||
await store.add_documents("to-delete", docs, [[1.0]])
|
await store.add_documents("to-delete", docs, [[1.0]])
|
||||||
|
|
||||||
@@ -104,6 +101,7 @@ class TestSqliteVecStore:
|
|||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_empty_collection(self) -> None:
|
async def test_search_empty_collection(self) -> None:
|
||||||
store = SqliteVecStore.ephemeral()
|
store = SqliteVecStore.ephemeral()
|
||||||
|
await store.ensure_collection("empty-test", "test.pdf")
|
||||||
docs = [Document(id="1", text="test", metadata={})]
|
docs = [Document(id="1", text="test", metadata={})]
|
||||||
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
|
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
|
||||||
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
|
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
|
||||||
@@ -117,7 +115,7 @@ class TestSqliteVecStore:
|
|||||||
await store.add_documents("bad", docs, [[1.0], [2.0]])
|
await store.add_documents("bad", docs, [[1.0], [2.0]])
|
||||||
|
|
||||||
|
|
||||||
# ── RagService (with stub embedder) ────────────────────────────────────
|
# DocumentService (with stub embedder)
|
||||||
|
|
||||||
|
|
||||||
class StubEmbeddingService:
|
class StubEmbeddingService:
|
||||||
@@ -127,7 +125,6 @@ class StubEmbeddingService:
|
|||||||
self._dim = dim
|
self._dim = dim
|
||||||
|
|
||||||
async def embed_query(self, text: str) -> list[float]:
|
async def embed_query(self, text: str) -> list[float]:
|
||||||
# Deterministic embedding based on hash of text
|
|
||||||
h = hash(text) % 1000
|
h = hash(text) % 1000
|
||||||
return [(h + i) / 1000.0 for i in range(self._dim)]
|
return [(h + i) / 1000.0 for i in range(self._dim)]
|
||||||
|
|
||||||
@@ -140,8 +137,6 @@ class StubEmbeddingService:
|
|||||||
source: str = "",
|
source: str = "",
|
||||||
base_metadata: dict[str, str] | None = None,
|
base_metadata: dict[str, str] | None = None,
|
||||||
) -> list[Document]:
|
) -> list[Document]:
|
||||||
from stirling.rag.chunker import chunk_text
|
|
||||||
|
|
||||||
chunks = chunk_text(text, 100, 10)
|
chunks = chunk_text(text, 100, 10)
|
||||||
docs = []
|
docs = []
|
||||||
for i, chunk in enumerate(chunks):
|
for i, chunk in enumerate(chunks):
|
||||||
@@ -154,53 +149,113 @@ class StubEmbeddingService:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def rag_service() -> RagService:
|
def documents() -> DocumentService:
|
||||||
"""Each RagService test gets its own fresh ephemeral store to avoid dimension conflicts."""
|
"""Each DocumentService test gets its own fresh ephemeral store to avoid dimension conflicts."""
|
||||||
store = SqliteVecStore.ephemeral()
|
store = SqliteVecStore.ephemeral()
|
||||||
return RagService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
|
return DocumentService(embedder=StubEmbeddingService(), store=store, default_top_k=3) # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
class TestRagService:
|
def _pages(text: str) -> list[PageText]:
|
||||||
|
return [PageText(page_number=1, text=text)]
|
||||||
|
|
||||||
|
|
||||||
|
class TestDocumentService:
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_index_and_search(self, rag_service: RagService) -> None:
|
async def test_ingest_and_search(self, documents: DocumentService) -> 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(FileId("docs"), text, source="guide.pdf")
|
count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf")
|
||||||
assert count > 0
|
assert count > 0
|
||||||
|
|
||||||
results = await rag_service.search("Python libraries", collection=FileId("docs"))
|
results = await documents.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
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_index_empty_text_returns_zero(self, rag_service: RagService) -> None:
|
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None:
|
||||||
count = await rag_service.index_text(FileId("docs"), "", source="empty.pdf")
|
count = await documents.ingest(FileId("docs"), _pages(""), 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, documents: DocumentService) -> None:
|
||||||
results = await rag_service.search("anything", collection=FileId("nonexistent"))
|
results = await documents.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, documents: DocumentService) -> None:
|
||||||
await rag_service.index_text(FileId("col-a"), "Machine learning overview.", source="ml.pdf")
|
await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf")
|
||||||
await rag_service.index_text(FileId("col-b"), "Deep learning with neural networks.", source="dl.pdf")
|
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf")
|
||||||
|
|
||||||
results = await rag_service.search("neural networks")
|
results = await documents.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, documents: DocumentService) -> None:
|
||||||
await rag_service.index_text(FileId("temp"), "Temporary data.", source="tmp.pdf")
|
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf")
|
||||||
collections = await rag_service.list_collections()
|
collections = await documents.list_collections()
|
||||||
assert "temp" in collections
|
assert "temp" in collections
|
||||||
|
|
||||||
await rag_service.delete_collection(FileId("temp"))
|
await documents.delete_collection(FileId("temp"))
|
||||||
collections = await rag_service.list_collections()
|
collections = await documents.list_collections()
|
||||||
assert "temp" not in collections
|
assert "temp" not in collections
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_ingest_stores_pages_in_order(self, documents: DocumentService) -> None:
|
||||||
|
pages = [
|
||||||
|
PageText(page_number=1, text="First page text."),
|
||||||
|
PageText(page_number=2, text="Second page text."),
|
||||||
|
PageText(page_number=3, text="Third page text."),
|
||||||
|
]
|
||||||
|
await documents.ingest(FileId("ordered"), pages, source="ordered.pdf")
|
||||||
|
|
||||||
# ── RagCapability ──────────────────────────────────────────────────────
|
stored = await documents.read_pages(FileId("ordered"))
|
||||||
|
assert [p.page_number for p in stored] == [1, 2, 3]
|
||||||
|
assert stored[0].text == "First page text."
|
||||||
|
assert stored[0].char_count == len("First page text.")
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_read_pages_with_range(self, documents: DocumentService) -> None:
|
||||||
|
from stirling.contracts import PageRange
|
||||||
|
|
||||||
|
pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)]
|
||||||
|
await documents.ingest(FileId("ranged"), pages, source="r.pdf")
|
||||||
|
|
||||||
|
subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4))
|
||||||
|
assert [p.page_number for p in subset] == [2, 3, 4]
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_ingest_replaces_previous_pages(self, documents: DocumentService) -> None:
|
||||||
|
await documents.ingest(
|
||||||
|
FileId("doc"),
|
||||||
|
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
|
||||||
|
source="v1.pdf",
|
||||||
|
)
|
||||||
|
await documents.ingest(
|
||||||
|
FileId("doc"),
|
||||||
|
[PageText(page_number=1, text="new")],
|
||||||
|
source="v2.pdf",
|
||||||
|
)
|
||||||
|
|
||||||
|
stored = await documents.read_pages(FileId("doc"))
|
||||||
|
assert [p.page_number for p in stored] == [1]
|
||||||
|
assert stored[0].text == "new"
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_ingest_keeps_blank_pages_in_page_store(self, documents: DocumentService) -> None:
|
||||||
|
"""Blank pages are skipped for embedding but retained in the page store
|
||||||
|
so page numbering stays continuous when reading back."""
|
||||||
|
pages = [
|
||||||
|
PageText(page_number=1, text="Real text on page 1."),
|
||||||
|
PageText(page_number=2, text=" "),
|
||||||
|
PageText(page_number=3, text="Real text on page 3."),
|
||||||
|
]
|
||||||
|
await documents.ingest(FileId("with-blanks"), pages, source="blanks.pdf")
|
||||||
|
|
||||||
|
stored = await documents.read_pages(FileId("with-blanks"))
|
||||||
|
assert [p.page_number for p in stored] == [1, 2, 3]
|
||||||
|
assert stored[1].text.strip() == ""
|
||||||
|
|
||||||
|
|
||||||
|
# RagCapability
|
||||||
|
|
||||||
|
|
||||||
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
|
async def _invoke_search_knowledge(capability: RagCapability, query: str, max_results: int = 5) -> str:
|
||||||
@@ -210,27 +265,27 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
|
|||||||
toolset = capability.toolset
|
toolset = capability.toolset
|
||||||
assert isinstance(toolset, FunctionToolset)
|
assert isinstance(toolset, FunctionToolset)
|
||||||
tool = toolset.tools["search_knowledge"]
|
tool = toolset.tools["search_knowledge"]
|
||||||
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg] — pyright can't infer the generic tool function's kwargs
|
return await tool.function(query=query, max_results=max_results) # type: ignore[call-arg]
|
||||||
|
|
||||||
|
|
||||||
class TestRagCapability:
|
class TestRagCapability:
|
||||||
def test_instructions_static_when_collections_pinned(self, rag_service: RagService) -> None:
|
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None:
|
||||||
cap = RagCapability(rag_service, collections=[FileId("docs"), FileId("manuals")])
|
cap = RagCapability(documents, 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
|
||||||
assert "search_knowledge" in instructions
|
assert "search_knowledge" in instructions
|
||||||
|
|
||||||
def test_instructions_dynamic_when_no_collections(self, rag_service: RagService) -> None:
|
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
instructions = cap.instructions
|
instructions = cap.instructions
|
||||||
assert callable(instructions)
|
assert callable(instructions)
|
||||||
|
|
||||||
@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, documents: DocumentService) -> None:
|
||||||
await rag_service.index_text(FileId("col-a"), "Alpha content.", source="a.pdf")
|
await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf")
|
||||||
await rag_service.index_text(FileId("col-b"), "Beta content.", source="b.pdf")
|
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf")
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
instructions_fn = cap.instructions
|
instructions_fn = cap.instructions
|
||||||
assert callable(instructions_fn)
|
assert callable(instructions_fn)
|
||||||
text = await instructions_fn()
|
text = await instructions_fn()
|
||||||
@@ -238,23 +293,23 @@ class TestRagCapability:
|
|||||||
assert "col-b" in text
|
assert "col-b" in text
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_dynamic_instructions_when_store_empty(self, rag_service: RagService) -> None:
|
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None:
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
instructions_fn = cap.instructions
|
instructions_fn = cap.instructions
|
||||||
assert callable(instructions_fn)
|
assert callable(instructions_fn)
|
||||||
text = await instructions_fn()
|
text = await instructions_fn()
|
||||||
assert "empty" in text.lower()
|
assert "empty" in text.lower()
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_search_knowledge_returns_no_results_message_when_empty(self, rag_service: RagService) -> None:
|
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None:
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
output = await _invoke_search_knowledge(cap, "anything")
|
output = await _invoke_search_knowledge(cap, "anything")
|
||||||
assert output == "No relevant results found in the knowledge base."
|
assert output == "No relevant results found in the knowledge base."
|
||||||
|
|
||||||
@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, documents: DocumentService) -> None:
|
||||||
await rag_service.index_text(FileId("docs"), "Python is a programming language.", source="guide.pdf")
|
await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf")
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
output = await _invoke_search_knowledge(cap, "Python")
|
output = await _invoke_search_knowledge(cap, "Python")
|
||||||
assert "[Result 1" in output
|
assert "[Result 1" in output
|
||||||
assert "source: guide.pdf" in output
|
assert "source: guide.pdf" in output
|
||||||
@@ -262,43 +317,39 @@ class TestRagCapability:
|
|||||||
assert "relevance:" in output
|
assert "relevance:" in output
|
||||||
|
|
||||||
@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, documents: DocumentService) -> None:
|
||||||
await rag_service.index_text(FileId("pinned"), "Pinned collection content.", source="pinned.pdf")
|
await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf")
|
||||||
await rag_service.index_text(FileId("other"), "Content in another collection.", source="other.pdf")
|
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf")
|
||||||
|
|
||||||
cap = RagCapability(rag_service, collections=[FileId("pinned")])
|
cap = RagCapability(documents, 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
|
||||||
|
|
||||||
@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, documents: DocumentService) -> 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(FileId("bulk"), paragraphs, source="bulk.pdf")
|
await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf")
|
||||||
|
|
||||||
cap = RagCapability(rag_service)
|
cap = RagCapability(documents)
|
||||||
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
|
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
|
||||||
# Only two results requested, so only Result 1 and Result 2 should appear
|
|
||||||
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
|
@pytest.mark.anyio
|
||||||
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, rag_service: RagService) -> None:
|
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None:
|
||||||
"""The prepare callback must return None once max_searches has been reached
|
"""The prepare callback must return None once max_searches has been reached
|
||||||
so the agent can no longer call the tool on subsequent turns."""
|
so the agent can no longer call the tool on subsequent turns."""
|
||||||
await rag_service.index_text(FileId("docs"), "Some content.", source="x.pdf")
|
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf")
|
||||||
cap = RagCapability(rag_service, max_searches=2)
|
cap = RagCapability(documents, max_searches=2)
|
||||||
tool_def = _dummy_tool_def()
|
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]
|
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")
|
||||||
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]
|
assert await cap._prepare_search_knowledge(None, tool_def) is None # type: ignore[arg-type]
|
||||||
|
|
||||||
|
|
||||||
@@ -6,9 +6,9 @@ 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_service
|
from stirling.api.dependencies import get_document_service
|
||||||
|
from stirling.documents import Document, DocumentService, SqliteVecStore
|
||||||
from stirling.models import FileId
|
from stirling.models import FileId
|
||||||
from stirling.rag import Document, RagService, SqliteVecStore
|
|
||||||
|
|
||||||
|
|
||||||
class StubEmbedder:
|
class StubEmbedder:
|
||||||
@@ -30,7 +30,7 @@ class StubEmbedder:
|
|||||||
source: str = "",
|
source: str = "",
|
||||||
base_metadata: dict[str, str] | None = None,
|
base_metadata: dict[str, str] | None = None,
|
||||||
) -> list[Document]:
|
) -> list[Document]:
|
||||||
from stirling.rag.chunker import chunk_text
|
from stirling.documents.chunker import chunk_text
|
||||||
|
|
||||||
chunks = chunk_text(text, 100, 10)
|
chunks = chunk_text(text, 100, 10)
|
||||||
docs = []
|
docs = []
|
||||||
@@ -43,8 +43,8 @@ class StubEmbedder:
|
|||||||
return docs
|
return docs
|
||||||
|
|
||||||
|
|
||||||
def _build_service() -> RagService:
|
def _build_service() -> DocumentService:
|
||||||
return RagService(
|
return DocumentService(
|
||||||
embedder=StubEmbedder(), # type: ignore[arg-type]
|
embedder=StubEmbedder(), # type: ignore[arg-type]
|
||||||
store=SqliteVecStore.ephemeral(),
|
store=SqliteVecStore.ephemeral(),
|
||||||
default_top_k=3,
|
default_top_k=3,
|
||||||
@@ -52,25 +52,25 @@ def _build_service() -> RagService:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def service() -> RagService:
|
def service() -> DocumentService:
|
||||||
return _build_service()
|
return _build_service()
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def client(service: RagService) -> Iterator[TestClient]:
|
def client(service: DocumentService) -> Iterator[TestClient]:
|
||||||
app.dependency_overrides[get_rag_service] = lambda: service
|
app.dependency_overrides[get_document_service] = lambda: service
|
||||||
try:
|
try:
|
||||||
yield TestClient(app)
|
yield TestClient(app)
|
||||||
finally:
|
finally:
|
||||||
app.dependency_overrides.pop(get_rag_service, None)
|
app.dependency_overrides.pop(get_document_service, None)
|
||||||
|
|
||||||
|
|
||||||
# ── POST /documents ─────────────────────────────────────────────────────
|
# ── POST /documents ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
def test_ingest_document_indexes_page_text(client: TestClient, service: RagService) -> None:
|
def test_ingest_document_indexes_page_text(client: TestClient, service: DocumentService) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "doc-123",
|
"documentId": "doc-123",
|
||||||
"source": "report.pdf",
|
"source": "report.pdf",
|
||||||
@@ -87,9 +87,9 @@ def test_ingest_document_indexes_page_text(client: TestClient, service: RagServi
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_ingest_document_replaces_existing_content(client: TestClient, service: RagService) -> None:
|
async def test_ingest_document_replaces_existing_content(client: TestClient, service: DocumentService) -> None:
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "replace-me",
|
"documentId": "replace-me",
|
||||||
"source": "replace-me.pdf",
|
"source": "replace-me.pdf",
|
||||||
@@ -98,7 +98,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
|
|||||||
)
|
)
|
||||||
# Second ingest with different content should replace the first entirely
|
# Second ingest with different content should replace the first entirely
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "replace-me",
|
"documentId": "replace-me",
|
||||||
"source": "replace-me.pdf",
|
"source": "replace-me.pdf",
|
||||||
@@ -115,7 +115,7 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
|
|||||||
|
|
||||||
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "mixed",
|
"documentId": "mixed",
|
||||||
"source": "mixed.pdf",
|
"source": "mixed.pdf",
|
||||||
@@ -130,14 +130,14 @@ def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
|
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"})
|
response = client.post("/api/v1/documents", json={"documentId": "empty", "source": "empty.pdf"})
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert response.json()["chunksIndexed"] == 0
|
assert response.json()["chunksIndexed"] == 0
|
||||||
|
|
||||||
|
|
||||||
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
@@ -145,7 +145,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
|||||||
|
|
||||||
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
@@ -153,7 +153,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
|||||||
|
|
||||||
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
||||||
)
|
)
|
||||||
assert response.status_code == 422
|
assert response.status_code == 422
|
||||||
@@ -161,7 +161,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
|||||||
|
|
||||||
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
|
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
|
||||||
response = client.post(
|
response = client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "bad-page",
|
"documentId": "bad-page",
|
||||||
"source": "bad-page.pdf",
|
"source": "bad-page.pdf",
|
||||||
@@ -176,30 +176,30 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) ->
|
|||||||
|
|
||||||
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
|
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={
|
json={
|
||||||
"documentId": "to-delete",
|
"documentId": "to-delete",
|
||||||
"source": "to-delete.pdf",
|
"source": "to-delete.pdf",
|
||||||
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
response = client.delete("/api/v1/rag/documents/to-delete")
|
response = client.delete("/api/v1/documents/to-delete")
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert response.json() == {"documentId": "to-delete", "deleted": True}
|
assert response.json() == {"documentId": "to-delete", "deleted": True}
|
||||||
|
|
||||||
|
|
||||||
def test_delete_document_is_idempotent(client: TestClient) -> None:
|
def test_delete_document_is_idempotent(client: TestClient) -> None:
|
||||||
response = client.delete("/api/v1/rag/documents/never-existed")
|
response = client.delete("/api/v1/documents/never-existed")
|
||||||
assert response.status_code == 200
|
assert response.status_code == 200
|
||||||
assert response.json() == {"documentId": "never-existed", "deleted": False}
|
assert response.json() == {"documentId": "never-existed", "deleted": False}
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_delete_document_removes_collection(client: TestClient, service: RagService) -> None:
|
async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
|
||||||
client.post(
|
client.post(
|
||||||
"/api/v1/rag/documents",
|
"/api/v1/documents",
|
||||||
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
|
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]},
|
||||||
)
|
)
|
||||||
assert await service.has_collection(FileId("gone"))
|
assert await service.has_collection(FileId("gone"))
|
||||||
client.delete("/api/v1/rag/documents/gone")
|
client.delete("/api/v1/documents/gone")
|
||||||
assert not await service.has_collection(FileId("gone"))
|
assert not await service.has_collection(FileId("gone"))
|
||||||
@@ -9,6 +9,7 @@ from stirling.contracts import (
|
|||||||
AiFile,
|
AiFile,
|
||||||
ExtractedFileText,
|
ExtractedFileText,
|
||||||
NeedIngestResponse,
|
NeedIngestResponse,
|
||||||
|
PageText,
|
||||||
PdfContentType,
|
PdfContentType,
|
||||||
PdfQuestionAnswerResponse,
|
PdfQuestionAnswerResponse,
|
||||||
PdfQuestionNotFoundResponse,
|
PdfQuestionNotFoundResponse,
|
||||||
@@ -17,8 +18,8 @@ from stirling.contracts import (
|
|||||||
PdfTextSelection,
|
PdfTextSelection,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
)
|
)
|
||||||
|
from stirling.documents import Document, DocumentService, SqliteVecStore
|
||||||
from stirling.models import FileId
|
from stirling.models import FileId
|
||||||
from stirling.rag import Document, RagService, SqliteVecStore
|
|
||||||
from stirling.services.runtime import AppRuntime
|
from stirling.services.runtime import AppRuntime
|
||||||
|
|
||||||
|
|
||||||
@@ -41,7 +42,7 @@ class StubEmbedder:
|
|||||||
source: str = "",
|
source: str = "",
|
||||||
base_metadata: dict[str, str] | None = None,
|
base_metadata: dict[str, str] | None = None,
|
||||||
) -> list[Document]:
|
) -> list[Document]:
|
||||||
from stirling.rag.chunker import chunk_text
|
from stirling.documents.chunker import chunk_text
|
||||||
|
|
||||||
chunks = chunk_text(text, 100, 10)
|
chunks = chunk_text(text, 100, 10)
|
||||||
docs: list[Document] = []
|
docs: list[Document] = []
|
||||||
@@ -65,13 +66,13 @@ class StubPdfQuestionAgent(PdfQuestionAgent):
|
|||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
|
def runtime_with_stub_rag(runtime: AppRuntime) -> AppRuntime:
|
||||||
"""A runtime whose RAG service uses a stub embedder + ephemeral store."""
|
"""A runtime whose document service uses a stub embedder + ephemeral store."""
|
||||||
stub = RagService(
|
stub = DocumentService(
|
||||||
embedder=StubEmbedder(), # type: ignore[arg-type]
|
embedder=StubEmbedder(), # type: ignore[arg-type]
|
||||||
store=SqliteVecStore.ephemeral(),
|
store=SqliteVecStore.ephemeral(),
|
||||||
default_top_k=runtime.settings.rag_default_top_k,
|
default_top_k=runtime.settings.rag_default_top_k,
|
||||||
)
|
)
|
||||||
return replace(runtime, rag_service=stub)
|
return replace(runtime, documents=stub)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
@@ -89,9 +90,9 @@ async def test_requests_ingest_when_file_missing_from_rag(runtime_with_stub_rag:
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
|
async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
await runtime_with_stub_rag.rag_service.index_text(
|
await runtime_with_stub_rag.documents.ingest(
|
||||||
collection=FileId("present-id"),
|
FileId("present-id"),
|
||||||
text="Invoice total: 120.00.",
|
[PageText(page_number=1, text="Invoice total: 120.00.")],
|
||||||
source="present.pdf",
|
source="present.pdf",
|
||||||
)
|
)
|
||||||
agent = PdfQuestionAgent(runtime_with_stub_rag)
|
agent = PdfQuestionAgent(runtime_with_stub_rag)
|
||||||
@@ -106,9 +107,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) ->
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub_rag: AppRuntime) -> None:
|
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(
|
await runtime_with_stub_rag.documents.ingest(
|
||||||
collection=FileId("invoice-id"),
|
FileId("invoice-id"),
|
||||||
text="Invoice total: 120.00.",
|
[PageText(page_number=1, text="Invoice total: 120.00.")],
|
||||||
source="invoice.pdf",
|
source="invoice.pdf",
|
||||||
)
|
)
|
||||||
agent = StubPdfQuestionAgent(
|
agent = StubPdfQuestionAgent(
|
||||||
@@ -137,9 +138,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub
|
|||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
|
async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: AppRuntime) -> None:
|
||||||
await runtime_with_stub_rag.rag_service.index_text(
|
await runtime_with_stub_rag.documents.ingest(
|
||||||
collection=FileId("shipping-id"),
|
FileId("shipping-id"),
|
||||||
text="This page contains only a shipping address.",
|
[PageText(page_number=1, text="This page contains only a shipping address.")],
|
||||||
source="shipping.pdf",
|
source="shipping.pdf",
|
||||||
)
|
)
|
||||||
agent = StubPdfQuestionAgent(
|
agent = StubPdfQuestionAgent(
|
||||||
|
|||||||
@@ -1,3 +1,7 @@
|
|||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
from conftest import build_app_settings
|
from conftest import build_app_settings
|
||||||
from fastapi.testclient import TestClient
|
from fastapi.testclient import TestClient
|
||||||
|
|
||||||
@@ -25,8 +29,11 @@ from stirling.contracts import (
|
|||||||
PdfQuestionNotFoundResponse,
|
PdfQuestionNotFoundResponse,
|
||||||
PdfQuestionRequest,
|
PdfQuestionRequest,
|
||||||
SupportedCapability,
|
SupportedCapability,
|
||||||
|
WholeDocReadStarted,
|
||||||
|
WholeDocSliceDone,
|
||||||
)
|
)
|
||||||
from stirling.models.tool_models import Angle, RotatePdfParams
|
from stirling.models.tool_models import Angle, RotatePdfParams
|
||||||
|
from stirling.services import emit_progress
|
||||||
|
|
||||||
|
|
||||||
class StubOrchestratorAgent:
|
class StubOrchestratorAgent:
|
||||||
@@ -40,6 +47,34 @@ class StubOrchestratorAgent:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class StubProgressOrchestratorAgent:
|
||||||
|
"""Orchestrator stub that emits two progress events before returning.
|
||||||
|
|
||||||
|
Used to verify the streaming endpoint plumbs the ContextVar emitter through
|
||||||
|
to deep callees and forwards events as NDJSON in order.
|
||||||
|
"""
|
||||||
|
|
||||||
|
async def handle(self, request: OrchestratorRequest) -> NeedContentResponse:
|
||||||
|
await emit_progress(WholeDocReadStarted(question="x", pages=10, slices=2))
|
||||||
|
await emit_progress(
|
||||||
|
WholeDocSliceDone(
|
||||||
|
completed=1,
|
||||||
|
total=2,
|
||||||
|
pages="pages=1-5",
|
||||||
|
duration_ms=42,
|
||||||
|
excerpts=2,
|
||||||
|
facts=3,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return NeedContentResponse(
|
||||||
|
resume_with=SupportedCapability.PDF_QUESTION,
|
||||||
|
reason=request.user_message,
|
||||||
|
files=[],
|
||||||
|
max_pages=1,
|
||||||
|
max_characters=1000,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class StubPdfEditAgent:
|
class StubPdfEditAgent:
|
||||||
async def handle(self, request: PdfEditRequest) -> EditCannotDoResponse:
|
async def handle(self, request: PdfEditRequest) -> EditCannotDoResponse:
|
||||||
return EditCannotDoResponse(reason=request.user_message)
|
return EditCannotDoResponse(reason=request.user_message)
|
||||||
@@ -87,14 +122,89 @@ def test_health_route() -> None:
|
|||||||
assert response.json()["status"] == "ok"
|
assert response.json()["status"] == "ok"
|
||||||
|
|
||||||
|
|
||||||
def test_orchestrator_route() -> None:
|
def test_orchestrator_route_streams_result_only_when_no_progress() -> None:
|
||||||
response = client.post(
|
"""The orchestrator endpoint always streams NDJSON. An agent that emits no
|
||||||
|
progress events still produces a single ``result`` frame with the typed
|
||||||
|
response body."""
|
||||||
|
with client.stream(
|
||||||
|
"POST",
|
||||||
"/api/v1/orchestrator",
|
"/api/v1/orchestrator",
|
||||||
json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
|
json={"userMessage": "route this", "files": [{"id": "test-id", "name": "test.pdf"}]},
|
||||||
)
|
) as response:
|
||||||
|
assert response.status_code == 200
|
||||||
|
events = [json.loads(line) for line in response.iter_lines() if line]
|
||||||
|
|
||||||
assert response.status_code == 200
|
assert [e["event"] for e in events] == ["result"]
|
||||||
assert response.json()["outcome"] == "need_content"
|
body = events[0]["response"]
|
||||||
|
assert body["outcome"] == "need_content"
|
||||||
|
|
||||||
|
|
||||||
|
def test_orchestrator_route_streams_progress_then_result() -> None:
|
||||||
|
"""When an agent emits progress via the ContextVar emitter, those frames
|
||||||
|
arrive on the wire before the final result frame."""
|
||||||
|
app.dependency_overrides[get_orchestrator_agent] = lambda: StubProgressOrchestratorAgent()
|
||||||
|
try:
|
||||||
|
with client.stream(
|
||||||
|
"POST",
|
||||||
|
"/api/v1/orchestrator",
|
||||||
|
json={"userMessage": "stream this", "files": [{"id": "test-id", "name": "test.pdf"}]},
|
||||||
|
) as response:
|
||||||
|
assert response.status_code == 200
|
||||||
|
events = [json.loads(line) for line in response.iter_lines() if line]
|
||||||
|
finally:
|
||||||
|
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
|
||||||
|
|
||||||
|
progress = [e for e in events if e["event"] == "progress"]
|
||||||
|
results = [e for e in events if e["event"] == "result"]
|
||||||
|
assert [p["phase"] for p in progress] == ["whole_doc_read_started", "whole_doc_slice_done"]
|
||||||
|
assert progress[1]["completed"] == 1
|
||||||
|
assert len(results) == 1
|
||||||
|
response = results[0]["response"]
|
||||||
|
assert response["outcome"] == "need_content"
|
||||||
|
# Wire format must be camelCase: Java's Jackson deserializer expects camelCase
|
||||||
|
# field names. ``maxPages`` here doubles as a regression guard against the
|
||||||
|
# snake_case bug that surfaced as "need_ingest without listing any files to ingest".
|
||||||
|
assert "maxPages" in response
|
||||||
|
assert "max_pages" not in response
|
||||||
|
|
||||||
|
|
||||||
|
def test_orchestrator_route_emits_heartbeats_while_agent_is_busy() -> None:
|
||||||
|
"""While the agent is in flight, the streaming endpoint emits heartbeat
|
||||||
|
frames at the configured cadence so each layer of the connection stays
|
||||||
|
visibly alive and disconnects propagate within bounded latency."""
|
||||||
|
|
||||||
|
class _SlowAgent:
|
||||||
|
async def handle(self, _request: OrchestratorRequest) -> NeedContentResponse:
|
||||||
|
# Sleep long enough for several heartbeats at the patched cadence.
|
||||||
|
await asyncio.sleep(0.2)
|
||||||
|
return NeedContentResponse(
|
||||||
|
resume_with=SupportedCapability.PDF_QUESTION,
|
||||||
|
reason="ok",
|
||||||
|
files=[],
|
||||||
|
max_pages=1,
|
||||||
|
max_characters=1000,
|
||||||
|
)
|
||||||
|
|
||||||
|
app.dependency_overrides[get_orchestrator_agent] = lambda: _SlowAgent()
|
||||||
|
try:
|
||||||
|
with patch("stirling.api.routes.orchestrator.HEARTBEAT_INTERVAL_SECONDS", 0.03):
|
||||||
|
with client.stream(
|
||||||
|
"POST",
|
||||||
|
"/api/v1/orchestrator",
|
||||||
|
json={"userMessage": "wait", "files": [{"id": "test-id", "name": "test.pdf"}]},
|
||||||
|
) as response:
|
||||||
|
assert response.status_code == 200
|
||||||
|
events = [json.loads(line) for line in response.iter_lines() if line]
|
||||||
|
finally:
|
||||||
|
app.dependency_overrides[get_orchestrator_agent] = lambda: StubOrchestratorAgent()
|
||||||
|
|
||||||
|
heartbeats = [e for e in events if e["event"] == "heartbeat"]
|
||||||
|
results = [e for e in events if e["event"] == "result"]
|
||||||
|
# At least a couple of heartbeats fired during the 0.2s agent sleep at 0.03s cadence.
|
||||||
|
assert len(heartbeats) >= 2
|
||||||
|
# The result still arrives after the agent finishes.
|
||||||
|
assert len(results) == 1
|
||||||
|
assert results[0]["response"]["outcome"] == "need_content"
|
||||||
|
|
||||||
|
|
||||||
def test_pdf_edit_route() -> None:
|
def test_pdf_edit_route() -> None:
|
||||||
|
|||||||
@@ -92,6 +92,10 @@ def test_app_settings_accepts_model_configuration() -> None:
|
|||||||
rag_chunk_overlap=64,
|
rag_chunk_overlap=64,
|
||||||
rag_default_top_k=5,
|
rag_default_top_k=5,
|
||||||
rag_max_searches=5,
|
rag_max_searches=5,
|
||||||
|
chunked_reasoner_chars_per_slice=16_000,
|
||||||
|
chunked_reasoner_concurrency=10,
|
||||||
|
chunked_reasoner_worker_timeout_seconds=60.0,
|
||||||
|
chunked_reasoner_notes_char_budget=250_000,
|
||||||
max_pages=200,
|
max_pages=200,
|
||||||
max_characters=200_000,
|
max_characters=200_000,
|
||||||
posthog_enabled=False,
|
posthog_enabled=False,
|
||||||
|
|||||||
@@ -2682,6 +2682,10 @@ executing_tool_step = "Running {{tool}} (step {{step}} of {{total}})..."
|
|||||||
extracting_content = "Extracting content from your documents..."
|
extracting_content = "Extracting content from your documents..."
|
||||||
processing = "Processing extracted content..."
|
processing = "Processing extracted content..."
|
||||||
thinking = "Thinking..."
|
thinking = "Thinking..."
|
||||||
|
whole_doc_compression_round = "Consolidating notes..."
|
||||||
|
whole_doc_read_done = "Finished reading the document..."
|
||||||
|
whole_doc_read_started = "Reading the document..."
|
||||||
|
whole_doc_slice_done = "Reading the document... ({{percent}}% complete)"
|
||||||
|
|
||||||
[chat.toolsUsed]
|
[chat.toolsUsed]
|
||||||
summary = "Ran {{count}} tools"
|
summary = "Ran {{count}} tools"
|
||||||
|
|||||||
@@ -37,6 +37,79 @@ export enum AiWorkflowPhase {
|
|||||||
EXTRACTING_CONTENT = "extracting_content",
|
EXTRACTING_CONTENT = "extracting_content",
|
||||||
EXECUTING_TOOL = "executing_tool",
|
EXECUTING_TOOL = "executing_tool",
|
||||||
PROCESSING = "processing",
|
PROCESSING = "processing",
|
||||||
|
ENGINE_PROGRESS = "engine_progress",
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Engine-side progress detail for ENGINE_PROGRESS events. Mirrors the Python
|
||||||
|
* {@code ProgressEvent} discriminated union (engine/src/stirling/contracts/progress.py)
|
||||||
|
* and the Java {@code AiEngineProgressDetail} sealed interface; the {@code phase}
|
||||||
|
* string is the discriminator. Field names are camelCase because the engine
|
||||||
|
* serialises by alias.
|
||||||
|
*/
|
||||||
|
export interface WholeDocReadStartedDetail {
|
||||||
|
phase: "whole_doc_read_started";
|
||||||
|
question: string;
|
||||||
|
pages: number;
|
||||||
|
slices: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WholeDocSliceDoneDetail {
|
||||||
|
phase: "whole_doc_slice_done";
|
||||||
|
completed: number;
|
||||||
|
total: number;
|
||||||
|
/** Page-range label, e.g. "pages=1-5". */
|
||||||
|
pages: string;
|
||||||
|
durationMs: number;
|
||||||
|
excerpts: number;
|
||||||
|
facts: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WholeDocCompressionRoundDetail {
|
||||||
|
phase: "whole_doc_compression_round";
|
||||||
|
roundNumber: number;
|
||||||
|
notesIn: number;
|
||||||
|
groups: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface WholeDocReadDoneDetail {
|
||||||
|
phase: "whole_doc_read_done";
|
||||||
|
completed: number;
|
||||||
|
slices: number;
|
||||||
|
durationSeconds: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export type EngineProgressDetail =
|
||||||
|
| WholeDocReadStartedDetail
|
||||||
|
| WholeDocSliceDoneDetail
|
||||||
|
| WholeDocCompressionRoundDetail
|
||||||
|
| WholeDocReadDoneDetail;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* What we actually carry across the wire boundary: a known typed variant, or a
|
||||||
|
* forward-compat shape with just the discriminator string. The "unknown" arm
|
||||||
|
* exists so a new engine-side phase rolling out before a frontend update keeps
|
||||||
|
* rendering the generic processing message instead of crashing the union.
|
||||||
|
*/
|
||||||
|
export interface UnknownEngineProgressDetail {
|
||||||
|
phase: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export type AnyEngineProgressDetail =
|
||||||
|
| EngineProgressDetail
|
||||||
|
| UnknownEngineProgressDetail;
|
||||||
|
|
||||||
|
const KNOWN_ENGINE_PHASES = new Set<string>([
|
||||||
|
"whole_doc_read_started",
|
||||||
|
"whole_doc_slice_done",
|
||||||
|
"whole_doc_compression_round",
|
||||||
|
"whole_doc_read_done",
|
||||||
|
]);
|
||||||
|
|
||||||
|
export function isKnownEngineProgressDetail(
|
||||||
|
detail: AnyEngineProgressDetail,
|
||||||
|
): detail is EngineProgressDetail {
|
||||||
|
return KNOWN_ENGINE_PHASES.has(detail.phase);
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface AiWorkflowProgress {
|
export interface AiWorkflowProgress {
|
||||||
@@ -47,6 +120,11 @@ export interface AiWorkflowProgress {
|
|||||||
stepIndex?: number;
|
stepIndex?: number;
|
||||||
/** Total number of plan steps, for EXECUTING_TOOL events. */
|
/** Total number of plan steps, for EXECUTING_TOOL events. */
|
||||||
stepCount?: number;
|
stepCount?: number;
|
||||||
|
/**
|
||||||
|
* Engine-side event payload, for ENGINE_PROGRESS events. Typed sub-phase
|
||||||
|
* record (e.g. {@link WholeDocSliceDoneDetail}) the UI can render in detail.
|
||||||
|
*/
|
||||||
|
engineDetail?: AnyEngineProgressDetail;
|
||||||
}
|
}
|
||||||
|
|
||||||
type AiWorkflowOutcome =
|
type AiWorkflowOutcome =
|
||||||
@@ -159,6 +237,7 @@ interface ProgressEvent {
|
|||||||
tool?: string;
|
tool?: string;
|
||||||
stepIndex?: number;
|
stepIndex?: number;
|
||||||
stepCount?: number;
|
stepCount?: number;
|
||||||
|
engineDetail?: AnyEngineProgressDetail;
|
||||||
}
|
}
|
||||||
|
|
||||||
async function consumeSSEStream(
|
async function consumeSSEStream(
|
||||||
@@ -378,6 +457,7 @@ export function ChatProvider({ children }: { children: ReactNode }) {
|
|||||||
tool: data.tool,
|
tool: data.tool,
|
||||||
stepIndex: data.stepIndex,
|
stepIndex: data.stepIndex,
|
||||||
stepCount: data.stepCount,
|
stepCount: data.stepCount,
|
||||||
|
engineDetail: data.engineDetail,
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -29,7 +29,9 @@ import ExpandLessIcon from "@mui/icons-material/ExpandLess";
|
|||||||
import {
|
import {
|
||||||
useChat,
|
useChat,
|
||||||
AiWorkflowPhase,
|
AiWorkflowPhase,
|
||||||
|
isKnownEngineProgressDetail,
|
||||||
type AiWorkflowProgress,
|
type AiWorkflowProgress,
|
||||||
|
type AnyEngineProgressDetail,
|
||||||
} from "@app/components/chat/ChatContext";
|
} from "@app/components/chat/ChatContext";
|
||||||
import { useTranslatedToolCatalog } from "@app/data/useTranslatedToolRegistry";
|
import { useTranslatedToolCatalog } from "@app/data/useTranslatedToolRegistry";
|
||||||
import "@app/components/chat/ChatPanel.css";
|
import "@app/components/chat/ChatPanel.css";
|
||||||
@@ -90,9 +92,41 @@ function formatProgress(
|
|||||||
})
|
})
|
||||||
: t("chat.progress.executing_tool_generic");
|
: t("chat.progress.executing_tool_generic");
|
||||||
}
|
}
|
||||||
|
if (progress.phase === AiWorkflowPhase.ENGINE_PROGRESS) {
|
||||||
|
return formatEngineProgress(progress.engineDetail, t);
|
||||||
|
}
|
||||||
return t(`chat.progress.${progress.phase}`);
|
return t(`chat.progress.${progress.phase}`);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Render an engine-side progress event (e.g. chunked-reasoner slice progress) into a user-facing
|
||||||
|
* message. Falls through to the generic processing label for unknown sub-phases so adding new
|
||||||
|
* engine events doesn't break the UI before the frontend learns about them.
|
||||||
|
*/
|
||||||
|
function formatEngineProgress(
|
||||||
|
detail: AnyEngineProgressDetail | undefined,
|
||||||
|
t: TranslateFn,
|
||||||
|
): string {
|
||||||
|
if (!detail || !isKnownEngineProgressDetail(detail)) {
|
||||||
|
return t("chat.progress.processing");
|
||||||
|
}
|
||||||
|
switch (detail.phase) {
|
||||||
|
case "whole_doc_read_started":
|
||||||
|
return t("chat.progress.whole_doc_read_started");
|
||||||
|
case "whole_doc_slice_done": {
|
||||||
|
const percent =
|
||||||
|
detail.total > 0
|
||||||
|
? Math.round((detail.completed / detail.total) * 100)
|
||||||
|
: 0;
|
||||||
|
return t("chat.progress.whole_doc_slice_done", { percent });
|
||||||
|
}
|
||||||
|
case "whole_doc_compression_round":
|
||||||
|
return t("chat.progress.whole_doc_compression_round");
|
||||||
|
case "whole_doc_read_done":
|
||||||
|
return t("chat.progress.whole_doc_read_done");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
function ToolsUsedBlock({
|
function ToolsUsedBlock({
|
||||||
tools,
|
tools,
|
||||||
resolveToolName,
|
resolveToolName,
|
||||||
|
|||||||
Reference in New Issue
Block a user