arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-14 10:34:06 +02:00
Code Issues Actions 4 Packages Projects Releases 1 Wiki Activity

Set up document management for Stirling Engine (#6476)

Browse Source 
# Description of Changes
Change Stirling Engine to support deleting documents automatically. This
happens both on user logout and after an amount of time specified by the
Java when ingesting a document (allowing for personal documents to have
short lifetimes but org documents to be left in the db with no expiry
date). Also sets up an [ACL
policy](https://en.wikipedia.org/wiki/Access-control_list) for the
documents so the database knows which users have access to which
documents. This is not fully implemented in the Java, so currently all
docs are treated as having a single owner, the uploader, but
theoretically when we need to support org storage, we shouldn't need to
change the db schema.
This commit is contained in:
James Brunton
2026-06-03 11:52:11 +00:00
committed by GitHub
parent 71633861d0
commit 1264f4cfed
48 changed files with 2039 additions and 411 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/proprietary/src/main/java/stirling/software/proprietary/controller/api/AiEngineController.java
+12 -3
View File
@@ -5,6 +5,7 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.Executor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpStatus;
@@ -30,6 +31,7 @@ import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.job.ResultFile;
import stirling.software.common.service.JobOwnershipService;
import stirling.software.common.service.TaskManager;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
@@ -58,6 +60,7 @@ public class AiEngineController {
private final TaskManager taskManager;
private final JobOwnershipService jobOwnershipService;
private final AiEngineEndpointResolver endpointResolver;
private final UserServiceInterface userService;
/**
* SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a
@@ -74,7 +77,8 @@ public class AiEngineController {
@Qualifier("aiStreamExecutor") Executor aiStreamExecutor,
TaskManager taskManager,
JobOwnershipService jobOwnershipService,
AiEngineEndpointResolver endpointResolver) {
AiEngineEndpointResolver endpointResolver,
@Autowired(required = false) UserServiceInterface userService) {
this.aiEngineClient = aiEngineClient;
this.aiWorkflowService = aiWorkflowService;
this.objectMapper = objectMapper;
@@ -82,6 +86,11 @@ public class AiEngineController {
this.taskManager = taskManager;
this.jobOwnershipService = jobOwnershipService;
this.endpointResolver = endpointResolver;
this.userService = userService;
}
private String currentUserId() {
return userService != null ? userService.getCurrentUsername() : null;
}
@GetMapping("/health")
@@ -89,7 +98,7 @@ public class AiEngineController {
summary = "AI engine health check",
description = "Returns the health status of the AI engine including configured models")
public ResponseEntity<String> health() throws IOException {
String response = aiEngineClient.get("/health");
String response = aiEngineClient.get("/health", currentUserId());
return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response);
}
@@ -243,7 +252,7 @@ public class AiEngineController {
HttpStatus.BAD_REQUEST, "Request body must be a JSON object");
}
String forwardedBody = withEnabledEndpoints((ObjectNode) parsed);
String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody);
String response = aiEngineClient.post("/api/v1/pdf/edit", forwardedBody, currentUserId());
return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response);
}
app/proprietary/src/main/java/stirling/software/proprietary/model/api/ai/AiDocumentIngestRequest.java
+13
View File
@@ -1,5 +1,6 @@
package stirling.software.proprietary.model.api.ai;
import java.time.Instant;
import java.util.List;
import lombok.AllArgsConstructor;
@@ -10,6 +11,12 @@ 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.
*
* <p>{@code ownerId} is the tenant the doc belongs to (a user for personal uploads, an org for
* shared content). {@code readPrincipals} is the explicit list of principals granted read access.
* {@code expiresAt} is when the engine's reaper should delete this doc; {@code null} means
* "persistent until explicit delete" (used for org-shared content). Java picks the value per doc;
* the engine does not default it.
*/
@Data
@NoArgsConstructor
@@ -21,4 +28,10 @@ public class AiDocumentIngestRequest {
private String source;
private List<AiPageText> pageText;
private String ownerId;
private List<String> readPrincipals;
private Instant expiresAt;
}
app/proprietary/src/main/java/stirling/software/proprietary/security/CustomLogoutSuccessHandler.java
+33
View File
@@ -8,6 +8,8 @@ import java.util.List;
import java.util.Locale;
import org.springframework.core.io.Resource;
import org.springframework.security.authentication.AuthenticationTrustResolver;
import org.springframework.security.authentication.AuthenticationTrustResolverImpl;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
@@ -36,6 +38,7 @@ import stirling.software.proprietary.audit.Audited;
import stirling.software.proprietary.security.saml2.CertificateUtils;
import stirling.software.proprietary.security.saml2.CustomSaml2AuthenticatedPrincipal;
import stirling.software.proprietary.security.service.JwtServiceInterface;
import stirling.software.proprietary.service.AiUserDataService;
@Slf4j
@RequiredArgsConstructor
@@ -49,12 +52,22 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
private final JwtServiceInterface jwtService;
private final AiUserDataService aiUserDataService;
private static final AuthenticationTrustResolver TRUST_RESOLVER =
new AuthenticationTrustResolverImpl();
@Override
@Audited(type = AuditEventType.USER_LOGOUT, level = AuditLevel.BASIC)
public void onLogoutSuccess(
HttpServletRequest request, HttpServletResponse response, Authentication authentication)
throws IOException {
String username = resolveUsername(request, authentication);
if (username != null) {
aiUserDataService.purgeUserDocuments(username);
}
if (!response.isCommitted()) {
if (authentication != null) {
if (authentication instanceof Saml2Authentication samlAuthentication) {
@@ -88,6 +101,26 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
}
}
/**
* Pick the right name to purge under. JWT cookie wins if present and parseable; we fall through
* to whatever Spring handed us only when there's no cookie. Spring's anonymous principal is
* filtered out via {@link AuthenticationTrustResolver} so we don't purge under that
* pseudo-user.
*/
private String resolveUsername(HttpServletRequest request, Authentication authentication) {
if (jwtService != null) {
String fromCookie = jwtService.extractUsernameFromRequestAllowExpired(request);
if (fromCookie != null) {
return fromCookie;
}
}
if (authentication == null || TRUST_RESOLVER.isAnonymous(authentication)) {
return null;
}
String name = authentication.getName();
return (name != null && !name.isBlank()) ? name : null;
}
// Redirect for SAML2 authentication logout
private void getRedirect_saml2(
HttpServletRequest request,
app/proprietary/src/main/java/stirling/software/proprietary/security/configuration/SecurityConfiguration.java
+8 -2
View File
@@ -93,6 +93,7 @@ public class SecurityConfiguration {
licenseSettingsService;
private final ClientRegistrationRepository clientRegistrationRepository;
private final PasswordEncoder passwordEncoder;
private final stirling.software.proprietary.service.AiUserDataService aiUserDataService;
public SecurityConfiguration(
PersistentLoginRepository persistentLoginRepository,
@@ -115,7 +116,8 @@ public class SecurityConfiguration {
OpenSaml5AuthenticationRequestResolver saml2AuthenticationRequestResolver,
@Autowired(required = false) ClientRegistrationRepository clientRegistrationRepository,
stirling.software.proprietary.service.UserLicenseSettingsService licenseSettingsService,
PasswordEncoder passwordEncoder) {
PasswordEncoder passwordEncoder,
stirling.software.proprietary.service.AiUserDataService aiUserDataService) {
this.userDetailsService = userDetailsService;
this.userService = userService;
this.loginEnabledValue = loginEnabledValue;
@@ -135,6 +137,7 @@ public class SecurityConfiguration {
this.clientRegistrationRepository = clientRegistrationRepository;
this.licenseSettingsService = licenseSettingsService;
this.passwordEncoder = passwordEncoder;
this.aiUserDataService = aiUserDataService;
}
/**
@@ -322,7 +325,10 @@ public class SecurityConfiguration {
.matcher("/logout"))
.logoutSuccessHandler(
new CustomLogoutSuccessHandler(
securityProperties, appConfig, jwtService))
securityProperties,
appConfig,
jwtService,
aiUserDataService))
.clearAuthentication(true)
.invalidateHttpSession(true)
.deleteCookies("JSESSIONID", "remember-me", "stirling_jwt"));
app/proprietary/src/main/java/stirling/software/proprietary/security/controller/api/AuthController.java
+6 -2
View File
@@ -44,6 +44,7 @@ import stirling.software.proprietary.security.service.RefreshRateLimitService;
import stirling.software.proprietary.security.service.TotpService;
import stirling.software.proprietary.security.service.UserService;
import stirling.software.proprietary.security.util.DesktopClientUtils;
import stirling.software.proprietary.service.AiUserDataService;
/** REST API Controller for authentication operations. */
@RestController
@@ -62,6 +63,7 @@ public class AuthController {
private final RefreshRateLimitService refreshRateLimitService;
private final ApplicationProperties.Security securityProperties;
private final ApplicationProperties applicationProperties;
private final AiUserDataService aiUserDataService;
/**
* Login endpoint - replaces Supabase signInWithPassword
@@ -281,11 +283,13 @@ public class AuthController {
*/
@PreAuthorize("!hasAuthority('ROLE_DEMO_USER')")
@PostMapping("/logout")
public ResponseEntity<?> logout(HttpServletResponse response) {
public ResponseEntity<?> logout(HttpServletRequest request, HttpServletResponse response) {
try {
String username = jwtService.extractUsernameFromRequestAllowExpired(request);
SecurityContextHolder.clearContext();
aiUserDataService.purgeUserDocuments(username);
log.debug("User logged out successfully");
log.debug("User logged out successfully (username={})", username);
return ResponseEntity.ok(Map.of("message", "Logged out successfully"));
app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtService.java
+15
View File
@@ -349,6 +349,21 @@ public class JwtService implements JwtServiceInterface {
return null;
}
@Override
public String extractUsernameFromRequestAllowExpired(HttpServletRequest request) {
try {
String token = extractToken(request);
if (token == null || token.isBlank()) {
return null;
}
String username = extractUsernameAllowExpired(token);
return (username != null && !username.isBlank()) ? username : null;
} catch (Exception e) {
log.debug("Could not extract username from request JWT: {}", e.getMessage());
return null;
}
}
@Override
public boolean isJwtEnabled() {
return v2Enabled;
app/proprietary/src/main/java/stirling/software/proprietary/security/service/JwtServiceInterface.java
+13
View File
@@ -93,6 +93,19 @@ public interface JwtServiceInterface {
*/
String extractToken(HttpServletRequest request);
/**
* Read the username off the request's JWT, allowing an expired token. Returns null when no
* token is present, the token can't be parsed, or the resulting username is blank.
*
* <p>Used by flows that need to identify the user without depending on {@code
* SecurityContextHolder} - for example logout, where the security filter chain may have left
* the anonymous principal in place by the time the handler runs.
*
* @param request HTTP servlet request
* @return username from the token, or null when one can't be safely derived
*/
String extractUsernameFromRequestAllowExpired(HttpServletRequest request);
/**
* Check if JWT authentication is enabled
*
app/proprietary/src/main/java/stirling/software/proprietary/service/AiEngineClient.java
+61 -20
View File
@@ -43,22 +43,23 @@ public class AiEngineClient {
this.httpClient = httpClient;
}
public String post(String path, String jsonBody) throws IOException {
public String post(String path, String jsonBody, String userId) throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
return postWithTimeout(path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds()));
return postWithTimeout(
path, jsonBody, Duration.ofSeconds(config.getTimeoutSeconds()), userId);
}
/**
* POST with an explicit per-call timeout, for heavy operations (e.g. RAG ingestion of a large
* document) that legitimately take longer than the default timeout.
*/
public String postLongRunning(String path, String jsonBody) throws IOException {
public String postLongRunning(String path, String jsonBody, String userId) throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
return postWithTimeout(
path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds()));
path, jsonBody, Duration.ofSeconds(config.getLongRunningTimeoutSeconds()), userId);
}
private String postWithTimeout(String path, String jsonBody, Duration timeout)
private String postWithTimeout(String path, String jsonBody, Duration timeout, String userId)
throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
@@ -69,22 +70,32 @@ public class AiEngineClient {
String url = config.getUrl().stripTrailing() + path;
log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds());
HttpRequest request =
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Content-Type", "application/json")
.header("Accept", "application/json")
.timeout(timeout)
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
.build();
HttpResponse<String> response = sendRequest(request);
.POST(HttpRequest.BodyPublishers.ofString(jsonBody));
addUserHeader(builder, userId);
HttpResponse<String> response = sendRequest(builder.build());
log.debug("AI engine responded with status {}", response.statusCode());
checkResponseStatus(response);
return response.body();
}
/**
* Attach the X-User-Id header so the engine can scope per-user storage (RAG documents, search
* results) to the caller. Skipped when {@code userId} is blank: the engine treats the request
* as anonymous and refuses any route that requires tenancy.
*/
private static void addUserHeader(HttpRequest.Builder builder, String userId) {
if (userId != null && !userId.isBlank()) {
builder.header("X-User-Id", userId);
}
}
/**
* 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.
@@ -94,7 +105,8 @@ public class AiEngineClient {
* 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)
public void streamPost(
String path, String jsonBody, String userId, Consumer<String> lineConsumer)
throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
@@ -109,14 +121,15 @@ public class AiEngineClient {
url,
timeout.toSeconds());
HttpRequest request =
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Content-Type", "application/json")
.header("Accept", "application/x-ndjson")
.timeout(timeout)
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
.build();
.POST(HttpRequest.BodyPublishers.ofString(jsonBody));
addUserHeader(builder, userId);
HttpRequest request = builder.build();
HttpResponse<Stream<String>> response;
try {
@@ -149,7 +162,36 @@ public class AiEngineClient {
}
}
public String get(String path) throws IOException {
/**
* DELETE with no body. Used for purging the caller's RAG content on logout. Wraps the same
* error envelope as {@link #post} / {@link #get} so callers see a consistent set of {@code
* ResponseStatusException}s.
*/
public String delete(String path, String userId) throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
throw new ResponseStatusException(
HttpStatus.SERVICE_UNAVAILABLE, "AI engine is not enabled");
}
String url = config.getUrl().stripTrailing() + path;
log.debug("Proxying AI engine DELETE request to {}", url);
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Accept", "application/json")
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
.DELETE();
addUserHeader(builder, userId);
HttpResponse<String> response = sendRequest(builder.build());
log.debug("AI engine responded with status {}", response.statusCode());
checkResponseStatus(response);
return response.body();
}
public String get(String path, String userId) throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) {
throw new ResponseStatusException(
@@ -159,15 +201,14 @@ public class AiEngineClient {
String url = config.getUrl().stripTrailing() + path;
log.debug("Proxying AI engine GET request to {}", url);
HttpRequest request =
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Accept", "application/json")
.timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
.GET()
.build();
HttpResponse<String> response = sendRequest(request);
.GET();
addUserHeader(builder, userId);
HttpResponse<String> response = sendRequest(builder.build());
log.debug("AI engine responded with status {}", response.statusCode());
checkResponseStatus(response);
app/proprietary/src/main/java/stirling/software/proprietary/service/AiUserDataService.java
+50
View File
@@ -0,0 +1,50 @@
package stirling.software.proprietary.service;
import java.io.IOException;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.web.server.ResponseStatusException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
/**
* Lifecycle hooks for a user's AI document data on the Python engine.
*
* <p>Today: cleanup on logout. The engine also runs a TTL reaper that catches sessions ended
* without a clean logout (tab close, JWT expiry, engine restart), so this service is the happy-path
* purge, not a hard guarantee. Calls are fire-and-forget on a background thread (Spring's default
* {@code @Async} executor) so an unavailable engine never delays the caller's response.
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class AiUserDataService {
private static final String PURGE_PATH = "/api/v1/documents/by-owner";
private final AiEngineClient aiEngineClient;
/**
* Tell the engine to delete every collection owned by {@code userId}: vector chunks, page text,
* ACL rows, and the owner row itself. Runs asynchronously so the calling thread (typically a
* logout handler) returns immediately; engine errors are logged on the worker thread and never
* propagated. The engine's TTL reaper backstops any miss within ~24h.
*/
@Async
public void purgeUserDocuments(String userId) {
if (userId == null || userId.isBlank()) {
log.debug("Skipping user document purge: no user id");
return;
}
try {
aiEngineClient.delete(PURGE_PATH, userId);
log.debug("Requested document purge for user {}", userId);
} catch (ResponseStatusException e) {
log.warn("AI engine refused document purge for {}: {}", userId, e.getReason());
} catch (IOException e) {
log.warn("Failed to purge documents for {}: {}", userId, e.getMessage());
}
}
}
app/proprietary/src/main/java/stirling/software/proprietary/service/AiWorkflowService.java
+70 -4
View File
@@ -1,6 +1,8 @@
package stirling.software.proprietary.service;
import java.io.IOException;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
@@ -9,6 +11,7 @@ import java.util.stream.Collectors;
import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders;
@@ -24,14 +27,15 @@ import org.springframework.web.multipart.MultipartFile;
import io.github.pixee.security.Filenames;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.FileStorage;
import stirling.software.common.service.InternalApiClient;
import stirling.software.common.service.InternalApiTimeoutException;
import stirling.software.common.service.ToolMetadataService;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.common.util.ExceptionUtils;
import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager;
@@ -49,6 +53,7 @@ import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
import stirling.software.proprietary.security.util.DesktopClientUtils;
import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile;
import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult;
import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact;
@@ -59,7 +64,6 @@ import tools.jackson.databind.ObjectMapper;
@Slf4j
@Service
@RequiredArgsConstructor
public class AiWorkflowService {
private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents";
@@ -74,6 +78,56 @@ public class AiWorkflowService {
private final TempFileManager tempFileManager;
private final FileIdStrategy fileIdStrategy;
private final AiEngineEndpointResolver endpointResolver;
private final UserServiceInterface userService;
private final ApplicationProperties applicationProperties;
public AiWorkflowService(
CustomPDFDocumentFactory pdfDocumentFactory,
AiEngineClient aiEngineClient,
PdfContentExtractor pdfContentExtractor,
ObjectMapper objectMapper,
InternalApiClient internalApiClient,
FileStorage fileStorage,
ToolMetadataService toolMetadataService,
TempFileManager tempFileManager,
FileIdStrategy fileIdStrategy,
AiEngineEndpointResolver endpointResolver,
@Autowired(required = false) UserServiceInterface userService,
ApplicationProperties applicationProperties) {
this.pdfDocumentFactory = pdfDocumentFactory;
this.aiEngineClient = aiEngineClient;
this.pdfContentExtractor = pdfContentExtractor;
this.objectMapper = objectMapper;
this.internalApiClient = internalApiClient;
this.fileStorage = fileStorage;
this.toolMetadataService = toolMetadataService;
this.tempFileManager = tempFileManager;
this.fileIdStrategy = fileIdStrategy;
this.endpointResolver = endpointResolver;
this.userService = userService;
this.applicationProperties = applicationProperties;
}
/**
* How long an AI-workflow-ingested personal doc lives on the engine before the reaper deletes
* it. Mirrors the configured web JWT lifetime, so a stale cookie can never see data the user
* has lost their session to. Org-shared content (when we add it) bypasses this and sends a null
* {@code expiresAt} so it's persistent.
*/
private Duration personalDocTtl() {
int minutes = DesktopClientUtils.getWebTokenExpiryMinutes(applicationProperties);
return Duration.ofMinutes(minutes);
}
/**
* Resolve the currently-authenticated user's id for X-User-Id propagation to the AI engine.
* Returns null when security is disabled (no UserServiceInterface bean) or no one is logged in.
* The engine rejects per-user routes (ingest, search) when this is null; non-tenant routes
* (health, orchestrate without RAG) still work.
*/
private String currentUserId() {
return userService != null ? userService.getCurrentUsername() : null;
}
@FunctionalInterface
public interface ProgressListener {
@@ -298,10 +352,21 @@ public class AiWorkflowService {
}
}
}
// Personal-doc semantics for AI workflows today: caller owns the doc and is its only
// grantee, with a session-bounded expiry so the reaper cleans up if logout misses.
// When org / shared-doc ingestion lands, the caller chooses owner, grantees, and
// expiry (null = persistent) explicitly.
String callerId = currentUserId();
AiDocumentIngestRequest ingestRequest =
new AiDocumentIngestRequest(file.getId(), file.getName(), pages);
new AiDocumentIngestRequest(
file.getId(),
file.getName(),
pages,
callerId,
callerId == null ? List.of() : List.of(callerId),
Instant.now().plus(personalDocTtl()));
String body = objectMapper.writeValueAsString(ingestRequest);
aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body);
aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body, callerId);
log.debug(
"Ingested document: id={}, name={}, pages={}",
file.getId(),
@@ -699,6 +764,7 @@ public class AiWorkflowService {
aiEngineClient.streamPost(
"/api/v1/orchestrator",
requestBody,
currentUserId(),
line -> handleStreamLine(line, listener, resultHolder, errorHolder));
if (errorHolder[0] != null) {
app/proprietary/src/main/java/stirling/software/proprietary/service/MathAuditorOrchestrator.java
+22 -4
View File
@@ -8,13 +8,14 @@ import java.util.List;
import java.util.UUID;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.proprietary.model.api.ai.Evidence;
import stirling.software.proprietary.model.api.ai.Folio;
import stirling.software.proprietary.model.api.ai.FolioManifest;
@@ -40,7 +41,6 @@ import tools.jackson.databind.ObjectMapper;
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class MathAuditorOrchestrator {
private static final String EXAMINE_PATH = "/api/v1/ai/math-auditor-agent/examine";
@@ -50,6 +50,24 @@ public class MathAuditorOrchestrator {
private final CustomPDFDocumentFactory pdfDocumentFactory;
private final PdfContentExtractor pdfContentExtractor;
private final ObjectMapper objectMapper;
private final UserServiceInterface userService;
public MathAuditorOrchestrator(
AiEngineClient aiEngineClient,
CustomPDFDocumentFactory pdfDocumentFactory,
PdfContentExtractor pdfContentExtractor,
ObjectMapper objectMapper,
@Autowired(required = false) UserServiceInterface userService) {
this.aiEngineClient = aiEngineClient;
this.pdfDocumentFactory = pdfDocumentFactory;
this.pdfContentExtractor = pdfContentExtractor;
this.objectMapper = objectMapper;
this.userService = userService;
}
private String currentUserId() {
return userService != null ? userService.getCurrentUsername() : null;
}
/**
* Run a full math audit against the supplied PDF file.
@@ -110,7 +128,7 @@ public class MathAuditorOrchestrator {
EXAMINE_PATH,
manifest.sessionId(),
manifest.round());
String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody);
String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody, currentUserId());
return objectMapper.readValue(responseBody, Requisition.class);
}
@@ -123,7 +141,7 @@ public class MathAuditorOrchestrator {
evidence.sessionId(),
evidence.round(),
evidence.finalRound());
String responseBody = aiEngineClient.post(path, requestBody);
String responseBody = aiEngineClient.post(path, requestBody, currentUserId());
return objectMapper.readValue(responseBody, Verdict.class);
}
app/proprietary/src/main/java/stirling/software/proprietary/service/PdfCommentAgentOrchestrator.java
+23 -3
View File
@@ -10,18 +10,19 @@ import java.util.UUID;
import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.api.comments.AnnotationLocation;
import stirling.software.common.model.api.comments.StickyNoteSpec;
import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.PdfAnnotationService;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineRequest;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineResponse;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentInstruction;
@@ -50,7 +51,6 @@ import tools.jackson.databind.ObjectMapper;
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class PdfCommentAgentOrchestrator {
private static final String GENERATE_PATH = "/api/v1/ai/pdf-comment-agent/generate";
@@ -80,6 +80,26 @@ public class PdfCommentAgentOrchestrator {
private final CustomPDFDocumentFactory pdfDocumentFactory;
private final ObjectMapper objectMapper;
private final PdfAnnotationService pdfAnnotationService;
private final UserServiceInterface userService;
public PdfCommentAgentOrchestrator(
AiEngineClient aiEngineClient,
PdfTextChunkExtractor pdfTextChunkExtractor,
CustomPDFDocumentFactory pdfDocumentFactory,
ObjectMapper objectMapper,
PdfAnnotationService pdfAnnotationService,
@Autowired(required = false) UserServiceInterface userService) {
this.aiEngineClient = aiEngineClient;
this.pdfTextChunkExtractor = pdfTextChunkExtractor;
this.pdfDocumentFactory = pdfDocumentFactory;
this.objectMapper = objectMapper;
this.pdfAnnotationService = pdfAnnotationService;
this.userService = userService;
}
private String currentUserId() {
return userService != null ? userService.getCurrentUsername() : null;
}
/**
* Run the full PDF comment generation flow.
@@ -164,7 +184,7 @@ public class PdfCommentAgentOrchestrator {
PdfCommentEngineRequest engineRequest =
new PdfCommentEngineRequest(sessionId, prompt, chunks);
String requestBody = objectMapper.writeValueAsString(engineRequest);
String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody);
String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody, currentUserId());
PdfCommentEngineResponse engineResponse =
objectMapper.readValue(responseBody, PdfCommentEngineResponse.class);
app/proprietary/src/test/java/stirling/software/proprietary/security/controller/api/AuthControllerLoginTest.java
+2 -1
View File
@@ -79,7 +79,8 @@ class AuthControllerLoginTest {
totpService,
refreshRateLimitService,
securityProperties,
applicationProperties);
applicationProperties,
new stirling.software.proprietary.service.AiUserDataService(null));
mockMvc = MockMvcBuilders.standaloneSetup(controller).build();
}
app/proprietary/src/test/java/stirling/software/proprietary/security/service/JwtServiceTest.java
+29
View File
@@ -251,6 +251,35 @@ class JwtServiceTest {
assertNull(jwtService.extractToken(request));
}
@Test
void testExtractUsernameFromRequestAllowExpiredReturnsUsername() throws Exception {
String username = "alice";
when(keystoreService.getActiveKey()).thenReturn(testVerificationKey);
when(keystoreService.getKeyPair("test-key-id")).thenReturn(Optional.of(testKeyPair));
String token = jwtService.generateToken(username, Collections.emptyMap());
when(request.getHeader("Authorization")).thenReturn("Bearer " + token);
assertEquals(username, jwtService.extractUsernameFromRequestAllowExpired(request));
}
@Test
void testExtractUsernameFromRequestAllowExpiredReturnsNullWhenNoToken() {
when(request.getHeader("Authorization")).thenReturn(null);
assertNull(jwtService.extractUsernameFromRequestAllowExpired(request));
}
@Test
void testExtractUsernameFromRequestAllowExpiredReturnsNullOnGarbageToken() {
// The logout path depends on this helper returning null - never throwing - so a
// malformed JWT never blocks a user from logging out. Any parse/validation problem
// collapses to "no resolvable user".
when(request.getHeader("Authorization")).thenReturn("Bearer not-a-real-jwt");
assertNull(jwtService.extractUsernameFromRequestAllowExpired(request));
}
@Test
void testGenerateTokenWithKeyId() throws Exception {
String username = "testuser";
app/proprietary/src/test/java/stirling/software/proprietary/service/AiEngineClientTest.java
+4 -4
View File
@@ -46,7 +46,7 @@ class AiEngineClientTest {
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}"));
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
assertSame(cause, ex.getCause(), "Original cause should be preserved for diagnostics");
@@ -58,7 +58,7 @@ class AiEngineClientTest {
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}"));
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.GATEWAY_TIMEOUT, ex.getStatusCode());
assertSame(cause, ex.getCause());
@@ -70,7 +70,7 @@ class AiEngineClientTest {
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.get("/x"));
assertThrows(ResponseStatusException.class, () -> client.get("/x", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
}
@@ -80,7 +80,7 @@ class AiEngineClientTest {
applicationProperties.getAiEngine().setEnabled(false);
ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}"));
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
}
app/proprietary/src/test/java/stirling/software/proprietary/service/AiUserDataServiceTest.java
+76
View File
@@ -0,0 +1,76 @@
package stirling.software.proprietary.service;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import java.io.IOException;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.http.HttpStatus;
import org.springframework.web.server.ResponseStatusException;
/**
* Tests the wiring of {@link AiUserDataService}: the path/userId combination forwarded to the
* engine, plus the swallowing-not-throwing behaviour for engine failures. Doesn't try to assert
* {@code @Async} dispatch - that's Spring infrastructure and only fires through the proxy, which a
* unit-level test bypasses by design.
*/
@ExtendWith(MockitoExtension.class)
class AiUserDataServiceTest {
private static final String PURGE_PATH = "/api/v1/documents/by-owner";
@Mock private AiEngineClient aiEngineClient;
private AiUserDataService service;
@BeforeEach
void setUp() {
service = new AiUserDataService(aiEngineClient);
}
@Test
void purgesUnderTheGivenUserIdOnTheByOwnerEndpoint() throws IOException {
service.purgeUserDocuments("alice");
verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice"));
}
@Test
void noOpWhenUserIdIsNull() throws IOException {
service.purgeUserDocuments(null);
verify(aiEngineClient, never()).delete(eq(PURGE_PATH), eq(null));
}
@Test
void noOpWhenUserIdIsBlank() throws IOException {
service.purgeUserDocuments(" ");
verify(aiEngineClient, never()).delete(eq(PURGE_PATH), eq(" "));
}
@Test
void swallowsIoExceptionFromEngine() throws IOException {
doThrow(new IOException("connection refused"))
.when(aiEngineClient)
.delete(eq(PURGE_PATH), eq("alice"));
// Must not throw - the logout path depends on this swallowing failures so the user
// can still log out when the engine is unreachable.
service.purgeUserDocuments("alice");
verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice"));
}
@Test
void swallowsResponseStatusExceptionFromEngine() throws IOException {
doThrow(new ResponseStatusException(HttpStatus.BAD_GATEWAY, "engine returned 502"))
.when(aiEngineClient)
.delete(eq(PURGE_PATH), eq("alice"));
service.purgeUserDocuments("alice");
verify(aiEngineClient, times(1)).delete(eq(PURGE_PATH), eq("alice"));
}
}
app/proprietary/src/test/java/stirling/software/proprietary/service/AiWorkflowServiceTest.java
+12 -7
View File
@@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.never;
@@ -117,7 +118,9 @@ class AiWorkflowServiceTest {
toolMetadataService,
tempFileManager,
fileIdStrategy,
endpointResolver);
endpointResolver,
null,
new ApplicationProperties());
when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of());
}
@@ -479,18 +482,20 @@ class AiWorkflowServiceTest {
{"outcome":"answer","answer":"done","evidence":[]}
""";
}
Consumer<String> consumer = inv.getArgument(2);
Consumer<String> consumer = inv.getArgument(3);
consumer.accept(wrapAsResultEvent(responseJson));
return null;
})
.when(aiEngineClient)
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
.streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any());
AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString());
verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any());
verify(aiEngineClient, times(1))
.postLongRunning(eq("/api/v1/documents"), anyString(), nullable(String.class));
verify(aiEngineClient, times(2))
.streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any());
}
// --- helpers ---
@@ -498,12 +503,12 @@ class AiWorkflowServiceTest {
private void stubOrchestrator(String responseJson) throws IOException {
doAnswer(
inv -> {
Consumer<String> consumer = inv.getArgument(2);
Consumer<String> consumer = inv.getArgument(3);
consumer.accept(wrapAsResultEvent(responseJson));
return null;
})
.when(aiEngineClient)
.streamPost(eq("/api/v1/orchestrator"), anyString(), any());
.streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any());
}
/**
app/proprietary/src/test/java/stirling/software/proprietary/service/PdfCommentAgentOrchestratorTest.java
+9 -7
View File
@@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -72,7 +73,8 @@ class PdfCommentAgentOrchestratorTest {
pdfTextChunkExtractor,
pdfDocumentFactory,
objectMapper,
pdfAnnotationService);
pdfAnnotationService,
null);
}
@Test
@@ -96,7 +98,7 @@ class PdfCommentAgentOrchestratorTest {
new PdfCommentInstruction(
"p1-c0", "Comment on page 1", null, null)),
"reviewed");
when(aiEngineClient.post(anyString(), anyString()))
when(aiEngineClient.post(anyString(), anyString(), nullable(String.class)))
.thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "please comment");
@@ -130,7 +132,7 @@ class PdfCommentAgentOrchestratorTest {
new PdfCommentInstruction("p0-c0", "Valid", null, null),
new PdfCommentInstruction("p999-c999", "Bogus", null, null)),
"mixed");
when(aiEngineClient.post(anyString(), anyString()))
when(aiEngineClient.post(anyString(), anyString(), nullable(String.class)))
.thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "test");
@@ -156,7 +158,7 @@ class PdfCommentAgentOrchestratorTest {
PdfCommentEngineResponse engineResponse =
new PdfCommentEngineResponse("s", List.of(), "no comments worth making");
when(aiEngineClient.post(anyString(), anyString()))
when(aiEngineClient.post(anyString(), anyString(), nullable(String.class)))
.thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "test");
@@ -184,7 +186,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class,
() -> orchestrator.applyComments(input, "whatever"));
assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString());
verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
}
@Test
@@ -197,7 +199,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class,
() -> orchestrator.applyComments(input, tooLong));
assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString());
verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
}
@Test
@@ -209,7 +211,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class,
() -> orchestrator.applyComments(input, " "));
assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString());
verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
}
// ---------------------------------------------------------------------
engine/src/stirling/agents/contradiction/capability.py
+4 -1
View File
@@ -21,6 +21,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.contradiction.detector import ContradictionDetector
from stirling.contracts import AiFile
from stirling.contracts.contradiction import Claim, ContradictionReport
from stirling.models import PrincipalId
logger = logging.getLogger(__name__)
@@ -52,6 +53,7 @@ class ContradictionCapability:
self,
detector: ContradictionDetector,
files: list[AiFile],
principals: list[PrincipalId],
*,
max_audits: int = DEFAULT_MAX_AUDITS,
) -> None:
@@ -59,6 +61,7 @@ class ContradictionCapability:
raise ValueError("max_audits must be >= 1")
self._detector = detector
self._files = files
self._principals = principals
self._max_audits = max_audits
self._audit_count = 0
toolset: FunctionToolset[None] = FunctionToolset()
@@ -121,7 +124,7 @@ class ContradictionCapability:
if not self._files:
return "No documents attached to audit."
report = await self._detector.detect(self._files, query=query)
report = await self._detector.detect(self._files, principals=self._principals, query=query)
formatted = self.format_report(report)
logger.info(
"[contradiction-capability] audit query=%r files=%d -> %d findings, %d chars",
engine/src/stirling/agents/contradiction/detector.py
+11 -4
View File
@@ -41,6 +41,7 @@ from stirling.contracts.contradiction import (
ContradictionSeverity,
)
from stirling.contracts.documents import Page
from stirling.models import PrincipalId
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
@@ -205,8 +206,13 @@ class ContradictionDetector:
# Public entry point
# ------------------------------------------------------------------
async def detect(self, files: list[AiFile], query: str | None = None) -> ContradictionReport:
"""Run the full pipeline over the supplied files.
async def detect(
self,
files: list[AiFile],
principals: list[PrincipalId],
query: str | None = None,
) -> ContradictionReport:
"""Run the full pipeline over the supplied files for ``principals``.
``files`` must have already been ingested (the caller is
responsible for the ``has_collection`` precheck — the question
@@ -234,7 +240,7 @@ class ContradictionDetector:
effective_query = query or "extract claims"
per_file_results = await asyncio.gather(
*(self._extract_claims_for_file(file, effective_query) for file in files),
*(self._extract_claims_for_file(file, effective_query, principals) for file in files),
return_exceptions=True,
)
@@ -312,6 +318,7 @@ class ContradictionDetector:
self,
file: AiFile,
query: str,
principals: list[PrincipalId],
) -> _FileExtractionResult:
"""Run the per-chunk extractor over one file's pages.
@@ -324,7 +331,7 @@ class ContradictionDetector:
``asyncio.gather`` and the mapper's internal semaphore — this
helper itself awaits each step sequentially within one file.
"""
file_pages = await self._runtime.documents.read_pages(file.id)
file_pages = await self._runtime.documents.read_pages(file.id, principals=principals)
if not file_pages:
logger.info(
"[contradiction] no stored pages for %s (id=%s); skipping",
engine/src/stirling/agents/pdf_questions.py
+8 -2
View File
@@ -27,8 +27,9 @@ from stirling.contracts import (
format_file_names,
)
from stirling.documents import RagCapability
from stirling.models import PrincipalId
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
from stirling.services import AppRuntime
from stirling.services import AppRuntime, require_current_user_id
logger = logging.getLogger(__name__)
@@ -161,9 +162,10 @@ class PdfQuestionAgent:
)
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = []
for file in files:
if not await self.runtime.documents.has_collection(file.id):
if not await self.runtime.documents.has_collection(file.id, principals=principals):
missing.append(file)
return missing
@@ -175,8 +177,10 @@ class PdfQuestionAgent:
upstream classifier keeps that judgement in the same call that writes
the answer, and lets the agent mix tools when the question warrants it.
"""
principals = [PrincipalId(require_current_user_id())]
rag = RagCapability(
documents=self.runtime.documents,
principals=principals,
collections=[file.id for file in request.files],
top_k=self.runtime.settings.rag_default_top_k,
max_searches=self.runtime.settings.rag_max_searches,
@@ -184,11 +188,13 @@ class PdfQuestionAgent:
whole_doc = WholeDocReaderCapability(
runtime=self.runtime,
files=request.files,
principals=principals,
reasoner=self._chunked_reasoner,
)
contradiction = ContradictionCapability(
detector=self._contradiction_detector,
files=request.files,
principals=principals,
)
agent = Agent(
model=self.runtime.smart_model,
engine/src/stirling/agents/pdf_review.py
+9 -4
View File
@@ -49,14 +49,14 @@ from stirling.contracts import (
Verdict,
)
from stirling.contracts.ledger import Discrepancy
from stirling.models import ApiModel, ToolEndpoint
from stirling.models import ApiModel, PrincipalId, ToolEndpoint
from stirling.models.agent_tool_models import (
AgentToolId,
MathAuditorAgentParams,
PdfCommentAgentParams,
)
from stirling.models.tool_models import AddCommentsParams
from stirling.services import AppRuntime
from stirling.services import AppRuntime, require_current_user_id
# Fallback right-margin placement used when a finding has no usable
# anchor text. A4/Letter portrait assumed.
@@ -158,7 +158,11 @@ class PdfReviewAgent:
files_to_ingest=missing,
content_types=[PdfContentType.PAGE_TEXT],
)
report = await self._contradiction_detector.detect(request.files, query=request.user_message)
report = await self._contradiction_detector.detect(
request.files,
principals=[PrincipalId(require_current_user_id())],
query=request.user_message,
)
comments_json = await self._build_contradiction_comments_payload(request.user_message, report)
return EditPlanResponse(
summary="",
@@ -193,9 +197,10 @@ class PdfReviewAgent:
)
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = []
for file in files:
if not await self.runtime.documents.has_collection(file.id):
if not await self.runtime.documents.has_collection(file.id, principals=principals):
missing.append(file)
return missing
engine/src/stirling/agents/shared/whole_doc_reader.py
+4 -1
View File
@@ -19,6 +19,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner
from stirling.contracts import AiFile
from stirling.models import PrincipalId
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
@@ -47,12 +48,14 @@ class WholeDocReaderCapability:
self,
runtime: AppRuntime,
files: list[AiFile],
principals: list[PrincipalId],
*,
reasoner: ChunkedReasoner | None = None,
max_reads: int = DEFAULT_MAX_READS,
) -> None:
self._runtime = runtime
self._files = files
self._principals = principals
self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime)
self._max_reads = max_reads
self._read_count = 0
@@ -115,7 +118,7 @@ class WholeDocReaderCapability:
sections: list[str] = []
for file in self._files:
pages = await self._runtime.documents.read_pages(file.id)
pages = await self._runtime.documents.read_pages(file.id, principals=self._principals)
if not pages:
logger.info(
"[whole-doc-reader] no stored pages for %s (id=%s); skipping",
engine/src/stirling/api/app.py
+51
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import asyncio
import logging
from contextlib import asynccontextmanager
from typing import Annotated
@@ -29,8 +31,45 @@ from stirling.api.routes import (
)
from stirling.config import AppSettings, load_settings
from stirling.contracts import HealthResponse
from stirling.documents import DocumentService
from stirling.services import build_runtime, setup_posthog_tracking
logger = logging.getLogger(__name__)
async def _run_expired_doc_reaper(
documents: DocumentService,
interval_seconds: int,
) -> None:
"""Periodically delete documents whose ``expires_at`` has passed.
A reaped collection drops everything rooted at that document. Backstop
for the explicit logout purge: catches sessions that ended without a
clean logout (tab close, JWT expiry, engine restart). Persistent rows
(``expires_at`` null, the shape we use for org-shared docs) are never
touched. Runs until cancelled by the lifespan teardown.
"""
await _reap(documents)
while True:
await asyncio.sleep(interval_seconds)
await _reap(documents)
async def _reap(documents: DocumentService) -> None:
"""One reaper iteration. Logs the deleted count on success and the full
exception with traceback on failure; never re-raises non-cancel errors so
a bad iteration doesn't kill the loop. ``asyncio.CancelledError`` is
re-raised so the lifespan teardown can cancel the task cleanly.
"""
try:
deleted = await documents.reap_expired()
if deleted:
logger.info("Reaped %d expired document collection(s)", deleted)
except asyncio.CancelledError:
raise
except Exception:
logger.exception("Document reaper iteration failed; will retry on next interval")
def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
override = fast_api.dependency_overrides.get(load_settings)
@@ -56,7 +95,19 @@ async def lifespan(fast_api: FastAPI):
tracer_provider = setup_posthog_tracking(settings)
if tracer_provider:
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
reaper_task = asyncio.create_task(
_run_expired_doc_reaper(
runtime.documents,
interval_seconds=settings.documents_reaper_interval_seconds,
),
name="expired-document-reaper",
)
yield
reaper_task.cancel()
try:
await reaper_task
except asyncio.CancelledError:
pass
await runtime.documents.close()
if tracer_provider:
tracer_provider.shutdown()
engine/src/stirling/api/dependencies.py
+20 -2
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from fastapi import Request
from fastapi import HTTPException, Request, status
from stirling.agents import (
ExecutionPlanningAgent,
@@ -12,7 +12,8 @@ from stirling.agents import (
from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.documents import DocumentService
from stirling.services import AppRuntime
from stirling.models import UserId
from stirling.services import AppRuntime, current_user_id
def get_runtime(request: Request) -> AppRuntime:
@@ -49,3 +50,20 @@ def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
def get_pdf_comment_agent(request: Request) -> PdfCommentAgent:
return request.app.state.pdf_comment_agent
def require_user_id() -> UserId:
"""FastAPI dependency for routes that touch per-user storage.
Reads ``X-User-Id`` (already extracted into a ContextVar by ``UserIdMiddleware``)
and returns it. Returns HTTP 401 if the caller didn't supply the header. Apply
to any route that ingests, searches, reads, or deletes document content so
the tenancy gate is enforced at the API boundary.
"""
user_id = current_user_id.get()
if user_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="X-User-Id header is required",
)
return user_id
engine/src/stirling/api/middleware.py
+2 -1
View File
@@ -4,6 +4,7 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin
from starlette.requests import Request
from starlette.responses import Response
from stirling.models import UserId
from stirling.services.tracking import current_user_id
_USER_ID_HEADER = "X-User-Id"
@@ -15,7 +16,7 @@ class UserIdMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
user_id = request.headers.get(_USER_ID_HEADER)
if user_id:
token = current_user_id.set(user_id)
token = current_user_id.set(UserId(user_id))
try:
return await call_next(request)
finally:
engine/src/stirling/api/routes/documents.py
+47 -13
View File
@@ -1,49 +1,83 @@
from __future__ import annotations
import logging
from typing import Annotated
from fastapi import APIRouter, Depends
from stirling.api.dependencies import get_document_service
from stirling.api.dependencies import get_document_service, require_user_id
from stirling.contracts import (
DeleteDocumentResponse,
IngestDocumentRequest,
IngestDocumentResponse,
)
from stirling.contracts.documents import PurgeOwnerResponse
from stirling.documents import DocumentService
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId, UserId
router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
logger = logging.getLogger(__name__)
@router.post("", response_model=IngestDocumentResponse)
async def ingest_document(
request: IngestDocumentRequest,
documents: Annotated[DocumentService, Depends(get_document_service)],
_user_id: Annotated[UserId, Depends(require_user_id)],
) -> IngestDocumentResponse:
"""Replace-ingest a document's content under ``document_id``.
"""Replace-ingest a document's content under ``(owner_id, read_principals)``.
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.
``owner_id`` and ``read_principals`` are required on the request body and
are stored verbatim.
``X-User-Id`` is still required so the caller is authenticated and
PostHog tracks the right user, but the value is not used to constrain
the body.
"""
pages = request.page_text or []
chunks_indexed = await documents.ingest(
collection=request.document_id,
pages=pages,
pages=request.page_text or [],
source=request.source,
owner_id=request.owner_id,
read_principals=request.read_principals,
expires_at=request.expires_at,
)
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed)
@router.delete("/{document_id}", response_model=DeleteDocumentResponse)
@router.delete("/by-id/{document_id}", response_model=DeleteDocumentResponse)
async def delete_document(
document_id: FileId,
documents: Annotated[DocumentService, Depends(get_document_service)],
user_id: Annotated[UserId, Depends(require_user_id)],
) -> DeleteDocumentResponse:
"""Remove a document's content. Idempotent."""
existed = await documents.has_collection(document_id)
"""Remove the caller's copy of this document.
Owner is inferred from the caller - only personal-doc deletes go through
here. Org-doc deletes will need an explicit owner_id once we add the
admin endpoints for them; for now this route can't reach docs owned by
a different principal.
"""
owner_id = OwnerId(user_id)
principals = [PrincipalId(user_id)]
existed = await documents.has_collection(document_id, principals=principals)
if existed:
await documents.delete_collection(document_id)
await documents.delete_collection(document_id, owner_id=owner_id)
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
@router.delete("/by-owner", response_model=PurgeOwnerResponse)
async def purge_caller_documents(
documents: Annotated[DocumentService, Depends(get_document_service)],
user_id: Annotated[UserId, Depends(require_user_id)],
) -> PurgeOwnerResponse:
"""Delete every personal-doc collection owned by the caller.
Called by Java on logout so a user's document content disappears as soon as
the session ends. Org-owned docs (where the caller is a reader but not
the owner) are not touched - only collections whose ``owner_id`` matches
the calling user are removed.
"""
deleted = await documents.purge_owner(OwnerId(user_id))
logger.info("Purged %d collection(s) for owner=%s", deleted, user_id)
return PurgeOwnerResponse(owner_id=OwnerId(user_id), deleted=deleted)
engine/src/stirling/config/settings.py
+4
View File
@@ -37,6 +37,10 @@ class AppSettings(BaseSettings):
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
documents_reaper_interval_seconds: int = Field(
default=900,
validation_alias="STIRLING_DOCUMENTS_REAPER_INTERVAL_SECONDS",
)
# Chunked reasoner settings (whole-document map-reduce).
chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE")
engine/src/stirling/contracts/__init__.py
+2
View File
@@ -42,6 +42,7 @@ from .documents import (
Page,
PageRange,
PageText,
PurgeOwnerResponse,
)
from .execution import (
AgentExecutionRequest,
@@ -137,6 +138,7 @@ __all__ = [
"ContradictionSeverity",
"ConversationMessage",
"DeleteDocumentResponse",
"PurgeOwnerResponse",
"PdfToMarkdownCannotDoResponse",
"PdfToMarkdownOrchestrateResponse",
"PdfToMarkdownRequest",
engine/src/stirling/contracts/documents.py
+19 -3
View File
@@ -1,8 +1,10 @@
from __future__ import annotations
from datetime import datetime
from pydantic import Field
from stirling.models import ApiModel
from stirling.models import ApiModel, OwnerId, PrincipalId
from .common import FileId
@@ -35,20 +37,27 @@ class PageRange(ApiModel):
class IngestDocumentRequest(ApiModel):
"""Replace-ingest a document's content under the given ``document_id``.
"""Replace-ingest a document's content under ``(document_id, owner_id)``.
Each call wipes any previously-stored content for the document and writes
Each call wipes any previously-stored content for the pair 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.
``owner_id`` and ``read_principals`` are required: the engine never
defaults them. Callers must declare ownership and access explicitly.
"""
document_id: FileId = Field(min_length=1)
source: str = Field(min_length=1)
page_text: list[PageText] | None = None
owner_id: OwnerId = Field(min_length=1)
read_principals: list[PrincipalId] = Field(min_length=1)
# When to delete this doc. ``None`` means "persistent" (keep until an explicit delete)
expires_at: datetime | None
class IngestDocumentResponse(ApiModel):
@@ -59,3 +68,10 @@ class IngestDocumentResponse(ApiModel):
class DeleteDocumentResponse(ApiModel):
document_id: FileId
deleted: bool
class PurgeOwnerResponse(ApiModel):
"""Returned by ``DELETE /api/v1/documents/by-owner``."""
owner_id: OwnerId
deleted: int = Field(ge=0)
engine/src/stirling/documents/pgvector_store.py
+228 -60
View File
@@ -1,24 +1,34 @@
from __future__ import annotations
import json
from datetime import datetime
import psycopg
from pgvector.psycopg import register_vector_async
from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
class PgVectorStore(DocumentStore):
"""PostgreSQL + pgvector backed store.
"""PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads.
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.
Holds two tables under the same connection:
Owned tables:
* ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
* ``rag_documents`` - vector chunks for RAG search.
* ``document_pages`` - ordered page text for whole-document reading.
* ``document_acl`` - principals (users, groups, orgs) granted permissions
on a ``(collection, owner_id)`` pair. Read methods take the caller's
principal set and join through this table.
Writes are owner-scoped; reads are ACL-scoped.
"""
def __init__(self, dsn: str) -> None:
@@ -41,60 +51,141 @@ class PgVectorStore(DocumentStore):
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY,
source TEXT NOT NULL
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
source TEXT NOT NULL,
expires_at TIMESTAMP WITH TIME ZONE,
PRIMARY KEY (collection, owner_id)
)
"""
)
# Partial index over rows that can actually expire keeps the reaper
# scan tight even when most rows are persistent (org docs).
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_meta_expires_at "
"ON documents_meta(expires_at) WHERE expires_at IS NOT NULL"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS rag_documents (
id TEXT NOT NULL,
collection TEXT NOT NULL
REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
text TEXT NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
embedding vector NOT NULL,
PRIMARY KEY (id, collection)
PRIMARY KEY (id, collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
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_owner ON rag_documents(collection, owner_id)"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS document_pages (
collection TEXT NOT NULL
REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
page_number INTEGER NOT NULL,
text TEXT NOT NULL,
char_count INTEGER NOT NULL,
PRIMARY KEY (collection, page_number)
PRIMARY KEY (collection, owner_id, page_number),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS document_acl (
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
permission TEXT NOT NULL,
PRIMARY KEY (collection, owner_id, principal_id, permission),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
# Hot path: every read joins ``WHERE principal_id = ANY(...) AND permission = ?``.
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
)
await conn.commit()
self._initialized = True
async def ensure_collection(self, collection: str, source: str) -> None:
# ── lifecycle of the (collection, owner_id) row ────────────────────────
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> 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
INSERT INTO documents_meta (collection, owner_id, source, expires_at)
VALUES (%s, %s, %s, %s)
ON CONFLICT (collection, owner_id) DO UPDATE SET
source = EXCLUDED.source,
expires_at = EXCLUDED.expires_at
""",
(collection, source),
(collection, owner_id, source, expires_at),
)
await conn.commit()
async def purge_owner(self, owner_id: OwnerId) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM documents_meta WHERE owner_id = %s",
(owner_id,),
)
deleted = cur.rowcount
await conn.commit()
return deleted
async def reap_expired(self) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()")
deleted = cur.rowcount
await conn.commit()
return deleted
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
# Cascade FKs handle rag_documents, document_pages, document_acl.
await cur.execute(
"DELETE FROM documents_meta WHERE collection = %s AND owner_id = %s",
(collection, owner_id),
)
deleted = cur.rowcount > 0
await conn.commit()
return deleted
# ── write paths ────────────────────────────────────────────────────────
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
@@ -107,36 +198,100 @@ class PgVectorStore(DocumentStore):
for doc, emb in zip(documents, embeddings):
await cur.execute(
"""
INSERT INTO rag_documents (id, collection, text, metadata, embedding)
VALUES (%s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection)
INSERT INTO rag_documents (id, collection, owner_id, text, metadata, embedding)
VALUES (%s, %s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection, owner_id)
DO UPDATE SET
text = EXCLUDED.text,
metadata = EXCLUDED.metadata,
embedding = EXCLUDED.embedding
""",
(doc.id, collection, doc.text, json.dumps(doc.metadata), emb),
(doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), emb),
)
await conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> 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 AND owner_id = %s",
(collection, owner_id),
)
if pages:
await cur.executemany(
"""
INSERT INTO document_pages (collection, owner_id, page_number, text, char_count)
VALUES (%s, %s, %s, %s, %s)
""",
[(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages],
)
await conn.commit()
# ── ACL management ─────────────────────────────────────────────────────
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
if not principals:
return
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.executemany(
"""
INSERT INTO document_acl (collection, owner_id, principal_id, permission)
VALUES (%s, %s, %s, %s)
ON CONFLICT (collection, owner_id, principal_id, permission) DO NOTHING
""",
[(collection, owner_id, p, _READ_PERMISSION) for p in principals],
)
await conn.commit()
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s",
(collection, owner_id, principal),
)
await conn.commit()
# ── read paths (ACL-gated) ─────────────────────────────────────────────
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
await cur.execute(
"""
SELECT id, text, metadata, 1 - (embedding <=> %s) AS score
FROM rag_documents
WHERE collection = %s
WHERE collection = %s AND owner_id = %s
ORDER BY embedding <=> %s
LIMIT %s
""",
(query_embedding, collection, query_embedding, top_k),
(query_embedding, collection, owner_id, query_embedding, top_k),
)
rows = await cur.fetchall()
@@ -148,71 +303,84 @@ class PgVectorStore(DocumentStore):
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()
@staticmethod
async def _readable_owner_for(
cur: psycopg.AsyncCursor,
collection: str,
principals: list[PrincipalId],
) -> str | None:
"""Resolve which owner_id this caller is reading. ``None`` means no access."""
await cur.execute(
"""
SELECT owner_id FROM document_acl
WHERE collection = %s
AND permission = %s
AND principal_id = ANY(%s)
ORDER BY owner_id
LIMIT 1
""",
(collection, _READ_PERMISSION, list(principals)),
)
row = await cur.fetchone()
return row[0] if row else None
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
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,),
"WHERE collection = %s AND owner_id = %s ORDER BY page_number",
(collection, owner_id),
)
else:
await cur.execute(
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = %s AND page_number BETWEEN %s AND %s "
"WHERE collection = %s AND owner_id = %s "
"AND page_number BETWEEN %s AND %s "
"ORDER BY page_number",
(collection, page_range.start, page_range.end),
(collection, owner_id, 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 has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals:
return False
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
# Cascade FKs handle rag_documents and document_pages.
await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,))
await conn.commit()
owner_id = await self._readable_owner_for(cur, collection, principals)
return owner_id is not None
async def list_collections(self) -> list[str]:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("SELECT collection FROM documents_meta ORDER BY collection")
rows = await cur.fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"SELECT 1 FROM documents_meta WHERE collection = %s",
(collection,),
"""
SELECT DISTINCT collection FROM document_acl
WHERE permission = %s
AND principal_id = ANY(%s)
ORDER BY collection
""",
(_READ_PERMISSION, list(principals)),
)
row = await cur.fetchone()
return row is not None
rows = await cur.fetchall()
return [r[0] for r in rows]
async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release.
engine/src/stirling/documents/rag_capability.py
+14 -8
View File
@@ -8,7 +8,7 @@ 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, PrincipalId
logger = logging.getLogger(__name__)
@@ -16,9 +16,13 @@ logger = logging.getLogger(__name__)
class RagCapability:
"""Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection.
Agents consume this as::
Agents construct one per run, scoped to the caller's principal set, as::
rag = runtime.rag_capability
rag = RagCapability(
documents=self.runtime.documents,
principals=current_principals(),
collections=[file.id for file in request.files],
)
Agent(
...,
instructions=rag.instructions,
@@ -26,20 +30,22 @@ class RagCapability:
)
When no collections are pinned, the instructions are generated dynamically at
run time so the agent sees the current list of collections in the store.
run time so the agent sees the collections the caller can read.
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a
single agent run.
single agent run and binds to one caller's principal set.
"""
def __init__(
self,
documents: DocumentService,
principals: list[PrincipalId],
collections: list[FileId] | None = None,
top_k: int = 5,
max_searches: int = 5,
) -> None:
self._documents = documents
self._principals = principals
self._collections = collections
self._top_k = top_k
self._max_searches = max_searches
@@ -74,7 +80,7 @@ class RagCapability:
)
async def _dynamic_instructions(self) -> str:
collections = await self._documents.list_collections()
collections = await self._documents.list_collections(self._principals)
if collections:
names = ", ".join(collections)
collection_desc = f"the following knowledge base collections: {names}"
@@ -115,12 +121,12 @@ class RagCapability:
if self._collections:
all_results = []
for col in self._collections:
col_results = await self._documents.search(query, collection=col, top_k=k)
col_results = await self._documents.search(query, principals=self._principals, collection=col, top_k=k)
all_results.extend(col_results)
all_results.sort(key=lambda r: r.score, reverse=True)
results = all_results[:k]
else:
results = await self._documents.search(query, top_k=k)
results = await self._documents.search(query, principals=self._principals, top_k=k)
if not results:
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
engine/src/stirling/documents/service.py
+90 -32
View File
@@ -1,11 +1,12 @@
from __future__ import annotations
import logging
from datetime import datetime
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
from stirling.models import FileId, OwnerId, PrincipalId
logger = logging.getLogger(__name__)
@@ -17,7 +18,8 @@ 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``:
Holds two representations of every document under a single
``(collection, owner_id)`` pair:
* **Vector chunks** for RAG-style semantic retrieval (``search``).
* **Ordered pages** for whole-document reading (``read_pages``).
@@ -25,6 +27,15 @@ class DocumentService:
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.
**Owner vs principal.** ``owner_id`` is the tenant (a user or an org).
``principals`` is the caller's accessible principal set, matched against
the ACL on every read. Personal-doc behaviour: ingest with
``owner_id=user:bob`` and the route grants ``user:bob`` read access; Bob's
later searches pass ``principals=[user:bob]`` and find his docs.
Org-doc behaviour: ingest with ``owner_id=org:acme`` and grant read to
whichever groups should see it (``group:engineering``, etc.); members'
principal sets pick the doc up automatically.
"""
def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
@@ -37,20 +48,27 @@ class DocumentService:
collection: FileId,
pages: list[PageText],
source: str,
owner_id: OwnerId,
read_principals: list[PrincipalId],
expires_at: datetime | None,
) -> 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.
Wipes any previously-stored content for ``(collection, owner_id)`` 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)
if not read_principals:
raise ValueError("read_principals must not be empty - every doc needs at least one reader")
await self._store.delete_collection(collection, owner_id)
await self._store.ensure_collection(collection, source, owner_id, expires_at)
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)
await self._store.add_pages(collection, stored_pages, owner_id)
await self._store.grant_read(collection, owner_id, read_principals)
chunks: list[Document] = []
for page in pages:
@@ -70,63 +88,103 @@ class DocumentService:
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)
await self._store.add_documents(collection, chunks, embeddings, owner_id)
return len(chunks)
async def search(
self,
query: str,
principals: list[PrincipalId],
collection: FileId | None = None,
top_k: int | None = None,
) -> list[SearchResult]:
"""Embed query and search across one or all collections.
"""Embed query and search collections readable by ``principals``.
If collection is None, searches all available collections and merges results.
If ``collection`` is supplied, search only that collection (and only
if at least one principal can read it). If ``None``, search every
collection any principal in ``principals`` can read, merge, and
return the top-k.
"""
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):
if not await self._store.has_collection(collection, principals):
return []
return await self._store.search(collection, query_embedding, k)
return await self._store.search(collection, query_embedding, k, principals)
# Search all collections, skipping any that error (e.g. dimension mismatch)
collections = await self._store.list_collections()
# Search every collection the caller can read, skipping any that error
# (e.g. dimension mismatch).
collections = await self._store.list_collections(principals)
all_results: list[SearchResult] = []
for col_name in collections:
try:
results = await self._store.search(col_name, query_embedding, k)
results = await self._store.search(col_name, query_embedding, k, principals)
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)
logger.warning(
"Skipping collection %s during cross-collection search",
col_name,
exc_info=True,
)
# Sort by score descending, return top_k across all collections
# Sort by score descending, return top_k across all the caller's collections
all_results.sort(key=lambda r: r.score, reverse=True)
return all_results[:k]
async def read_pages(
self,
collection: FileId,
principals: list[PrincipalId],
page_range: PageRange | None = None,
) -> list[Page]:
"""Return ordered page text for ``collection``.
"""Return ordered page text for ``collection`` if any principal can read it."""
return await self._store.read_pages(collection, page_range, principals)
Empty list if the collection has no stored pages.
async def delete_collection(self, collection: FileId, owner_id: OwnerId) -> bool:
"""Remove a collection (chunks, pages, ACL).
Returns ``True`` if the collection was found and deleted, ``False`` if
no row matched ``(collection, owner_id)``.
"""
return await self._store.read_pages(collection, page_range)
return await self._store.delete_collection(collection, owner_id)
async def delete_collection(self, collection: FileId) -> None:
"""Remove a collection's chunks and pages."""
await self._store.delete_collection(collection)
async def purge_owner(self, owner_id: OwnerId) -> int:
"""Remove every collection ``owner_id`` owns, including vector chunks,
page text, and ACL rows. Returns the number of collections purged."""
return await self._store.purge_owner(owner_id)
async def has_collection(self, collection: FileId) -> bool:
"""Check whether a collection exists."""
return await self._store.has_collection(collection)
async def reap_expired(self) -> int:
"""Delete collections whose ``expires_at`` is set and in the past.
Persistent collections (``expires_at=null``) are never touched.
Returns the number of collections deleted."""
return await self._store.reap_expired()
async def list_collections(self) -> list[FileId]:
"""List all available collections."""
return [FileId(name) for name in await self._store.list_collections()]
async def has_collection(self, collection: FileId, principals: list[PrincipalId]) -> bool:
"""Check whether at least one principal can read this collection."""
return await self._store.has_collection(collection, principals)
async def list_collections(self, principals: list[PrincipalId]) -> list[FileId]:
"""List collections readable by at least one of ``principals``."""
return [FileId(name) for name in await self._store.list_collections(principals)]
async def grant_read(
self,
collection: FileId,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
"""Grant read access to additional principals on an existing doc."""
await self._store.grant_read(collection, owner_id, principals)
async def revoke(
self,
collection: FileId,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
"""Revoke a principal's access on an existing doc."""
await self._store.revoke(collection, owner_id, principal)
async def close(self) -> None:
"""Release the underlying store's resources."""
engine/src/stirling/documents/sqlite_vec_store.py
+314 -86
View File
@@ -5,19 +5,35 @@ import json
import math
import re
import sqlite3
from datetime import UTC, datetime
from pathlib import Path
import sqlite_vec
from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
# sqlite stores TIMESTAMP as TEXT. We normalise to UTC ISO 8601 ``YYYY-MM-DD HH:MM:SS``
# so lexicographic comparison against ``datetime('now')`` matches chronological order.
_SQLITE_DATETIME_FMT = "%Y-%m-%d %H:%M:%S"
def _to_sqlite_utc(dt: datetime | None) -> str | None:
if dt is None:
return None
if dt.tzinfo is not None:
dt = dt.astimezone(UTC).replace(tzinfo=None)
return dt.strftime(_SQLITE_DATETIME_FMT)
class SqliteVecStore(DocumentStore):
"""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
(detected on first insert). Document metadata lives in a regular table joined by rowid.
Each ``(collection, owner_id)`` pair gets its own `vec0` virtual table with a
fixed embedding dimension (detected on first insert). Document metadata lives
in a regular table joined by rowid.
"""
def __init__(self, db_path: str | Path) -> None:
@@ -51,17 +67,30 @@ class SqliteVecStore(DocumentStore):
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY,
source TEXT NOT NULL
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
source TEXT NOT NULL,
expires_at TIMESTAMP,
PRIMARY KEY (collection, owner_id)
)
"""
)
# The reaper filters on ``expires_at IS NOT NULL AND expires_at < now`` so
# a partial index over non-null rows keeps the scan tight even when most
# rows are persistent (org docs).
self._conn.execute(
"CREATE INDEX IF NOT EXISTS idx_meta_expires_at ON documents_meta(expires_at) WHERE expires_at IS NOT NULL"
)
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS collections (
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
dim INTEGER NOT NULL,
table_name TEXT NOT NULL
table_name TEXT NOT NULL,
PRIMARY KEY (collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
@@ -69,47 +98,154 @@ class SqliteVecStore(DocumentStore):
"""
CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL,
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
text TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}',
vec_rowid INTEGER NOT NULL,
PRIMARY KEY (id, collection)
PRIMARY KEY (id, collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
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_owner ON documents(collection, owner_id)")
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS document_pages (
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
page_number INTEGER NOT NULL,
text TEXT NOT NULL,
char_count INTEGER NOT NULL,
PRIMARY KEY (collection, page_number)
PRIMARY KEY (collection, owner_id, page_number),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
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(
"CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
)
self._conn.execute(
"""
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
""",
(collection, source),
CREATE TABLE IF NOT EXISTS document_acl (
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
permission TEXT NOT NULL,
PRIMARY KEY (collection, owner_id, principal_id, permission),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
# Lookup by principal is the hot path for search/list (every read
# joins through this index). Composite ordering matches the WHERE.
self._conn.execute(
"CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
)
self._conn.commit()
# ── lifecycle of the (collection, owner_id) row ────────────────────────
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_ensure_collection, collection, source, owner_id, expires_at)
def _sync_ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
self._conn.execute(
"""
INSERT INTO documents_meta(collection, owner_id, source, expires_at)
VALUES (?, ?, ?, ?)
ON CONFLICT(collection, owner_id) DO UPDATE SET
source = excluded.source,
expires_at = excluded.expires_at
""",
(collection, owner_id, source, _to_sqlite_utc(expires_at)),
)
self._conn.commit()
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
async with self._lock:
return await asyncio.to_thread(self._sync_delete_collection, collection, owner_id)
def _sync_delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
# (collections, documents, document_pages, document_acl) when documents_meta is deleted.
row = self._conn.execute(
"SELECT table_name FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is not None:
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
cursor = self._conn.execute(
"DELETE FROM documents_meta WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
)
self._conn.commit()
return cursor.rowcount > 0
async def purge_owner(self, owner_id: OwnerId) -> int:
async with self._lock:
return await asyncio.to_thread(self._sync_purge_owner, owner_id)
def _sync_purge_owner(self, owner_id: OwnerId) -> int:
# Drop all vec0 virtual tables for this owner first (FK cascade can't reach them).
vec_tables = [
r[0]
for r in self._conn.execute("SELECT table_name FROM collections WHERE owner_id = ?", (owner_id,)).fetchall()
]
for name in vec_tables:
self._conn.execute(f"DROP TABLE IF EXISTS {name}")
cursor = self._conn.execute("DELETE FROM documents_meta WHERE owner_id = ?", (owner_id,))
self._conn.commit()
return cursor.rowcount
async def reap_expired(self) -> int:
async with self._lock:
return await asyncio.to_thread(self._sync_reap_expired)
def _sync_reap_expired(self) -> int:
# Drop vec0 virtual tables for expired collections first (FK cascade can't reach them).
vec_tables = [
r[0]
for r in self._conn.execute(
"""
SELECT c.table_name FROM collections c
JOIN documents_meta m
ON m.collection = c.collection AND m.owner_id = c.owner_id
WHERE m.expires_at IS NOT NULL AND m.expires_at < datetime('now')
"""
).fetchall()
]
for name in vec_tables:
self._conn.execute(f"DROP TABLE IF EXISTS {name}")
cursor = self._conn.execute(
"DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < datetime('now')"
)
self._conn.commit()
return cursor.rowcount
# ── write paths ────────────────────────────────────────────────────────
@staticmethod
def _sanitize_table_name(collection: str) -> str:
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
return f"vec_{safe}"
def _sanitize_table_name(collection: str, owner_id: OwnerId) -> str:
safe_col = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
safe_owner = re.sub(r"[^a-zA-Z0-9_]", "_", owner_id)
return f"vec_{safe_owner}_{safe_col}"
@staticmethod
def _normalize(vector: list[float]) -> list[float]:
@@ -123,6 +259,7 @@ class SqliteVecStore(DocumentStore):
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
@@ -130,34 +267,40 @@ class SqliteVecStore(DocumentStore):
return
async with self._lock:
await asyncio.to_thread(self._sync_add, collection, documents, embeddings)
await asyncio.to_thread(self._sync_add, collection, documents, embeddings, owner_id)
def _sync_add(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
dim = len(embeddings[0])
row = self._conn.execute("SELECT dim, table_name FROM collections WHERE name = ?", (collection,)).fetchone()
row = self._conn.execute(
"SELECT dim, table_name FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is None:
table_name = self._sanitize_table_name(collection)
table_name = self._sanitize_table_name(collection, owner_id)
self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])")
self._conn.execute(
"INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)",
(collection, dim, table_name),
"INSERT INTO collections(collection, owner_id, dim, table_name) VALUES (?, ?, ?, ?)",
(collection, owner_id, dim, table_name),
)
else:
existing_dim, table_name = row
if existing_dim != dim:
raise ValueError(f"Collection {collection} has dim {existing_dim}, got embedding of dim {dim}")
raise ValueError(
f"Collection {collection} for owner {owner_id} has dim {existing_dim}, got embedding of dim {dim}"
)
# Upsert: delete existing docs with matching IDs first
ids = [doc.id for doc in documents]
placeholders = ",".join("?" * len(ids))
existing = self._conn.execute(
f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
f"SELECT vec_rowid FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, owner_id, *ids),
).fetchall()
if existing:
vec_rowids = [r[0] for r in existing]
@@ -167,8 +310,8 @@ class SqliteVecStore(DocumentStore):
vec_rowids,
)
self._conn.execute(
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, owner_id, *ids),
)
for doc, emb in zip(documents, embeddings):
@@ -179,27 +322,101 @@ class SqliteVecStore(DocumentStore):
)
vec_rowid = cursor.lastrowid
self._conn.execute(
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)",
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid),
"INSERT INTO documents(id, collection, owner_id, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?, ?)",
(doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), vec_rowid),
)
self._conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_add_pages, collection, pages, owner_id)
def _sync_add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
self._conn.execute(
"DELETE FROM document_pages WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
)
if pages:
self._conn.executemany(
"INSERT INTO document_pages(collection, owner_id, page_number, text, char_count) "
"VALUES (?, ?, ?, ?, ?)",
[(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages],
)
self._conn.commit()
# ── ACL management ─────────────────────────────────────────────────────
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
if not principals:
return
async with self._lock:
await asyncio.to_thread(self._sync_grant_read, collection, owner_id, principals)
def _sync_grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
self._conn.executemany(
"""
INSERT INTO document_acl(collection, owner_id, principal_id, permission)
VALUES (?, ?, ?, ?)
ON CONFLICT(collection, owner_id, principal_id, permission) DO NOTHING
""",
[(collection, owner_id, p, _READ_PERMISSION) for p in principals],
)
self._conn.commit()
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_revoke, collection, owner_id, principal)
def _sync_revoke(self, collection: str, owner_id: OwnerId, principal: PrincipalId) -> None:
self._conn.execute(
"DELETE FROM document_acl WHERE collection = ? AND owner_id = ? AND principal_id = ?",
(collection, owner_id, principal),
)
self._conn.commit()
# ── read paths (ACL-gated) ─────────────────────────────────────────────
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
async with self._lock:
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k)
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k, principals)
def _sync_search(
self,
collection: str,
query_embedding: list[float],
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone()
if not principals:
return []
owner_id = self._readable_owner_for(collection, principals)
if owner_id is None:
return []
row = self._conn.execute(
"SELECT table_name, dim FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is None:
return []
table_name, dim = row
@@ -213,11 +430,14 @@ class SqliteVecStore(DocumentStore):
f"""
SELECT d.id, d.text, d.metadata, v.distance
FROM {table_name} v
JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ?
JOIN documents d
ON d.vec_rowid = v.rowid
AND d.collection = ?
AND d.owner_id = ?
WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance
""",
(collection, query_blob, top_k),
(collection, owner_id, query_blob, top_k),
).fetchall()
return [
@@ -233,77 +453,85 @@ class SqliteVecStore(DocumentStore):
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()
def _readable_owner_for(self, collection: str, principals: list[PrincipalId]) -> str | None:
"""Resolve which owner_id this caller is reading. ``None`` means no access."""
placeholders = ",".join("?" * len(principals))
row = self._conn.execute(
f"""
SELECT owner_id FROM document_acl
WHERE collection = ?
AND permission = ?
AND principal_id IN ({placeholders})
ORDER BY owner_id
LIMIT 1
""",
(collection, _READ_PERMISSION, *principals),
).fetchone()
return row[0] if row else None
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
async with self._lock:
return await asyncio.to_thread(self._sync_read_pages, collection, page_range)
return await asyncio.to_thread(self._sync_read_pages, collection, page_range, principals)
def _sync_read_pages(
self,
collection: str,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
if not principals:
return []
owner_id = self._readable_owner_for(collection, principals)
if owner_id is None:
return []
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,),
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = ? AND owner_id = ? ORDER BY page_number",
(collection, owner_id),
).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),
"WHERE collection = ? AND owner_id = ? AND page_number BETWEEN ? AND ? "
"ORDER BY page_number",
(collection, owner_id, 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 has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
async with self._lock:
await asyncio.to_thread(self._sync_delete_collection, collection)
return await asyncio.to_thread(self._sync_has_collection, collection, principals)
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()
if row is not None:
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,))
self._conn.commit()
def _sync_has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals:
return False
return self._readable_owner_for(collection, principals) is not None
async def list_collections(self) -> list[str]:
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
async with self._lock:
return await asyncio.to_thread(self._sync_list_collections)
return await asyncio.to_thread(self._sync_list_collections, principals)
def _sync_list_collections(self) -> list[str]:
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall()
def _sync_list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals:
return []
placeholders = ",".join("?" * len(principals))
rows = self._conn.execute(
f"""
SELECT DISTINCT collection FROM document_acl
WHERE permission = ?
AND principal_id IN ({placeholders})
ORDER BY collection
""",
(_READ_PERMISSION, *principals),
).fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
async with self._lock:
return await asyncio.to_thread(self._sync_has_collection, collection)
def _sync_has_collection(self, collection: str) -> bool:
row = self._conn.execute(
"SELECT 1 FROM documents_meta WHERE collection = ?",
(collection,),
).fetchone()
return row is not None
async def close(self) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_close)
engine/src/stirling/documents/store.py
+75 -33
View File
@@ -2,8 +2,10 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from stirling.contracts.documents import Page, PageRange
from stirling.models import OwnerId, PrincipalId
@dataclass
@@ -40,73 +42,113 @@ class DocumentStore(ABC):
* **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.
"""
# ── lifecycle of the (collection, owner_id) row ────────────────────────
@abstractmethod
async def ensure_collection(self, collection: str, source: str) -> None:
"""Upsert the top-level ``documents_meta`` row for this collection.
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
"""Upsert the top-level ``documents_meta`` row for ``(collection, owner_id)``.
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.
write into child tables that hold a foreign key to the parent row, so
it must exist first.
"""
@abstractmethod
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
"""Remove a collection's chunks, pages, and ACL rows.
Returns ``True`` when a row matched and was deleted, ``False`` when
``(collection, owner_id)`` didn't exist.
"""
@abstractmethod
async def purge_owner(self, owner_id: OwnerId) -> int:
"""Remove every collection (and ACL row) belonging to ``owner_id``.
Returns the number of collections purged.
"""
@abstractmethod
async def reap_expired(self) -> int:
"""Remove collections whose ``expires_at`` is non-null and in the past."""
# ── write paths (scoped by owner) ──────────────────────────────────────
@abstractmethod
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
"""Store vector chunks with their embeddings in the named collection."""
"""Store vector chunks with their embeddings under the owner's collection."""
@abstractmethod
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
"""Replace the stored pages for ``(collection, owner_id)`` with the supplied pages."""
# ── ACL management ─────────────────────────────────────────────────────
@abstractmethod
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
"""Grant read access on ``(collection, owner_id)`` to each principal."""
@abstractmethod
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
"""Remove every permission this principal has on ``(collection, owner_id)``."""
# ── read paths (scoped by ACL principal set) ───────────────────────────
@abstractmethod
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
"""Return the top_k most similar vector chunks from the collection."""
"""Top-k similar vector chunks from ``collection`` that any principal can read.
@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.
Returns an empty list when no principal in ``principals`` has a read
ACL row for this collection (regardless of which owner backs it).
"""
@abstractmethod
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> 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.
"""
"""Return ordered pages for ``collection``. Empty list when no read ACL match."""
@abstractmethod
async def delete_collection(self, collection: str) -> None:
"""Remove a collection's chunks and pages."""
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
"""Check whether any principal in ``principals`` can read this collection."""
@abstractmethod
async def list_collections(self) -> list[str]:
"""Return names of all existing collections."""
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
"""Return collection names readable by at least one of ``principals``."""
@abstractmethod
async def has_collection(self, collection: str) -> bool:
"""Check whether a collection exists."""
# ── lifecycle ──────────────────────────────────────────────────────────
@abstractmethod
async def close(self) -> None:
engine/src/stirling/models/__init__.py
+4 -1
View File
@@ -1,12 +1,15 @@
from . import tool_models
from .base import ApiModel, FileId
from .base import ApiModel, FileId, OwnerId, PrincipalId, UserId
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
__all__ = [
"ApiModel",
"FileId",
"OPERATIONS",
"OwnerId",
"ParamToolModel",
"PrincipalId",
"ToolEndpoint",
"UserId",
"tool_models",
]
engine/src/stirling/models/base.py
+10 -3
View File
@@ -5,11 +5,18 @@ from typing import NewType
from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel
# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's
# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key
# throughout the engine.
# Stable, opaque identifier for a file supplied by the caller
FileId = NewType("FileId", str)
# Stable, opaque identifier for the calling user
UserId = NewType("UserId", str)
# Tenant that owns a document
OwnerId = NewType("OwnerId", str)
# An entity that can hold permissions on a document
PrincipalId = NewType("PrincipalId", str)
class ApiModel(BaseModel):
"""Base for every contract model crossing a service boundary."""
engine/src/stirling/services/__init__.py
+3 -1
View File
@@ -7,14 +7,16 @@ from .progress import (
set_progress_emitter,
)
from .runtime import AppRuntime, build_model_settings, build_runtime
from .tracking import setup_posthog_tracking
from .tracking import current_user_id, require_current_user_id, setup_posthog_tracking
__all__ = [
"AppRuntime",
"ProgressEmitter",
"build_model_settings",
"build_runtime",
"current_user_id",
"emit_progress",
"require_current_user_id",
"reset_progress_emitter",
"set_progress_emitter",
"setup_posthog_tracking",
engine/src/stirling/services/runtime.py
+4 -11
View File
@@ -16,7 +16,6 @@ from stirling.documents import (
DocumentStore,
EmbeddingService,
PgVectorStore,
RagCapability,
SqliteVecStore,
)
@@ -47,7 +46,6 @@ class AppRuntime:
fast_model: Model
smart_model: Model
documents: DocumentService
rag_capability: RagCapability
@property
def fast_model_settings(self) -> ModelSettings:
@@ -89,8 +87,8 @@ def _build_document_store(settings: AppSettings) -> DocumentStore:
assert_never(settings.rag_backend)
def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]:
"""Build the document service and the RAG-search capability that wraps it."""
def _build_documents(settings: AppSettings) -> DocumentService:
"""Build the document service used by per-request RAG capabilities."""
logger.info("Documents: embedding_model=%s", settings.rag_embedding_model)
embedder = EmbeddingService(
model_name=settings.rag_embedding_model,
@@ -98,9 +96,7 @@ def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapabil
chunk_overlap=settings.rag_chunk_overlap,
)
store = _build_document_store(settings)
service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
capability = RagCapability(documents=service, top_k=settings.rag_default_top_k)
return service, capability
return DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
def build_runtime(settings: AppSettings) -> AppRuntime:
@@ -109,14 +105,11 @@ def build_runtime(settings: AppSettings) -> AppRuntime:
validate_structured_output_support(fast_model, settings.fast_model_name)
validate_structured_output_support(smart_model, settings.smart_model_name)
documents, rag_capability = _build_documents(settings)
return AppRuntime(
settings=settings,
fast_model=fast_model,
smart_model=smart_model,
documents=documents,
rag_capability=rag_capability,
documents=_build_documents(settings),
)
engine/src/stirling/services/tracking.py
+16 -1
View File
@@ -27,10 +27,25 @@ from opentelemetry.trace import Span
from posthog.client import Client as PostHogClient
from stirling.config import AppSettings
from stirling.models import UserId
# Per-request user ID, set by middleware from the X-User-Id header.
# When not set, PostHog generates a random ID and marks the event as personless.
current_user_id: ContextVar[str | None] = ContextVar("current_user_id", default=None)
current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None)
def require_current_user_id() -> UserId:
"""Return the request's user ID or raise if the X-User-Id header was missing.
Use at the boundary of any code path that touches per-user document
storage (vector chunks, page text, ACL rows). Routes prefer the FastAPI dependency form
(``Depends(require_user_id)``); agent internals that don't have a request
object in scope call this helper directly to fail closed.
"""
user_id = current_user_id.get()
if user_id is None:
raise RuntimeError("X-User-Id is required for this operation but was not set on the request")
return user_id
class LRUSet:
engine/tests/agents/test_whole_doc_reader.py
+22 -2
View File
@@ -15,9 +15,12 @@ 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.models import FileId, OwnerId, PrincipalId
from stirling.services.runtime import AppRuntime
OWNER = OwnerId("test-user")
PRINCIPALS = [PrincipalId("test-user")]
class StubEmbedder:
"""Deterministic embeddings so tests don't need a real provider."""
@@ -75,7 +78,9 @@ async def test_read_full_document_returns_formatted_notes_for_single_file(
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")
await runtime_with_stub_docs.documents.ingest(
FileId("doc-id"), pages, source="doc.pdf", owner_id=OWNER, read_principals=PRINCIPALS, expires_at=None
)
reasoner = ChunkedReasoner(runtime_with_stub_docs)
canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])]
@@ -83,6 +88,7 @@ async def test_read_full_document_returns_formatted_notes_for_single_file(
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-id", "doc.pdf")],
principals=PRINCIPALS,
reasoner=reasoner,
)
result = await capability._read_full_document("what is in the document?")
@@ -106,6 +112,9 @@ async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs
FileId(cid),
[PageText(page_number=1, text=f"contents of {cid}")],
source=source,
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
)
reasoner = ChunkedReasoner(runtime_with_stub_docs)
@@ -121,6 +130,7 @@ async def test_read_full_document_iterates_multiple_files(runtime_with_stub_docs
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")],
principals=PRINCIPALS,
reasoner=reasoner,
)
result = await capability._read_full_document("compare them")
@@ -140,6 +150,9 @@ async def test_read_full_document_skips_files_without_pages(runtime_with_stub_do
FileId("present"),
[PageText(page_number=1, text="real content")],
source="present.pdf",
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
)
# 'missing' is never ingested -> read_pages returns [].
@@ -149,6 +162,7 @@ async def test_read_full_document_skips_files_without_pages(runtime_with_stub_do
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")],
principals=PRINCIPALS,
reasoner=reasoner,
)
result = await capability._read_full_document("anything")
@@ -167,6 +181,7 @@ async def test_read_full_document_returns_empty_message_when_no_pages_anywhere(
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("nope", "nope.pdf")],
principals=PRINCIPALS,
reasoner=reasoner,
)
result = await capability._read_full_document("anything")
@@ -186,12 +201,16 @@ async def test_read_full_document_budget_hides_tool_when_exhausted(
FileId("doc-id"),
[PageText(page_number=1, text="content")],
source="doc.pdf",
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
)
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")],
principals=PRINCIPALS,
reasoner=reasoner,
max_reads=1,
)
@@ -212,6 +231,7 @@ async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRu
capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")],
principals=PRINCIPALS,
)
text = capability.instructions
engine/tests/contradiction/test_capability.py
+8 -3
View File
@@ -17,9 +17,11 @@ from stirling.contracts.contradiction import (
ContradictionReport,
ContradictionSeverity,
)
from stirling.models import FileId
from stirling.models import FileId, PrincipalId
from stirling.services.runtime import AppRuntime
PRINCIPALS = [PrincipalId("test-user")]
def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name)
@@ -58,7 +60,7 @@ async def test_find_contradictions_returns_formatted_text(runtime: AppRuntime) -
canned = _canned_report()
detector.detect = AsyncMock(return_value=canned)
capability = ContradictionCapability(detector=detector, files=[_file("doc-a", "a.pdf")])
capability = ContradictionCapability(detector=detector, files=[_file("doc-a", "a.pdf")], principals=PRINCIPALS)
result = await capability._find_contradictions("are there inconsistent deadlines?")
detector.detect.assert_awaited_once()
@@ -84,6 +86,7 @@ async def test_budget_gate_hides_tool_after_first_audit(runtime: AppRuntime) ->
capability = ContradictionCapability(
detector=detector,
files=[_file("doc-a", "a.pdf")],
principals=PRINCIPALS,
max_audits=1,
)
# A real, minimal ToolDefinition — the prepare callback returns this
@@ -107,7 +110,7 @@ async def test_budget_gate_hides_tool_after_first_audit(runtime: AppRuntime) ->
async def test_find_contradictions_with_no_files_returns_message(runtime: AppRuntime) -> None:
detector = ContradictionDetector(runtime)
detector.detect = AsyncMock(return_value=_canned_report())
capability = ContradictionCapability(detector=detector, files=[])
capability = ContradictionCapability(detector=detector, files=[], principals=PRINCIPALS)
result = await capability._find_contradictions("anything")
@@ -120,6 +123,7 @@ def test_instructions_mention_attached_files(runtime: AppRuntime) -> None:
capability = ContradictionCapability(
detector=detector,
files=[_file("doc-a", "alpha.pdf"), _file("doc-b", "beta.pdf")],
principals=PRINCIPALS,
)
text = capability.instructions
@@ -149,6 +153,7 @@ def test_instructions_escape_filename_injection_attempt(runtime: AppRuntime) ->
capability = ContradictionCapability(
detector=detector,
files=[_file("doc-evil", evil_name)],
principals=PRINCIPALS,
)
text = capability.instructions
engine/tests/contradiction/test_detector.py
+24 -17
View File
@@ -27,7 +27,7 @@ from stirling.agents.shared.chunked_mapper import ChunkOutput
from stirling.contracts import AiFile
from stirling.contracts.contradiction import ContradictionSeverity
from stirling.contracts.documents import Page, PageRange
from stirling.models import FileId
from stirling.models import FileId, PrincipalId
from stirling.services.runtime import AppRuntime
@@ -58,10 +58,17 @@ def pages_a() -> list[Page]:
]
PRINCIPALS = [PrincipalId("test-user")]
def _install_documents_stub(runtime: AppRuntime, pages_by_id: dict[FileId, list[Page]]) -> None:
"""Patch ``runtime.documents.read_pages`` to return canned pages per file."""
async def _read(collection: FileId, page_range: PageRange | None = None) -> list[Page]:
async def _read(
collection: FileId,
principals: list[PrincipalId],
page_range: PageRange | None = None,
) -> list[Page]:
return pages_by_id.get(collection, [])
# AppRuntime is frozen; monkey-patch the documents service.
@@ -76,7 +83,7 @@ async def test_no_pages_returns_clean_empty_report(runtime: AppRuntime, file_a:
_install_documents_stub(runtime, {file_a.id: []})
detector = ContradictionDetector(runtime)
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert report.contradictions == []
assert report.pages_examined == []
@@ -126,7 +133,7 @@ async def test_happy_path_finds_contradiction_across_two_pages(
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("Examined 2 pages; found 1 contradiction."))
report = await detector.detect([file_a], query="check the deadline")
report = await detector.detect([file_a], principals=PRINCIPALS, query="check the deadline")
assert len(report.contradictions) == 1
c = report.contradictions[0]
@@ -150,12 +157,12 @@ async def test_zero_claims_returns_clean_report(runtime: AppRuntime, file_a: AiF
return_value=[ChunkOutput(pages=[1, 2], output=_ExtractedClaims(claims=[]), label="pages=1-2")]
)
# Stubbing the summary agent is unavoidable (the production code calls
# it on every detect()); we just don't assert on what it returns —
# asserting on the canned value here would only re-prove that AsyncMock
# it on every detect()); we just don't assert on what it returns.
# Asserting on the canned value here would only re-prove that AsyncMock
# works.
detector._summary_agent.run = AsyncMock(return_value=_stub_result("any text"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert report.contradictions == []
assert report.clean is True
@@ -204,7 +211,7 @@ async def test_canonicaliser_accepts_empty_alias_list(runtime: AppRuntime, file_
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert len(report.contradictions) == 1
@@ -332,7 +339,7 @@ async def test_canonicaliser_failure_falls_back_to_lexical_keys(
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
# Lexical key collapses both subjects so the bucket still forms.
assert len(report.contradictions) == 1
@@ -389,7 +396,7 @@ async def test_same_page_contradiction_is_surfaced(runtime: AppRuntime, file_a:
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert len(report.contradictions) == 1
assert report.contradictions[0].severity == ContradictionSeverity.ERROR
@@ -426,7 +433,7 @@ async def test_identical_quote_pair_is_still_dropped(runtime: AppRuntime, file_a
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert report.contradictions == []
@@ -455,7 +462,7 @@ async def test_summary_falls_back_to_deterministic_when_llm_unavailable(
)
detector._summary_agent.run = AsyncMock(side_effect=failure)
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
assert "No contradictions" in report.summary
assert report.clean is True
@@ -499,7 +506,7 @@ async def test_detector_chunk_timeout_falls_through(runtime: AppRuntime, file_a:
detector._pair_detector.run = AsyncMock(side_effect=TimeoutError("simulated"))
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
# Detector timed out so no pairs come back. Crucially: the pipeline
# reached the summary stage rather than crashing earlier, so
@@ -533,7 +540,7 @@ async def test_empty_chunk_with_substantial_content_logs_warning(
detector._summary_agent.run = AsyncMock(return_value=_stub_result("ok"))
with caplog.at_level(logging.WARNING, logger="stirling.agents.contradiction.detector"):
await detector.detect([file_a])
await detector.detect([file_a], principals=PRINCIPALS)
assert any(
"produced 0 claims" in record.getMessage() and "pages=1" in record.getMessage() for record in caplog.records
@@ -580,7 +587,7 @@ async def test_pages_examined_includes_every_attempted_page(runtime: AppRuntime,
detector._pair_detector.run = AsyncMock(return_value=_stub_result(_BucketContradictions(pairs=[])))
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
# Every page the extractor ran against is reported, even page 2
# (which produced no claim).
@@ -655,7 +662,7 @@ async def test_oversized_bucket_windows_translate_indices_globally(runtime: AppR
detector._pair_detector.run = _stub_detector # type: ignore[method-assign]
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done"))
report = await detector.detect([file_a])
report = await detector.detect([file_a], principals=PRINCIPALS)
# Both windows produced one valid pair each; dedup by global (i, j)
# leaves exactly two contradictions.
@@ -791,7 +798,7 @@ async def test_multi_file_pages_dont_collide_in_validation(runtime: AppRuntime)
)
detector._summary_agent.run = AsyncMock(return_value=_stub_result("ok"))
report = await detector.detect([file_a, file_b])
report = await detector.detect([file_a, file_b], principals=PRINCIPALS)
# Both claims validated as verbatim — each against the right file's
# page text. A collision bug would have produced "paraphrased" for at
engine/tests/contradiction/test_question_integration.py
+19 -1
View File
@@ -9,6 +9,7 @@ detector when invoked.
from __future__ import annotations
from collections.abc import Iterator
from dataclasses import replace
import pytest
@@ -22,10 +23,24 @@ from stirling.contracts import (
)
from stirling.contracts.contradiction import Claim
from stirling.documents import DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId, UserId
from stirling.services import current_user_id
from stirling.services.runtime import AppRuntime
from tests.test_pdf_question_agent import StubEmbedder
USER = UserId("test-user")
OWNER = OwnerId("test-user")
OWNER_PRINCIPALS = [PrincipalId("test-user")]
@pytest.fixture(autouse=True)
def _set_user_context() -> Iterator[None]:
token = current_user_id.set(USER)
try:
yield
finally:
current_user_id.reset(token)
def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name)
@@ -69,6 +84,9 @@ async def test_run_answer_agent_builds_agent_with_three_toolsets(
file.id,
[PageText(page_number=1, text="content")],
source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = PdfQuestionAgent(runtime_with_stub_docs)
engine/tests/contradiction/test_review_integration.py
+25 -1
View File
@@ -8,6 +8,7 @@ contradiction and the right cross-references and anchor handling.
from __future__ import annotations
import json
from collections.abc import Iterator
from dataclasses import replace
from typing import Literal
from unittest.mock import AsyncMock
@@ -27,11 +28,25 @@ from stirling.contracts import (
)
from stirling.contracts.contradiction import Claim
from stirling.documents import DocumentService, SqliteVecStore
from stirling.models import FileId, ToolEndpoint
from stirling.models import FileId, OwnerId, PrincipalId, ToolEndpoint, UserId
from stirling.models.tool_models import AddCommentsParams
from stirling.services import current_user_id
from stirling.services.runtime import AppRuntime
from tests.test_pdf_question_agent import StubEmbedder
USER = UserId("test-user")
OWNER = OwnerId("test-user")
OWNER_PRINCIPALS = [PrincipalId("test-user")]
@pytest.fixture(autouse=True)
def _set_user_context() -> Iterator[None]:
token = current_user_id.set(USER)
try:
yield
finally:
current_user_id.reset(token)
def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name)
@@ -88,6 +103,9 @@ async def test_localiser_prompt_escapes_verdict_tag_injection(
file.id,
[PageText(page_number=1, text="x")],
source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = PdfReviewAgent(runtime_with_stub_docs)
@@ -159,6 +177,9 @@ async def test_contradiction_intent_emits_add_comments_plan(
file.id,
[PageText(page_number=1, text="ignored"), PageText(page_number=5, text="ignored")],
source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = PdfReviewAgent(runtime_with_stub_docs)
@@ -264,6 +285,9 @@ async def test_contradiction_takes_precedence_over_math(
file.id,
[PageText(page_number=1, text="x")],
source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = PdfReviewAgent(runtime_with_stub_docs)
engine/tests/test_documents.py
+343 -53
View File
@@ -8,7 +8,17 @@ 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, OwnerId, PrincipalId
# Personal-doc tests reuse the same opaque string in all three roles — keeps the
# defaults of ingest() exercised and tests honest about what the same id means
# (caller, tenant, single ACL grantee). Org-doc tests construct ad-hoc owners
# and principal sets inline.
OWNER = OwnerId("test-user")
OWNER_PRINCIPALS = [PrincipalId("test-user")]
OTHER_OWNER = OwnerId("other-user")
OTHER_OWNER_PRINCIPALS = [PrincipalId("other-user")]
# chunk_text
@@ -56,7 +66,8 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("test-col", "test.pdf")
await store.ensure_collection("test-col", "test.pdf", OWNER, None)
await store.grant_read("test-col", OWNER, OWNER_PRINCIPALS)
docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
@@ -67,9 +78,9 @@ class TestSqliteVecStore:
[0.9, 0.1, 0.0],
[0.0, 0.0, 1.0],
]
await store.add_documents("test-col", docs, embeddings)
await store.add_documents("test-col", docs, embeddings, OWNER)
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, principals=OWNER_PRINCIPALS)
assert len(results) == 2
assert isinstance(results[0], SearchResult)
assert results[0].document.id == "1"
@@ -78,33 +89,36 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("my-collection", "test.pdf")
await store.ensure_collection("my-collection", "test.pdf", OWNER, None)
await store.grant_read("my-collection", OWNER, OWNER_PRINCIPALS)
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]], OWNER)
collections = await store.list_collections()
collections = await store.list_collections(OWNER_PRINCIPALS)
assert "my-collection" in collections
assert await store.has_collection("my-collection") is True
assert await store.has_collection("nonexistent") is False
assert await store.has_collection("my-collection", OWNER_PRINCIPALS) is True
assert await store.has_collection("nonexistent", OWNER_PRINCIPALS) is False
@pytest.mark.anyio
async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("to-delete", "test.pdf")
await store.ensure_collection("to-delete", "test.pdf", OWNER, None)
await store.grant_read("to-delete", OWNER, OWNER_PRINCIPALS)
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]], OWNER)
assert await store.has_collection("to-delete") is True
await store.delete_collection("to-delete")
assert await store.has_collection("to-delete") is False
assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is True
await store.delete_collection("to-delete", OWNER)
assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is False
@pytest.mark.anyio
async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("empty-test", "test.pdf")
await store.ensure_collection("empty-test", "test.pdf", OWNER, None)
await store.grant_read("empty-test", OWNER, OWNER_PRINCIPALS)
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
await store.add_documents("empty-test", docs, [[1.0, 0.0]], OWNER)
results = await store.search("empty-test", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS)
assert len(results) == 1
@pytest.mark.anyio
@@ -112,7 +126,118 @@ class TestSqliteVecStore:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
with pytest.raises(ValueError, match="documents.*embeddings"):
await store.add_documents("bad", docs, [[1.0], [2.0]])
await store.add_documents("bad", docs, [[1.0], [2.0]], OWNER)
@pytest.mark.anyio
async def test_collections_isolated_by_owner(self) -> None:
"""Two owners can store the same collection id; reads stay scoped to ACL."""
store = SqliteVecStore.ephemeral()
await store.ensure_collection("shared-id", "alice.pdf", OWNER, None)
await store.grant_read("shared-id", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("shared-id", "bob.pdf", OTHER_OWNER, None)
await store.grant_read("shared-id", OTHER_OWNER, OTHER_OWNER_PRINCIPALS)
await store.add_documents(
"shared-id",
[Document(id="1", text="alice content", metadata={})],
[[1.0, 0.0]],
OWNER,
)
await store.add_documents(
"shared-id",
[Document(id="1", text="bob content", metadata={})],
[[1.0, 0.0]],
OTHER_OWNER,
)
alice_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS)
bob_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OTHER_OWNER_PRINCIPALS)
assert [r.document.text for r in alice_results] == ["alice content"]
assert [r.document.text for r in bob_results] == ["bob content"]
assert await store.list_collections(OWNER_PRINCIPALS) == ["shared-id"]
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["shared-id"]
await store.delete_collection("shared-id", OWNER)
assert await store.has_collection("shared-id", OWNER_PRINCIPALS) is False
assert await store.has_collection("shared-id", OTHER_OWNER_PRINCIPALS) is True
@pytest.mark.anyio
async def test_purge_owner_removes_only_that_owners_collections(self) -> None:
"""Logout path: one owner's purge must not touch another owner's docs."""
store = SqliteVecStore.ephemeral()
for owner, principals, name in (
(OWNER, OWNER_PRINCIPALS, "a.pdf"),
(OWNER, OWNER_PRINCIPALS, "b.pdf"),
(OTHER_OWNER, OTHER_OWNER_PRINCIPALS, "c.pdf"),
):
await store.ensure_collection(name, name, owner, None)
await store.grant_read(name, owner, principals)
await store.add_documents(name, [Document(id="1", text="x", metadata={})], [[1.0, 0.0]], owner)
deleted = await store.purge_owner(OWNER)
assert deleted == 2
assert await store.list_collections(OWNER_PRINCIPALS) == []
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["c.pdf"]
@pytest.mark.anyio
async def test_reap_expired_drops_collections_past_expires_at(self) -> None:
"""TTL backstop: rows with ``expires_at`` in the past go away on reap."""
from datetime import UTC, datetime, timedelta
store = SqliteVecStore.ephemeral()
now = datetime.now(UTC)
await store.ensure_collection("fresh", "fresh.pdf", OWNER, now + timedelta(hours=1))
await store.grant_read("fresh", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("stale", "stale.pdf", OWNER, now - timedelta(seconds=1))
await store.grant_read("stale", OWNER, OWNER_PRINCIPALS)
deleted = await store.reap_expired()
assert deleted == 1
assert await store.list_collections(OWNER_PRINCIPALS) == ["fresh"]
@pytest.mark.anyio
async def test_reap_expired_keeps_persistent_and_unexpired(self) -> None:
"""Rows with ``expires_at`` in the future stay. Rows with null
``expires_at`` (org docs) are persistent and never touched, regardless
of age."""
from datetime import UTC, datetime, timedelta
store = SqliteVecStore.ephemeral()
await store.ensure_collection("session", "s.pdf", OWNER, datetime.now(UTC) + timedelta(hours=1))
await store.grant_read("session", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("persistent", "p.pdf", OTHER_OWNER, None)
await store.grant_read("persistent", OTHER_OWNER, OTHER_OWNER_PRINCIPALS)
deleted = await store.reap_expired()
assert deleted == 0
assert await store.list_collections(OWNER_PRINCIPALS) == ["session"]
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["persistent"]
@pytest.mark.anyio
async def test_acl_grants_read_to_extra_principal(self) -> None:
"""A principal without an ACL row can't read; once granted, they can."""
store = SqliteVecStore.ephemeral()
team_principal = PrincipalId("group:engineering")
await store.ensure_collection("doc", "engineering-runbook.pdf", OWNER, None)
await store.grant_read("doc", OWNER, OWNER_PRINCIPALS)
await store.add_documents(
"doc",
[Document(id="1", text="how to deploy", metadata={})],
[[1.0, 0.0]],
OWNER,
)
# Engineering can't see it yet.
assert await store.has_collection("doc", [team_principal]) is False
# Owner grants engineering read access.
await store.grant_read("doc", OWNER, [team_principal])
assert await store.has_collection("doc", [team_principal]) is True
# Revoke kills it.
await store.revoke("doc", OWNER, team_principal)
assert await store.has_collection("doc", [team_principal]) is False
# Owner still can.
assert await store.has_collection("doc", OWNER_PRINCIPALS) is True
# DocumentService (with stub embedder)
@@ -163,39 +288,131 @@ class TestDocumentService:
@pytest.mark.anyio
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."
count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf")
count = await documents.ingest(
FileId("docs"),
_pages(text),
source="guide.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
assert count > 0
results = await documents.search("Python libraries", collection=FileId("docs"))
results = await documents.search("Python libraries", principals=OWNER_PRINCIPALS, collection=FileId("docs"))
assert len(results) > 0
assert results[0].document.text
@pytest.mark.anyio
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None:
count = await documents.ingest(FileId("docs"), _pages(""), source="empty.pdf")
count = await documents.ingest(
FileId("docs"),
_pages(""),
source="empty.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None:
results = await documents.search("anything", collection=FileId("nonexistent"))
results = await documents.search("anything", principals=OWNER_PRINCIPALS, collection=FileId("nonexistent"))
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf")
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf")
async def test_search_all_collections_for_principal(self, documents: DocumentService) -> None:
await documents.ingest(
FileId("col-a"),
_pages("Machine learning overview."),
source="ml.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Deep learning with neural networks."),
source="dl.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
results = await documents.search("neural networks")
results = await documents.search("neural networks", principals=OWNER_PRINCIPALS)
assert len(results) > 0
@pytest.mark.anyio
async def test_search_does_not_cross_owner_boundary(self, documents: DocumentService) -> None:
"""A search by one principal set never returns docs owned by an unrelated owner."""
await documents.ingest(
FileId("col-a"),
_pages("Alice's private notes."),
source="alice.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Bob's private notes."),
source="bob.pdf",
owner_id=OTHER_OWNER,
read_principals=OTHER_OWNER_PRINCIPALS,
expires_at=None,
)
alice_results = await documents.search("notes", principals=OWNER_PRINCIPALS)
bob_results = await documents.search("notes", principals=OTHER_OWNER_PRINCIPALS)
alice_texts = [r.document.text for r in alice_results]
bob_texts = [r.document.text for r in bob_results]
assert any("Alice" in t for t in alice_texts)
assert not any("Bob" in t for t in alice_texts)
assert any("Bob" in t for t in bob_texts)
assert not any("Alice" in t for t in bob_texts)
@pytest.mark.anyio
async def test_org_doc_visible_only_to_granted_group(self, documents: DocumentService) -> None:
"""Org-owned doc with explicit ACL: only granted principals see it."""
org_owner = OwnerId("org:acme")
eng_group = PrincipalId("group:engineering")
hr_group = PrincipalId("group:hr")
await documents.ingest(
FileId("runbook"),
_pages("Production deploy steps."),
source="runbook.pdf",
owner_id=org_owner,
read_principals=[eng_group],
expires_at=None,
)
# Engineering can read.
eng_results = await documents.search("deploy", principals=[eng_group])
assert len(eng_results) > 0
# HR cannot.
hr_results = await documents.search("deploy", principals=[hr_group])
assert hr_results == []
# Caller with both memberships (or org-wide principal) still sees it via eng.
multi_results = await documents.search("deploy", principals=[hr_group, eng_group])
assert len(multi_results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, documents: DocumentService) -> None:
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf")
collections = await documents.list_collections()
await documents.ingest(
FileId("temp"),
_pages("Temporary data."),
source="tmp.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
collections = await documents.list_collections(OWNER_PRINCIPALS)
assert "temp" in collections
await documents.delete_collection(FileId("temp"))
collections = await documents.list_collections()
await documents.delete_collection(FileId("temp"), owner_id=OWNER)
collections = await documents.list_collections(OWNER_PRINCIPALS)
assert "temp" not in collections
@pytest.mark.anyio
@@ -205,9 +422,16 @@ class TestDocumentService:
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")
await documents.ingest(
FileId("ordered"),
pages,
source="ordered.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("ordered"))
stored = await documents.read_pages(FileId("ordered"), principals=OWNER_PRINCIPALS)
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.")
@@ -217,9 +441,13 @@ class TestDocumentService:
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")
await documents.ingest(
FileId("ranged"), pages, source="r.pdf", owner_id=OWNER, read_principals=OWNER_PRINCIPALS, expires_at=None
)
subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4))
subset = await documents.read_pages(
FileId("ranged"), principals=OWNER_PRINCIPALS, page_range=PageRange(start=2, end=4)
)
assert [p.page_number for p in subset] == [2, 3, 4]
@pytest.mark.anyio
@@ -228,14 +456,20 @@ class TestDocumentService:
FileId("doc"),
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
source="v1.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="new")],
source="v2.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("doc"))
stored = await documents.read_pages(FileId("doc"), principals=OWNER_PRINCIPALS)
assert [p.page_number for p in stored] == [1]
assert stored[0].text == "new"
@@ -248,9 +482,16 @@ class TestDocumentService:
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")
await documents.ingest(
FileId("with-blanks"),
pages,
source="blanks.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("with-blanks"))
stored = await documents.read_pages(FileId("with-blanks"), principals=OWNER_PRINCIPALS)
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[1].text.strip() == ""
@@ -270,22 +511,36 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None:
cap = RagCapability(documents, collections=[FileId("docs"), FileId("manuals")])
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("docs"), FileId("manuals")])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions = cap.instructions
assert callable(instructions)
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf")
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf")
cap = RagCapability(documents)
await documents.ingest(
FileId("col-a"),
_pages("Alpha content."),
source="a.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Beta content."),
source="b.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -294,7 +549,7 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -302,14 +557,21 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None:
await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf")
cap = RagCapability(documents)
await documents.ingest(
FileId("docs"),
_pages("Python is a programming language."),
source="guide.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
assert "source: guide.pdf" in output
@@ -318,10 +580,24 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf")
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf")
await documents.ingest(
FileId("pinned"),
_pages("Pinned collection content."),
source="pinned.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("other"),
_pages("Content in another collection."),
source="other.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, collections=[FileId("pinned")])
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("pinned")])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@@ -329,9 +605,16 @@ class TestRagCapability:
@pytest.mark.anyio
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))
await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf")
await documents.ingest(
FileId("bulk"),
_pages(paragraphs),
source="bulk.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
assert "[Result 1" in output
assert "[Result 2" in output
@@ -341,8 +624,15 @@ class TestRagCapability:
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
so the agent can no longer call the tool on subsequent turns."""
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf")
cap = RagCapability(documents, max_searches=2)
await documents.ingest(
FileId("docs"),
_pages("Some content."),
source="x.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, max_searches=2)
tool_def = _dummy_tool_def()
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
engine/tests/test_documents_routes.py
+174 -9
View File
@@ -8,7 +8,11 @@ from fastapi.testclient import TestClient
from stirling.api import app
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, PrincipalId, UserId
USER = UserId("test-user")
USER_PRINCIPALS = [PrincipalId("test-user")]
HEADERS = {"X-User-Id": USER}
class StubEmbedder:
@@ -78,7 +82,11 @@ def test_ingest_document_indexes_page_text(client: TestClient, service: Document
{"pageNumber": 1, "text": "The introduction covers the main topic."},
{"pageNumber": 2, "text": "The conclusion summarises the findings."},
],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 200
body = response.json()
@@ -94,7 +102,11 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
"documentId": "replace-me",
"source": "replace-me.pdf",
"pageText": [{"pageNumber": 1, "text": "Original content that existed before."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
# Second ingest with different content should replace the first entirely
response = client.post(
@@ -103,11 +115,15 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
"documentId": "replace-me",
"source": "replace-me.pdf",
"pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 200
results = await service.search("New content", collection=FileId("replace-me"), top_k=5)
results = await service.search("New content", principals=USER_PRINCIPALS, collection=FileId("replace-me"), top_k=5)
texts = [r.document.text for r in results]
assert any("New content" in t for t in texts)
assert not any("Original content" in t for t in texts)
@@ -123,14 +139,28 @@ def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
{"pageNumber": 1, "text": " "},
{"pageNumber": 2, "text": "Real content on page 2."},
],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 200
assert response.json()["chunksIndexed"] >= 1
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
response = client.post("/api/v1/documents", json={"documentId": "empty", "source": "empty.pdf"})
response = client.post(
"/api/v1/documents",
json={
"documentId": "empty",
"source": "empty.pdf",
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 200
assert response.json()["chunksIndexed"] == 0
@@ -139,6 +169,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
response = client.post(
"/api/v1/documents",
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
)
assert response.status_code == 422
@@ -147,6 +178,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
response = client.post(
"/api/v1/documents",
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
)
assert response.status_code == 422
@@ -155,6 +187,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
response = client.post(
"/api/v1/documents",
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
)
assert response.status_code == 422
@@ -166,11 +199,61 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) ->
"documentId": "bad-page",
"source": "bad-page.pdf",
"pageText": [{"pageNumber": 0, "text": "something"}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 422
def test_ingest_document_rejects_missing_owner_id(client: TestClient) -> None:
"""ownerId is required — never derived from the caller. Forgetting it must 422,
not silently fall back to personal-doc semantics."""
response = client.post(
"/api/v1/documents",
json={
"documentId": "no-owner",
"source": "no-owner.pdf",
"pageText": [{"pageNumber": 1, "text": "something"}],
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert response.status_code == 422
def test_ingest_document_rejects_empty_read_principals(client: TestClient) -> None:
"""readPrincipals is required and must not be empty — every doc needs at least one reader."""
response = client.post(
"/api/v1/documents",
json={
"documentId": "no-readers",
"source": "no-readers.pdf",
"pageText": [{"pageNumber": 1, "text": "something"}],
"ownerId": USER,
"readPrincipals": [],
},
headers=HEADERS,
)
assert response.status_code == 422
def test_ingest_document_rejects_missing_user_header(client: TestClient) -> None:
"""The route must refuse to write per-user data when the caller didn't identify themselves."""
response = client.post(
"/api/v1/documents",
json={
"documentId": "doc-1",
"source": "x.pdf",
"pageText": [{"pageNumber": 1, "text": "something"}],
},
)
assert response.status_code == 401
# ── DELETE /documents/{id} ──────────────────────────────────────────────
@@ -181,15 +264,19 @@ def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -
"documentId": "to-delete",
"source": "to-delete.pdf",
"pageText": [{"pageNumber": 1, "text": "Text."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
response = client.delete("/api/v1/documents/to-delete")
response = client.delete("/api/v1/documents/by-id/to-delete", headers=HEADERS)
assert response.status_code == 200
assert response.json() == {"documentId": "to-delete", "deleted": True}
def test_delete_document_is_idempotent(client: TestClient) -> None:
response = client.delete("/api/v1/documents/never-existed")
response = client.delete("/api/v1/documents/by-id/never-existed", headers=HEADERS)
assert response.status_code == 200
assert response.json() == {"documentId": "never-existed", "deleted": False}
@@ -198,8 +285,86 @@ def test_delete_document_is_idempotent(client: TestClient) -> None:
async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
client.post(
"/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."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
)
assert await service.has_collection(FileId("gone"))
client.delete("/api/v1/documents/gone")
assert not await service.has_collection(FileId("gone"))
assert await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS)
client.delete("/api/v1/documents/by-id/gone", headers=HEADERS)
assert not await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS)
def test_delete_document_rejects_missing_user_header(client: TestClient) -> None:
response = client.delete("/api/v1/documents/by-id/anything")
assert response.status_code == 401
@pytest.mark.anyio
async def test_purge_by_owner_removes_only_callers_collections(client: TestClient, service: DocumentService) -> None:
"""Logout path: DELETE /api/v1/documents/by-owner purges only the caller's docs."""
alice = {
"documentId": "alice-doc",
"source": "alice.pdf",
"pageText": [{"pageNumber": 1, "text": "alice content"}],
"ownerId": "alice",
"readPrincipals": ["alice"],
"expiresAt": None,
}
bob = {
"documentId": "bob-doc",
"source": "bob.pdf",
"pageText": [{"pageNumber": 1, "text": "bob content"}],
"ownerId": "bob",
"readPrincipals": ["bob"],
"expiresAt": None,
}
client.post("/api/v1/documents", json=alice, headers={"X-User-Id": "alice"})
client.post("/api/v1/documents", json=bob, headers={"X-User-Id": "bob"})
response = client.delete("/api/v1/documents/by-owner", headers={"X-User-Id": "alice"})
assert response.status_code == 200
assert response.json() == {"ownerId": "alice", "deleted": 1}
# Alice gone, Bob still there.
assert await service.has_collection(FileId("alice-doc"), principals=[PrincipalId("alice")]) is False
assert await service.has_collection(FileId("bob-doc"), principals=[PrincipalId("bob")]) is True
def test_purge_by_owner_is_idempotent(client: TestClient) -> None:
"""Calling purge with no docs is fine — deleted=0 and no error."""
response = client.delete("/api/v1/documents/by-owner", headers=HEADERS)
assert response.status_code == 200
assert response.json() == {"ownerId": USER, "deleted": 0}
def test_purge_by_owner_rejects_missing_user_header(client: TestClient) -> None:
response = client.delete("/api/v1/documents/by-owner")
assert response.status_code == 401
def test_delete_document_only_affects_calling_user(client: TestClient) -> None:
"""Two users with the same document id: one user's delete must not remove the other's."""
alice_body = {
"documentId": "shared",
"source": "shared.pdf",
"pageText": [{"pageNumber": 1, "text": "x"}],
"ownerId": "alice",
"readPrincipals": ["alice"],
"expiresAt": None,
}
bob_body = {**alice_body, "ownerId": "bob", "readPrincipals": ["bob"], "expiresAt": None}
client.post("/api/v1/documents", json=alice_body, headers={"X-User-Id": "alice"})
client.post("/api/v1/documents", json=bob_body, headers={"X-User-Id": "bob"})
alice_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "alice"})
assert alice_delete.json() == {"documentId": "shared", "deleted": True}
# Bob's copy is still there
bob_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "bob"})
assert bob_delete.json() == {"documentId": "shared", "deleted": True}
engine/tests/test_pdf_question_agent.py
+27 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from collections.abc import Iterator
from dataclasses import replace
import pytest
@@ -19,9 +20,25 @@ from stirling.contracts import (
SupportedCapability,
)
from stirling.documents import Document, DocumentService, SqliteVecStore
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId, UserId
from stirling.services import current_user_id
from stirling.services.runtime import AppRuntime
USER = UserId("test-user")
OWNER = OwnerId("test-user")
OWNER_PRINCIPALS = [PrincipalId("test-user")]
@pytest.fixture(autouse=True)
def _set_user_context() -> Iterator[None]:
"""Set current_user_id to a fixed test user for every test in this module so
agent code calling require_current_user_id() doesn't fail closed."""
token = current_user_id.set(USER)
try:
yield
finally:
current_user_id.reset(token)
class StubEmbedder:
"""Deterministic embeddings so RAG lookups work in tests without network."""
@@ -94,6 +111,9 @@ async def test_reports_only_missing_files(runtime_with_stub_rag: AppRuntime) ->
FileId("present-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="present.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = PdfQuestionAgent(runtime_with_stub_rag)
@@ -111,6 +131,9 @@ async def test_returns_grounded_answer_when_all_files_ingested(runtime_with_stub
FileId("invoice-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")],
source="invoice.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = StubPdfQuestionAgent(
runtime_with_stub_rag,
@@ -142,6 +165,9 @@ async def test_returns_not_found_when_answer_not_in_doc(runtime_with_stub_rag: A
FileId("shipping-id"),
[PageText(page_number=1, text="This page contains only a shipping address.")],
source="shipping.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
agent = StubPdfQuestionAgent(
runtime_with_stub_rag,