Set up document management for Stirling Engine (#6476)

# 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
@@ -5,6 +5,7 @@ import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.concurrent.Executor; 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.Qualifier;
import org.springframework.beans.factory.annotation.Value; import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpStatus; 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.model.job.ResultFile;
import stirling.software.common.service.JobOwnershipService; import stirling.software.common.service.JobOwnershipService;
import stirling.software.common.service.TaskManager; import stirling.software.common.service.TaskManager;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent; import stirling.software.proprietary.model.api.ai.AiWorkflowProgressEvent;
import stirling.software.proprietary.model.api.ai.AiWorkflowRequest; import stirling.software.proprietary.model.api.ai.AiWorkflowRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
@@ -58,6 +60,7 @@ public class AiEngineController {
private final TaskManager taskManager; private final TaskManager taskManager;
private final JobOwnershipService jobOwnershipService; private final JobOwnershipService jobOwnershipService;
private final AiEngineEndpointResolver endpointResolver; private final AiEngineEndpointResolver endpointResolver;
private final UserServiceInterface userService;
/** /**
* SSE emitter timeout. Long enough to accommodate multi-gigabyte PDF workflows (OCR on a * 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, @Qualifier("aiStreamExecutor") Executor aiStreamExecutor,
TaskManager taskManager, TaskManager taskManager,
JobOwnershipService jobOwnershipService, JobOwnershipService jobOwnershipService,
AiEngineEndpointResolver endpointResolver) { AiEngineEndpointResolver endpointResolver,
@Autowired(required = false) UserServiceInterface userService) {
this.aiEngineClient = aiEngineClient; this.aiEngineClient = aiEngineClient;
this.aiWorkflowService = aiWorkflowService; this.aiWorkflowService = aiWorkflowService;
this.objectMapper = objectMapper; this.objectMapper = objectMapper;
@@ -82,6 +86,11 @@ public class AiEngineController {
this.taskManager = taskManager; this.taskManager = taskManager;
this.jobOwnershipService = jobOwnershipService; this.jobOwnershipService = jobOwnershipService;
this.endpointResolver = endpointResolver; this.endpointResolver = endpointResolver;
this.userService = userService;
}
private String currentUserId() {
return userService != null ? userService.getCurrentUsername() : null;
} }
@GetMapping("/health") @GetMapping("/health")
@@ -89,7 +98,7 @@ public class AiEngineController {
summary = "AI engine health check", summary = "AI engine health check",
description = "Returns the health status of the AI engine including configured models") description = "Returns the health status of the AI engine including configured models")
public ResponseEntity<String> health() throws IOException { 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); 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"); HttpStatus.BAD_REQUEST, "Request body must be a JSON object");
} }
String forwardedBody = withEnabledEndpoints((ObjectNode) parsed); 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); return ResponseEntity.ok().contentType(MediaType.APPLICATION_JSON).body(response);
} }
@@ -1,5 +1,6 @@
package stirling.software.proprietary.model.api.ai; package stirling.software.proprietary.model.api.ai;
import java.time.Instant;
import java.util.List; import java.util.List;
import lombok.AllArgsConstructor; 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 * 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 * {@code need_ingest} and the requested document's extracted content must be stored before the
* workflow can continue. * 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 @Data
@NoArgsConstructor @NoArgsConstructor
@@ -21,4 +28,10 @@ public class AiDocumentIngestRequest {
private String source; private String source;
private List<AiPageText> pageText; private List<AiPageText> pageText;
private String ownerId;
private List<String> readPrincipals;
private Instant expiresAt;
} }
@@ -8,6 +8,8 @@ import java.util.List;
import java.util.Locale; import java.util.Locale;
import org.springframework.core.io.Resource; 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.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication; import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken; 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.CertificateUtils;
import stirling.software.proprietary.security.saml2.CustomSaml2AuthenticatedPrincipal; import stirling.software.proprietary.security.saml2.CustomSaml2AuthenticatedPrincipal;
import stirling.software.proprietary.security.service.JwtServiceInterface; import stirling.software.proprietary.security.service.JwtServiceInterface;
import stirling.software.proprietary.service.AiUserDataService;
@Slf4j @Slf4j
@RequiredArgsConstructor @RequiredArgsConstructor
@@ -49,12 +52,22 @@ public class CustomLogoutSuccessHandler extends SimpleUrlLogoutSuccessHandler {
private final JwtServiceInterface jwtService; private final JwtServiceInterface jwtService;
private final AiUserDataService aiUserDataService;
private static final AuthenticationTrustResolver TRUST_RESOLVER =
new AuthenticationTrustResolverImpl();
@Override @Override
@Audited(type = AuditEventType.USER_LOGOUT, level = AuditLevel.BASIC) @Audited(type = AuditEventType.USER_LOGOUT, level = AuditLevel.BASIC)
public void onLogoutSuccess( public void onLogoutSuccess(
HttpServletRequest request, HttpServletResponse response, Authentication authentication) HttpServletRequest request, HttpServletResponse response, Authentication authentication)
throws IOException { throws IOException {
String username = resolveUsername(request, authentication);
if (username != null) {
aiUserDataService.purgeUserDocuments(username);
}
if (!response.isCommitted()) { if (!response.isCommitted()) {
if (authentication != null) { if (authentication != null) {
if (authentication instanceof Saml2Authentication samlAuthentication) { 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 // Redirect for SAML2 authentication logout
private void getRedirect_saml2( private void getRedirect_saml2(
HttpServletRequest request, HttpServletRequest request,
@@ -93,6 +93,7 @@ public class SecurityConfiguration {
licenseSettingsService; licenseSettingsService;
private final ClientRegistrationRepository clientRegistrationRepository; private final ClientRegistrationRepository clientRegistrationRepository;
private final PasswordEncoder passwordEncoder; private final PasswordEncoder passwordEncoder;
private final stirling.software.proprietary.service.AiUserDataService aiUserDataService;
public SecurityConfiguration( public SecurityConfiguration(
PersistentLoginRepository persistentLoginRepository, PersistentLoginRepository persistentLoginRepository,
@@ -115,7 +116,8 @@ public class SecurityConfiguration {
OpenSaml5AuthenticationRequestResolver saml2AuthenticationRequestResolver, OpenSaml5AuthenticationRequestResolver saml2AuthenticationRequestResolver,
@Autowired(required = false) ClientRegistrationRepository clientRegistrationRepository, @Autowired(required = false) ClientRegistrationRepository clientRegistrationRepository,
stirling.software.proprietary.service.UserLicenseSettingsService licenseSettingsService, stirling.software.proprietary.service.UserLicenseSettingsService licenseSettingsService,
PasswordEncoder passwordEncoder) { PasswordEncoder passwordEncoder,
stirling.software.proprietary.service.AiUserDataService aiUserDataService) {
this.userDetailsService = userDetailsService; this.userDetailsService = userDetailsService;
this.userService = userService; this.userService = userService;
this.loginEnabledValue = loginEnabledValue; this.loginEnabledValue = loginEnabledValue;
@@ -135,6 +137,7 @@ public class SecurityConfiguration {
this.clientRegistrationRepository = clientRegistrationRepository; this.clientRegistrationRepository = clientRegistrationRepository;
this.licenseSettingsService = licenseSettingsService; this.licenseSettingsService = licenseSettingsService;
this.passwordEncoder = passwordEncoder; this.passwordEncoder = passwordEncoder;
this.aiUserDataService = aiUserDataService;
} }
/** /**
@@ -322,7 +325,10 @@ public class SecurityConfiguration {
.matcher("/logout")) .matcher("/logout"))
.logoutSuccessHandler( .logoutSuccessHandler(
new CustomLogoutSuccessHandler( new CustomLogoutSuccessHandler(
securityProperties, appConfig, jwtService)) securityProperties,
appConfig,
jwtService,
aiUserDataService))
.clearAuthentication(true) .clearAuthentication(true)
.invalidateHttpSession(true) .invalidateHttpSession(true)
.deleteCookies("JSESSIONID", "remember-me", "stirling_jwt")); .deleteCookies("JSESSIONID", "remember-me", "stirling_jwt"));
@@ -44,6 +44,7 @@ import stirling.software.proprietary.security.service.RefreshRateLimitService;
import stirling.software.proprietary.security.service.TotpService; import stirling.software.proprietary.security.service.TotpService;
import stirling.software.proprietary.security.service.UserService; import stirling.software.proprietary.security.service.UserService;
import stirling.software.proprietary.security.util.DesktopClientUtils; import stirling.software.proprietary.security.util.DesktopClientUtils;
import stirling.software.proprietary.service.AiUserDataService;
/** REST API Controller for authentication operations. */ /** REST API Controller for authentication operations. */
@RestController @RestController
@@ -62,6 +63,7 @@ public class AuthController {
private final RefreshRateLimitService refreshRateLimitService; private final RefreshRateLimitService refreshRateLimitService;
private final ApplicationProperties.Security securityProperties; private final ApplicationProperties.Security securityProperties;
private final ApplicationProperties applicationProperties; private final ApplicationProperties applicationProperties;
private final AiUserDataService aiUserDataService;
/** /**
* Login endpoint - replaces Supabase signInWithPassword * Login endpoint - replaces Supabase signInWithPassword
@@ -281,11 +283,13 @@ public class AuthController {
*/ */
@PreAuthorize("!hasAuthority('ROLE_DEMO_USER')") @PreAuthorize("!hasAuthority('ROLE_DEMO_USER')")
@PostMapping("/logout") @PostMapping("/logout")
public ResponseEntity<?> logout(HttpServletResponse response) { public ResponseEntity<?> logout(HttpServletRequest request, HttpServletResponse response) {
try { try {
String username = jwtService.extractUsernameFromRequestAllowExpired(request);
SecurityContextHolder.clearContext(); 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")); return ResponseEntity.ok(Map.of("message", "Logged out successfully"));
@@ -349,6 +349,21 @@ public class JwtService implements JwtServiceInterface {
return null; 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 @Override
public boolean isJwtEnabled() { public boolean isJwtEnabled() {
return v2Enabled; return v2Enabled;
@@ -93,6 +93,19 @@ public interface JwtServiceInterface {
*/ */
String extractToken(HttpServletRequest request); 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 * Check if JWT authentication is enabled
* *
@@ -43,22 +43,23 @@ public class AiEngineClient {
this.httpClient = httpClient; 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(); 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 * 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. * 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(); ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
return postWithTimeout( 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 { throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) { if (!config.isEnabled()) {
@@ -69,22 +70,32 @@ public class AiEngineClient {
String url = config.getUrl().stripTrailing() + path; String url = config.getUrl().stripTrailing() + path;
log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds()); log.debug("Proxying AI engine request to {} (timeout {}s)", url, timeout.toSeconds());
HttpRequest request = HttpRequest.Builder builder =
HttpRequest.newBuilder() HttpRequest.newBuilder()
.uri(URI.create(url)) .uri(URI.create(url))
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
.header("Accept", "application/json") .header("Accept", "application/json")
.timeout(timeout) .timeout(timeout)
.POST(HttpRequest.BodyPublishers.ofString(jsonBody)) .POST(HttpRequest.BodyPublishers.ofString(jsonBody));
.build(); addUserHeader(builder, userId);
HttpResponse<String> response = sendRequest(builder.build());
HttpResponse<String> response = sendRequest(request);
log.debug("AI engine responded with status {}", response.statusCode()); log.debug("AI engine responded with status {}", response.statusCode());
checkResponseStatus(response); checkResponseStatus(response);
return response.body(); 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 * 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. * {@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 * 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. * 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 { throws IOException {
ApplicationProperties.AiEngine config = applicationProperties.getAiEngine(); ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) { if (!config.isEnabled()) {
@@ -109,14 +121,15 @@ public class AiEngineClient {
url, url,
timeout.toSeconds()); timeout.toSeconds());
HttpRequest request = HttpRequest.Builder builder =
HttpRequest.newBuilder() HttpRequest.newBuilder()
.uri(URI.create(url)) .uri(URI.create(url))
.header("Content-Type", "application/json") .header("Content-Type", "application/json")
.header("Accept", "application/x-ndjson") .header("Accept", "application/x-ndjson")
.timeout(timeout) .timeout(timeout)
.POST(HttpRequest.BodyPublishers.ofString(jsonBody)) .POST(HttpRequest.BodyPublishers.ofString(jsonBody));
.build(); addUserHeader(builder, userId);
HttpRequest request = builder.build();
HttpResponse<Stream<String>> response; HttpResponse<Stream<String>> response;
try { 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(); ApplicationProperties.AiEngine config = applicationProperties.getAiEngine();
if (!config.isEnabled()) { if (!config.isEnabled()) {
throw new ResponseStatusException( throw new ResponseStatusException(
@@ -159,15 +201,14 @@ public class AiEngineClient {
String url = config.getUrl().stripTrailing() + path; String url = config.getUrl().stripTrailing() + path;
log.debug("Proxying AI engine GET request to {}", url); log.debug("Proxying AI engine GET request to {}", url);
HttpRequest request = HttpRequest.Builder builder =
HttpRequest.newBuilder() HttpRequest.newBuilder()
.uri(URI.create(url)) .uri(URI.create(url))
.header("Accept", "application/json") .header("Accept", "application/json")
.timeout(Duration.ofSeconds(config.getTimeoutSeconds())) .timeout(Duration.ofSeconds(config.getTimeoutSeconds()))
.GET() .GET();
.build(); addUserHeader(builder, userId);
HttpResponse<String> response = sendRequest(builder.build());
HttpResponse<String> response = sendRequest(request);
log.debug("AI engine responded with status {}", response.statusCode()); log.debug("AI engine responded with status {}", response.statusCode());
checkResponseStatus(response); checkResponseStatus(response);
@@ -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());
}
}
}
@@ -1,6 +1,8 @@
package stirling.software.proprietary.service; package stirling.software.proprietary.service;
import java.io.IOException; import java.io.IOException;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.LinkedHashMap; import java.util.LinkedHashMap;
import java.util.List; import java.util.List;
@@ -9,6 +11,7 @@ import java.util.stream.Collectors;
import org.apache.commons.io.FilenameUtils; import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument; import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.FileSystemResource; import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource; import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders; import org.springframework.http.HttpHeaders;
@@ -24,14 +27,15 @@ import org.springframework.web.multipart.MultipartFile;
import io.github.pixee.security.Filenames; import io.github.pixee.security.Filenames;
import lombok.Data; import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.ApplicationProperties;
import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.FileStorage; import stirling.software.common.service.FileStorage;
import stirling.software.common.service.InternalApiClient; import stirling.software.common.service.InternalApiClient;
import stirling.software.common.service.InternalApiTimeoutException; import stirling.software.common.service.InternalApiTimeoutException;
import stirling.software.common.service.ToolMetadataService; import stirling.software.common.service.ToolMetadataService;
import stirling.software.common.service.UserServiceInterface;
import stirling.software.common.util.ExceptionUtils; import stirling.software.common.util.ExceptionUtils;
import stirling.software.common.util.TempFile; import stirling.software.common.util.TempFile;
import stirling.software.common.util.TempFileManager; 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.AiWorkflowRequest;
import stirling.software.proprietary.model.api.ai.AiWorkflowResponse; import stirling.software.proprietary.model.api.ai.AiWorkflowResponse;
import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile; import stirling.software.proprietary.model.api.ai.AiWorkflowResultFile;
import stirling.software.proprietary.security.util.DesktopClientUtils;
import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile; import stirling.software.proprietary.service.PdfContentExtractor.LoadedFile;
import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult; import stirling.software.proprietary.service.PdfContentExtractor.PdfContentResult;
import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact; import stirling.software.proprietary.service.PdfContentExtractor.WorkflowArtifact;
@@ -59,7 +64,6 @@ import tools.jackson.databind.ObjectMapper;
@Slf4j @Slf4j
@Service @Service
@RequiredArgsConstructor
public class AiWorkflowService { public class AiWorkflowService {
private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents"; private static final String DOCUMENTS_ENDPOINT = "/api/v1/documents";
@@ -74,6 +78,56 @@ public class AiWorkflowService {
private final TempFileManager tempFileManager; private final TempFileManager tempFileManager;
private final FileIdStrategy fileIdStrategy; private final FileIdStrategy fileIdStrategy;
private final AiEngineEndpointResolver endpointResolver; 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 @FunctionalInterface
public interface ProgressListener { 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 = 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); String body = objectMapper.writeValueAsString(ingestRequest);
aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body); aiEngineClient.postLongRunning(DOCUMENTS_ENDPOINT, body, callerId);
log.debug( log.debug(
"Ingested document: id={}, name={}, pages={}", "Ingested document: id={}, name={}, pages={}",
file.getId(), file.getId(),
@@ -699,6 +764,7 @@ public class AiWorkflowService {
aiEngineClient.streamPost( aiEngineClient.streamPost(
"/api/v1/orchestrator", "/api/v1/orchestrator",
requestBody, requestBody,
currentUserId(),
line -> handleStreamLine(line, listener, resultHolder, errorHolder)); line -> handleStreamLine(line, listener, resultHolder, errorHolder));
if (errorHolder[0] != null) { if (errorHolder[0] != null) {
@@ -8,13 +8,14 @@ import java.util.List;
import java.util.UUID; import java.util.UUID;
import org.apache.pdfbox.pdmodel.PDDocument; import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service; import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile; import org.springframework.web.multipart.MultipartFile;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import stirling.software.common.service.CustomPDFDocumentFactory; 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.Evidence;
import stirling.software.proprietary.model.api.ai.Folio; import stirling.software.proprietary.model.api.ai.Folio;
import stirling.software.proprietary.model.api.ai.FolioManifest; import stirling.software.proprietary.model.api.ai.FolioManifest;
@@ -40,7 +41,6 @@ import tools.jackson.databind.ObjectMapper;
*/ */
@Slf4j @Slf4j
@Service @Service
@RequiredArgsConstructor
public class MathAuditorOrchestrator { public class MathAuditorOrchestrator {
private static final String EXAMINE_PATH = "/api/v1/ai/math-auditor-agent/examine"; 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 CustomPDFDocumentFactory pdfDocumentFactory;
private final PdfContentExtractor pdfContentExtractor; private final PdfContentExtractor pdfContentExtractor;
private final ObjectMapper objectMapper; 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. * Run a full math audit against the supplied PDF file.
@@ -110,7 +128,7 @@ public class MathAuditorOrchestrator {
EXAMINE_PATH, EXAMINE_PATH,
manifest.sessionId(), manifest.sessionId(),
manifest.round()); manifest.round());
String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody); String responseBody = aiEngineClient.post(EXAMINE_PATH, requestBody, currentUserId());
return objectMapper.readValue(responseBody, Requisition.class); return objectMapper.readValue(responseBody, Requisition.class);
} }
@@ -123,7 +141,7 @@ public class MathAuditorOrchestrator {
evidence.sessionId(), evidence.sessionId(),
evidence.round(), evidence.round(),
evidence.finalRound()); evidence.finalRound());
String responseBody = aiEngineClient.post(path, requestBody); String responseBody = aiEngineClient.post(path, requestBody, currentUserId());
return objectMapper.readValue(responseBody, Verdict.class); return objectMapper.readValue(responseBody, Verdict.class);
} }
@@ -10,18 +10,19 @@ import java.util.UUID;
import org.apache.commons.io.FilenameUtils; import org.apache.commons.io.FilenameUtils;
import org.apache.pdfbox.pdmodel.PDDocument; import org.apache.pdfbox.pdmodel.PDDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus; import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service; import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile; import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.server.ResponseStatusException; import org.springframework.web.server.ResponseStatusException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import stirling.software.common.model.api.comments.AnnotationLocation; import stirling.software.common.model.api.comments.AnnotationLocation;
import stirling.software.common.model.api.comments.StickyNoteSpec; import stirling.software.common.model.api.comments.StickyNoteSpec;
import stirling.software.common.service.CustomPDFDocumentFactory; import stirling.software.common.service.CustomPDFDocumentFactory;
import stirling.software.common.service.PdfAnnotationService; 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.PdfCommentEngineRequest;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineResponse; import stirling.software.proprietary.model.api.ai.comments.PdfCommentEngineResponse;
import stirling.software.proprietary.model.api.ai.comments.PdfCommentInstruction; import stirling.software.proprietary.model.api.ai.comments.PdfCommentInstruction;
@@ -50,7 +51,6 @@ import tools.jackson.databind.ObjectMapper;
*/ */
@Slf4j @Slf4j
@Service @Service
@RequiredArgsConstructor
public class PdfCommentAgentOrchestrator { public class PdfCommentAgentOrchestrator {
private static final String GENERATE_PATH = "/api/v1/ai/pdf-comment-agent/generate"; 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 CustomPDFDocumentFactory pdfDocumentFactory;
private final ObjectMapper objectMapper; private final ObjectMapper objectMapper;
private final PdfAnnotationService pdfAnnotationService; 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. * Run the full PDF comment generation flow.
@@ -164,7 +184,7 @@ public class PdfCommentAgentOrchestrator {
PdfCommentEngineRequest engineRequest = PdfCommentEngineRequest engineRequest =
new PdfCommentEngineRequest(sessionId, prompt, chunks); new PdfCommentEngineRequest(sessionId, prompt, chunks);
String requestBody = objectMapper.writeValueAsString(engineRequest); String requestBody = objectMapper.writeValueAsString(engineRequest);
String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody); String responseBody = aiEngineClient.post(GENERATE_PATH, requestBody, currentUserId());
PdfCommentEngineResponse engineResponse = PdfCommentEngineResponse engineResponse =
objectMapper.readValue(responseBody, PdfCommentEngineResponse.class); objectMapper.readValue(responseBody, PdfCommentEngineResponse.class);
@@ -79,7 +79,8 @@ class AuthControllerLoginTest {
totpService, totpService,
refreshRateLimitService, refreshRateLimitService,
securityProperties, securityProperties,
applicationProperties); applicationProperties,
new stirling.software.proprietary.service.AiUserDataService(null));
mockMvc = MockMvcBuilders.standaloneSetup(controller).build(); mockMvc = MockMvcBuilders.standaloneSetup(controller).build();
} }
@@ -251,6 +251,35 @@ class JwtServiceTest {
assertNull(jwtService.extractToken(request)); 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 @Test
void testGenerateTokenWithKeyId() throws Exception { void testGenerateTokenWithKeyId() throws Exception {
String username = "testuser"; String username = "testuser";
@@ -46,7 +46,7 @@ class AiEngineClientTest {
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause); when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex = ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
assertSame(cause, ex.getCause(), "Original cause should be preserved for diagnostics"); 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); when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex = ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.GATEWAY_TIMEOUT, ex.getStatusCode()); assertEquals(HttpStatus.GATEWAY_TIMEOUT, ex.getStatusCode());
assertSame(cause, ex.getCause()); assertSame(cause, ex.getCause());
@@ -70,7 +70,7 @@ class AiEngineClientTest {
when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause); when(httpClient.send(any(), any(HttpResponse.BodyHandler.class))).thenThrow(cause);
ResponseStatusException ex = ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.get("/x")); assertThrows(ResponseStatusException.class, () -> client.get("/x", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
} }
@@ -80,7 +80,7 @@ class AiEngineClientTest {
applicationProperties.getAiEngine().setEnabled(false); applicationProperties.getAiEngine().setEnabled(false);
ResponseStatusException ex = ResponseStatusException ex =
assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}")); assertThrows(ResponseStatusException.class, () -> client.post("/x", "{}", null));
assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode()); assertEquals(HttpStatus.SERVICE_UNAVAILABLE, ex.getStatusCode());
} }
@@ -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"));
}
}
@@ -8,6 +8,7 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.never; import static org.mockito.Mockito.never;
@@ -117,7 +118,9 @@ class AiWorkflowServiceTest {
toolMetadataService, toolMetadataService,
tempFileManager, tempFileManager,
fileIdStrategy, fileIdStrategy,
endpointResolver); endpointResolver,
null,
new ApplicationProperties());
when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of()); when(endpointResolver.getEnabledEndpointUrls()).thenReturn(List.of());
} }
@@ -479,18 +482,20 @@ class AiWorkflowServiceTest {
{"outcome":"answer","answer":"done","evidence":[]} {"outcome":"answer","answer":"done","evidence":[]}
"""; """;
} }
Consumer<String> consumer = inv.getArgument(2); Consumer<String> consumer = inv.getArgument(3);
consumer.accept(wrapAsResultEvent(responseJson)); consumer.accept(wrapAsResultEvent(responseJson));
return null; return null;
}) })
.when(aiEngineClient) .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")); AiWorkflowResponse result = service.orchestrate(requestFor(input, "summarise this"));
assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome()); assertEquals(AiWorkflowOutcome.ANSWER, result.getOutcome());
verify(aiEngineClient, times(1)).postLongRunning(eq("/api/v1/documents"), anyString()); verify(aiEngineClient, times(1))
verify(aiEngineClient, times(2)).streamPost(eq("/api/v1/orchestrator"), anyString(), any()); .postLongRunning(eq("/api/v1/documents"), anyString(), nullable(String.class));
verify(aiEngineClient, times(2))
.streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any());
} }
// --- helpers --- // --- helpers ---
@@ -498,12 +503,12 @@ class AiWorkflowServiceTest {
private void stubOrchestrator(String responseJson) throws IOException { private void stubOrchestrator(String responseJson) throws IOException {
doAnswer( doAnswer(
inv -> { inv -> {
Consumer<String> consumer = inv.getArgument(2); Consumer<String> consumer = inv.getArgument(3);
consumer.accept(wrapAsResultEvent(responseJson)); consumer.accept(wrapAsResultEvent(responseJson));
return null; return null;
}) })
.when(aiEngineClient) .when(aiEngineClient)
.streamPost(eq("/api/v1/orchestrator"), anyString(), any()); .streamPost(eq("/api/v1/orchestrator"), anyString(), nullable(String.class), any());
} }
/** /**
@@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.nullable;
import static org.mockito.Mockito.never; import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
@@ -72,7 +73,8 @@ class PdfCommentAgentOrchestratorTest {
pdfTextChunkExtractor, pdfTextChunkExtractor,
pdfDocumentFactory, pdfDocumentFactory,
objectMapper, objectMapper,
pdfAnnotationService); pdfAnnotationService,
null);
} }
@Test @Test
@@ -96,7 +98,7 @@ class PdfCommentAgentOrchestratorTest {
new PdfCommentInstruction( new PdfCommentInstruction(
"p1-c0", "Comment on page 1", null, null)), "p1-c0", "Comment on page 1", null, null)),
"reviewed"); "reviewed");
when(aiEngineClient.post(anyString(), anyString())) when(aiEngineClient.post(anyString(), anyString(), nullable(String.class)))
.thenReturn(objectMapper.writeValueAsString(engineResponse)); .thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "please comment"); AnnotatedPdf result = orchestrator.applyComments(input, "please comment");
@@ -130,7 +132,7 @@ class PdfCommentAgentOrchestratorTest {
new PdfCommentInstruction("p0-c0", "Valid", null, null), new PdfCommentInstruction("p0-c0", "Valid", null, null),
new PdfCommentInstruction("p999-c999", "Bogus", null, null)), new PdfCommentInstruction("p999-c999", "Bogus", null, null)),
"mixed"); "mixed");
when(aiEngineClient.post(anyString(), anyString())) when(aiEngineClient.post(anyString(), anyString(), nullable(String.class)))
.thenReturn(objectMapper.writeValueAsString(engineResponse)); .thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "test"); AnnotatedPdf result = orchestrator.applyComments(input, "test");
@@ -156,7 +158,7 @@ class PdfCommentAgentOrchestratorTest {
PdfCommentEngineResponse engineResponse = PdfCommentEngineResponse engineResponse =
new PdfCommentEngineResponse("s", List.of(), "no comments worth making"); 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)); .thenReturn(objectMapper.writeValueAsString(engineResponse));
AnnotatedPdf result = orchestrator.applyComments(input, "test"); AnnotatedPdf result = orchestrator.applyComments(input, "test");
@@ -184,7 +186,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class, ResponseStatusException.class,
() -> orchestrator.applyComments(input, "whatever")); () -> orchestrator.applyComments(input, "whatever"));
assertEquals(400, ex.getStatusCode().value()); assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString()); verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
} }
@Test @Test
@@ -197,7 +199,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class, ResponseStatusException.class,
() -> orchestrator.applyComments(input, tooLong)); () -> orchestrator.applyComments(input, tooLong));
assertEquals(400, ex.getStatusCode().value()); assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString()); verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
} }
@Test @Test
@@ -209,7 +211,7 @@ class PdfCommentAgentOrchestratorTest {
ResponseStatusException.class, ResponseStatusException.class,
() -> orchestrator.applyComments(input, " ")); () -> orchestrator.applyComments(input, " "));
assertEquals(400, ex.getStatusCode().value()); assertEquals(400, ex.getStatusCode().value());
verify(aiEngineClient, never()).post(anyString(), anyString()); verify(aiEngineClient, never()).post(anyString(), anyString(), nullable(String.class));
} }
// --------------------------------------------------------------------- // ---------------------------------------------------------------------
@@ -21,6 +21,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.contradiction.detector import ContradictionDetector from stirling.agents.contradiction.detector import ContradictionDetector
from stirling.contracts import AiFile from stirling.contracts import AiFile
from stirling.contracts.contradiction import Claim, ContradictionReport from stirling.contracts.contradiction import Claim, ContradictionReport
from stirling.models import PrincipalId
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -52,6 +53,7 @@ class ContradictionCapability:
self, self,
detector: ContradictionDetector, detector: ContradictionDetector,
files: list[AiFile], files: list[AiFile],
principals: list[PrincipalId],
*, *,
max_audits: int = DEFAULT_MAX_AUDITS, max_audits: int = DEFAULT_MAX_AUDITS,
) -> None: ) -> None:
@@ -59,6 +61,7 @@ class ContradictionCapability:
raise ValueError("max_audits must be >= 1") raise ValueError("max_audits must be >= 1")
self._detector = detector self._detector = detector
self._files = files self._files = files
self._principals = principals
self._max_audits = max_audits self._max_audits = max_audits
self._audit_count = 0 self._audit_count = 0
toolset: FunctionToolset[None] = FunctionToolset() toolset: FunctionToolset[None] = FunctionToolset()
@@ -121,7 +124,7 @@ class ContradictionCapability:
if not self._files: if not self._files:
return "No documents attached to audit." 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) formatted = self.format_report(report)
logger.info( logger.info(
"[contradiction-capability] audit query=%r files=%d -> %d findings, %d chars", "[contradiction-capability] audit query=%r files=%d -> %d findings, %d chars",
@@ -41,6 +41,7 @@ from stirling.contracts.contradiction import (
ContradictionSeverity, ContradictionSeverity,
) )
from stirling.contracts.documents import Page from stirling.contracts.documents import Page
from stirling.models import PrincipalId
from stirling.services import AppRuntime from stirling.services import AppRuntime
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -205,8 +206,13 @@ class ContradictionDetector:
# Public entry point # Public entry point
# ------------------------------------------------------------------ # ------------------------------------------------------------------
async def detect(self, files: list[AiFile], query: str | None = None) -> ContradictionReport: async def detect(
"""Run the full pipeline over the supplied files. 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 ``files`` must have already been ingested (the caller is
responsible for the ``has_collection`` precheck the question responsible for the ``has_collection`` precheck the question
@@ -234,7 +240,7 @@ class ContradictionDetector:
effective_query = query or "extract claims" effective_query = query or "extract claims"
per_file_results = await asyncio.gather( 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, return_exceptions=True,
) )
@@ -312,6 +318,7 @@ class ContradictionDetector:
self, self,
file: AiFile, file: AiFile,
query: str, query: str,
principals: list[PrincipalId],
) -> _FileExtractionResult: ) -> _FileExtractionResult:
"""Run the per-chunk extractor over one file's pages. """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 ``asyncio.gather`` and the mapper's internal semaphore — this
helper itself awaits each step sequentially within one file. 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: if not file_pages:
logger.info( logger.info(
"[contradiction] no stored pages for %s (id=%s); skipping", "[contradiction] no stored pages for %s (id=%s); skipping",
+8 -2
View File
@@ -27,8 +27,9 @@ from stirling.contracts import (
format_file_names, format_file_names,
) )
from stirling.documents import RagCapability from stirling.documents import RagCapability
from stirling.models import PrincipalId
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams 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__) logger = logging.getLogger(__name__)
@@ -161,9 +162,10 @@ class PdfQuestionAgent:
) )
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]: async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = [] missing: list[AiFile] = []
for file in files: 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) missing.append(file)
return missing return missing
@@ -175,8 +177,10 @@ class PdfQuestionAgent:
upstream classifier keeps that judgement in the same call that writes upstream classifier keeps that judgement in the same call that writes
the answer, and lets the agent mix tools when the question warrants it. the answer, and lets the agent mix tools when the question warrants it.
""" """
principals = [PrincipalId(require_current_user_id())]
rag = RagCapability( rag = RagCapability(
documents=self.runtime.documents, documents=self.runtime.documents,
principals=principals,
collections=[file.id for file in request.files], collections=[file.id for file in request.files],
top_k=self.runtime.settings.rag_default_top_k, top_k=self.runtime.settings.rag_default_top_k,
max_searches=self.runtime.settings.rag_max_searches, max_searches=self.runtime.settings.rag_max_searches,
@@ -184,11 +188,13 @@ class PdfQuestionAgent:
whole_doc = WholeDocReaderCapability( whole_doc = WholeDocReaderCapability(
runtime=self.runtime, runtime=self.runtime,
files=request.files, files=request.files,
principals=principals,
reasoner=self._chunked_reasoner, reasoner=self._chunked_reasoner,
) )
contradiction = ContradictionCapability( contradiction = ContradictionCapability(
detector=self._contradiction_detector, detector=self._contradiction_detector,
files=request.files, files=request.files,
principals=principals,
) )
agent = Agent( agent = Agent(
model=self.runtime.smart_model, model=self.runtime.smart_model,
+9 -4
View File
@@ -49,14 +49,14 @@ from stirling.contracts import (
Verdict, Verdict,
) )
from stirling.contracts.ledger import Discrepancy 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 ( from stirling.models.agent_tool_models import (
AgentToolId, AgentToolId,
MathAuditorAgentParams, MathAuditorAgentParams,
PdfCommentAgentParams, PdfCommentAgentParams,
) )
from stirling.models.tool_models import AddCommentsParams 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 # Fallback right-margin placement used when a finding has no usable
# anchor text. A4/Letter portrait assumed. # anchor text. A4/Letter portrait assumed.
@@ -158,7 +158,11 @@ class PdfReviewAgent:
files_to_ingest=missing, files_to_ingest=missing,
content_types=[PdfContentType.PAGE_TEXT], 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) comments_json = await self._build_contradiction_comments_payload(request.user_message, report)
return EditPlanResponse( return EditPlanResponse(
summary="", summary="",
@@ -193,9 +197,10 @@ class PdfReviewAgent:
) )
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]: async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = [] missing: list[AiFile] = []
for file in files: 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) missing.append(file)
return missing return missing
@@ -19,6 +19,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner from stirling.agents.shared.chunked_reasoner import ChunkedReasoner
from stirling.contracts import AiFile from stirling.contracts import AiFile
from stirling.models import PrincipalId
from stirling.services import AppRuntime from stirling.services import AppRuntime
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -47,12 +48,14 @@ class WholeDocReaderCapability:
self, self,
runtime: AppRuntime, runtime: AppRuntime,
files: list[AiFile], files: list[AiFile],
principals: list[PrincipalId],
*, *,
reasoner: ChunkedReasoner | None = None, reasoner: ChunkedReasoner | None = None,
max_reads: int = DEFAULT_MAX_READS, max_reads: int = DEFAULT_MAX_READS,
) -> None: ) -> None:
self._runtime = runtime self._runtime = runtime
self._files = files self._files = files
self._principals = principals
self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime) self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime)
self._max_reads = max_reads self._max_reads = max_reads
self._read_count = 0 self._read_count = 0
@@ -115,7 +118,7 @@ class WholeDocReaderCapability:
sections: list[str] = [] sections: list[str] = []
for file in self._files: 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: if not pages:
logger.info( logger.info(
"[whole-doc-reader] no stored pages for %s (id=%s); skipping", "[whole-doc-reader] no stored pages for %s (id=%s); skipping",
+51
View File
@@ -1,5 +1,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import logging
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from typing import Annotated from typing import Annotated
@@ -29,8 +31,45 @@ from stirling.api.routes import (
) )
from stirling.config import AppSettings, load_settings from stirling.config import AppSettings, load_settings
from stirling.contracts import HealthResponse from stirling.contracts import HealthResponse
from stirling.documents import DocumentService
from stirling.services import build_runtime, setup_posthog_tracking 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: def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
override = fast_api.dependency_overrides.get(load_settings) override = fast_api.dependency_overrides.get(load_settings)
@@ -56,7 +95,19 @@ async def lifespan(fast_api: FastAPI):
tracer_provider = setup_posthog_tracking(settings) tracer_provider = setup_posthog_tracking(settings)
if tracer_provider: if tracer_provider:
Agent.instrument_all(InstrumentationSettings(tracer_provider=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 yield
reaper_task.cancel()
try:
await reaper_task
except asyncio.CancelledError:
pass
await runtime.documents.close() await runtime.documents.close()
if tracer_provider: if tracer_provider:
tracer_provider.shutdown() tracer_provider.shutdown()
+20 -2
View File
@@ -1,6 +1,6 @@
from __future__ import annotations from __future__ import annotations
from fastapi import Request from fastapi import HTTPException, Request, status
from stirling.agents import ( from stirling.agents import (
ExecutionPlanningAgent, ExecutionPlanningAgent,
@@ -12,7 +12,8 @@ from stirling.agents import (
from stirling.agents.ledger import MathAuditorAgent from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.documents import DocumentService 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: 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: def get_pdf_comment_agent(request: Request) -> PdfCommentAgent:
return request.app.state.pdf_comment_agent 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
+2 -1
View File
@@ -4,6 +4,7 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin
from starlette.requests import Request from starlette.requests import Request
from starlette.responses import Response from starlette.responses import Response
from stirling.models import UserId
from stirling.services.tracking import current_user_id from stirling.services.tracking import current_user_id
_USER_ID_HEADER = "X-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: async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
user_id = request.headers.get(_USER_ID_HEADER) user_id = request.headers.get(_USER_ID_HEADER)
if user_id: if user_id:
token = current_user_id.set(user_id) token = current_user_id.set(UserId(user_id))
try: try:
return await call_next(request) return await call_next(request)
finally: finally:
+47 -13
View File
@@ -1,49 +1,83 @@
from __future__ import annotations from __future__ import annotations
import logging
from typing import Annotated from typing import Annotated
from fastapi import APIRouter, Depends 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 ( from stirling.contracts import (
DeleteDocumentResponse, DeleteDocumentResponse,
IngestDocumentRequest, IngestDocumentRequest,
IngestDocumentResponse, IngestDocumentResponse,
) )
from stirling.contracts.documents import PurgeOwnerResponse
from stirling.documents import DocumentService 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"]) router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
logger = logging.getLogger(__name__)
@router.post("", response_model=IngestDocumentResponse) @router.post("", response_model=IngestDocumentResponse)
async def ingest_document( async def ingest_document(
request: IngestDocumentRequest, request: IngestDocumentRequest,
documents: Annotated[DocumentService, Depends(get_document_service)], documents: Annotated[DocumentService, Depends(get_document_service)],
_user_id: Annotated[UserId, Depends(require_user_id)],
) -> IngestDocumentResponse: ) -> 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: ``owner_id`` and ``read_principals`` are required on the request body and
* embedded chunks for RAG search, are stored verbatim.
* ordered page text for whole-document reading.
Any previously-stored content for this document is removed first. ``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( chunks_indexed = await documents.ingest(
collection=request.document_id, collection=request.document_id,
pages=pages, pages=request.page_text or [],
source=request.source, 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) 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( async def delete_document(
document_id: FileId, document_id: FileId,
documents: Annotated[DocumentService, Depends(get_document_service)], documents: Annotated[DocumentService, Depends(get_document_service)],
user_id: Annotated[UserId, Depends(require_user_id)],
) -> DeleteDocumentResponse: ) -> DeleteDocumentResponse:
"""Remove a document's content. Idempotent.""" """Remove the caller's copy of this document.
existed = await documents.has_collection(document_id)
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: 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) 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)
+4
View File
@@ -37,6 +37,10 @@ class AppSettings(BaseSettings):
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP") rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K") rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES") 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 settings (whole-document map-reduce).
chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE") chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE")
@@ -42,6 +42,7 @@ from .documents import (
Page, Page,
PageRange, PageRange,
PageText, PageText,
PurgeOwnerResponse,
) )
from .execution import ( from .execution import (
AgentExecutionRequest, AgentExecutionRequest,
@@ -137,6 +138,7 @@ __all__ = [
"ContradictionSeverity", "ContradictionSeverity",
"ConversationMessage", "ConversationMessage",
"DeleteDocumentResponse", "DeleteDocumentResponse",
"PurgeOwnerResponse",
"PdfToMarkdownCannotDoResponse", "PdfToMarkdownCannotDoResponse",
"PdfToMarkdownOrchestrateResponse", "PdfToMarkdownOrchestrateResponse",
"PdfToMarkdownRequest", "PdfToMarkdownRequest",
+19 -3
View File
@@ -1,8 +1,10 @@
from __future__ import annotations from __future__ import annotations
from datetime import datetime
from pydantic import Field from pydantic import Field
from stirling.models import ApiModel from stirling.models import ApiModel, OwnerId, PrincipalId
from .common import FileId from .common import FileId
@@ -35,20 +37,27 @@ class PageRange(ApiModel):
class IngestDocumentRequest(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 both the vector-chunk and ordered-page representations from the supplied
pages. pages.
``source`` is a human-readable label (typically the original filename) ``source`` is a human-readable label (typically the original filename)
that flows into chunk metadata so search results are readable when that flows into chunk metadata so search results are readable when
``document_id`` is a hash. ``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) document_id: FileId = Field(min_length=1)
source: str = Field(min_length=1) source: str = Field(min_length=1)
page_text: list[PageText] | None = None 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): class IngestDocumentResponse(ApiModel):
@@ -59,3 +68,10 @@ class IngestDocumentResponse(ApiModel):
class DeleteDocumentResponse(ApiModel): class DeleteDocumentResponse(ApiModel):
document_id: FileId document_id: FileId
deleted: bool deleted: bool
class PurgeOwnerResponse(ApiModel):
"""Returned by ``DELETE /api/v1/documents/by-owner``."""
owner_id: OwnerId
deleted: int = Field(ge=0)
+228 -60
View File
@@ -1,24 +1,34 @@
from __future__ import annotations from __future__ import annotations
import json import json
from datetime import datetime
import psycopg import psycopg
from pgvector.psycopg import register_vector_async from pgvector.psycopg import register_vector_async
from stirling.contracts.documents import Page, PageRange from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
class PgVectorStore(DocumentStore): 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 Connects to an external Postgres instance (DSN provided via config) and uses the
`vector` extension for similarity search. The schema is created on first use. `vector` extension for similarity search. The schema is created on first use.
Holds two tables under the same connection: Owned tables:
* ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
* ``rag_documents`` - vector chunks for RAG search. * ``rag_documents`` - vector chunks for RAG search.
* ``document_pages`` - ordered page text for whole-document reading. * ``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: def __init__(self, dsn: str) -> None:
@@ -41,60 +51,141 @@ class PgVectorStore(DocumentStore):
await cur.execute( await cur.execute(
""" """
CREATE TABLE IF NOT EXISTS documents_meta ( CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY, collection TEXT NOT NULL,
source 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( await cur.execute(
""" """
CREATE TABLE IF NOT EXISTS rag_documents ( CREATE TABLE IF NOT EXISTS rag_documents (
id TEXT NOT NULL, id TEXT NOT NULL,
collection TEXT NOT NULL collection TEXT NOT NULL,
REFERENCES documents_meta(collection) ON DELETE CASCADE, owner_id TEXT NOT NULL,
text TEXT NOT NULL, text TEXT NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}'::jsonb, metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
embedding vector NOT NULL, embedding vector NOT NULL,
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( await cur.execute(
""" """
CREATE TABLE IF NOT EXISTS document_pages ( CREATE TABLE IF NOT EXISTS document_pages (
collection TEXT NOT NULL collection TEXT NOT NULL,
REFERENCES documents_meta(collection) ON DELETE CASCADE, owner_id TEXT NOT NULL,
page_number INTEGER NOT NULL, page_number INTEGER NOT NULL,
text TEXT NOT NULL, text TEXT NOT NULL,
char_count INTEGER 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() await conn.commit()
self._initialized = True 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() await self._ensure_schema()
async with await self._connect() as conn: async with await self._connect() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( await cur.execute(
""" """
INSERT INTO documents_meta (collection, source) INSERT INTO documents_meta (collection, owner_id, source, expires_at)
VALUES (%s, %s) VALUES (%s, %s, %s, %s)
ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source 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() 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( async def add_documents(
self, self,
collection: str, collection: str,
documents: list[Document], documents: list[Document],
embeddings: list[list[float]], embeddings: list[list[float]],
owner_id: OwnerId,
) -> None: ) -> None:
if len(documents) != len(embeddings): if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} 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): for doc, emb in zip(documents, embeddings):
await cur.execute( await cur.execute(
""" """
INSERT INTO rag_documents (id, collection, text, metadata, embedding) INSERT INTO rag_documents (id, collection, owner_id, text, metadata, embedding)
VALUES (%s, %s, %s, %s::jsonb, %s) VALUES (%s, %s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection) ON CONFLICT (id, collection, owner_id)
DO UPDATE SET DO UPDATE SET
text = EXCLUDED.text, text = EXCLUDED.text,
metadata = EXCLUDED.metadata, metadata = EXCLUDED.metadata,
embedding = EXCLUDED.embedding 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() 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( async def search(
self, self,
collection: str, collection: str,
query_embedding: list[float], query_embedding: list[float],
top_k: int = 5, top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]: ) -> list[SearchResult]:
if not principals:
return []
await self._ensure_schema() await self._ensure_schema()
async with await self._connect() as conn: async with await self._connect() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
await cur.execute( await cur.execute(
""" """
SELECT id, text, metadata, 1 - (embedding <=> %s) AS score SELECT id, text, metadata, 1 - (embedding <=> %s) AS score
FROM rag_documents FROM rag_documents
WHERE collection = %s WHERE collection = %s AND owner_id = %s
ORDER BY embedding <=> %s ORDER BY embedding <=> %s
LIMIT %s LIMIT %s
""", """,
(query_embedding, collection, query_embedding, top_k), (query_embedding, collection, owner_id, query_embedding, top_k),
) )
rows = await cur.fetchall() rows = await cur.fetchall()
@@ -148,71 +303,84 @@ class PgVectorStore(DocumentStore):
for r in rows for r in rows
] ]
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: @staticmethod
await self._ensure_schema() async def _readable_owner_for(
async with await self._connect() as conn: cur: psycopg.AsyncCursor,
async with conn.cursor() as cur: collection: str,
await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,)) principals: list[PrincipalId],
if pages: ) -> str | None:
await cur.executemany( """Resolve which owner_id this caller is reading. ``None`` means no access."""
""" await cur.execute(
INSERT INTO document_pages (collection, page_number, text, char_count) """
VALUES (%s, %s, %s, %s) SELECT owner_id FROM document_acl
""", WHERE collection = %s
[(collection, p.page_number, p.text, p.char_count) for p in pages], AND permission = %s
) AND principal_id = ANY(%s)
await conn.commit() 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( async def read_pages(
self, self,
collection: str, collection: str,
page_range: PageRange | None = None, page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]: ) -> list[Page]:
if not principals:
return []
await self._ensure_schema() await self._ensure_schema()
async with await self._connect() as conn: async with await self._connect() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
if page_range is None: if page_range is None:
await cur.execute( await cur.execute(
"SELECT page_number, text, char_count FROM document_pages " "SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = %s ORDER BY page_number", "WHERE collection = %s AND owner_id = %s ORDER BY page_number",
(collection,), (collection, owner_id),
) )
else: else:
await cur.execute( await cur.execute(
"SELECT page_number, text, char_count FROM document_pages " "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", "ORDER BY page_number",
(collection, page_range.start, page_range.end), (collection, owner_id, page_range.start, page_range.end),
) )
rows = await cur.fetchall() rows = await cur.fetchall()
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows] 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() await self._ensure_schema()
async with await self._connect() as conn: async with await self._connect() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
# Cascade FKs handle rag_documents and document_pages. owner_id = await self._readable_owner_for(cur, collection, principals)
await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,)) return owner_id is not None
await conn.commit()
async def list_collections(self) -> list[str]: async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
await self._ensure_schema() if not principals:
async with await self._connect() as conn: return []
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:
await self._ensure_schema() await self._ensure_schema()
async with await self._connect() as conn: async with await self._connect() as conn:
async with conn.cursor() as cur: async with conn.cursor() as cur:
await cur.execute( 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() rows = await cur.fetchall()
return row is not None return [r[0] for r in rows]
async def close(self) -> None: async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release. # Connections are opened and closed per call, so nothing persistent to release.
@@ -8,7 +8,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.documents.service import DocumentService from stirling.documents.service import DocumentService
from stirling.documents.store import SearchResult from stirling.documents.store import SearchResult
from stirling.models import FileId from stirling.models import FileId, PrincipalId
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -16,9 +16,13 @@ logger = logging.getLogger(__name__)
class RagCapability: class RagCapability:
"""Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection. """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( Agent(
..., ...,
instructions=rag.instructions, instructions=rag.instructions,
@@ -26,20 +30,22 @@ class RagCapability:
) )
When no collections are pinned, the instructions are generated dynamically at When no collections are pinned, the instructions are generated dynamically at
run time so the agent sees the current list of collections in the store. run time so the agent sees the collections the caller can read.
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a 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__( def __init__(
self, self,
documents: DocumentService, documents: DocumentService,
principals: list[PrincipalId],
collections: list[FileId] | None = None, collections: list[FileId] | None = None,
top_k: int = 5, top_k: int = 5,
max_searches: int = 5, max_searches: int = 5,
) -> None: ) -> None:
self._documents = documents self._documents = documents
self._principals = principals
self._collections = collections self._collections = collections
self._top_k = top_k self._top_k = top_k
self._max_searches = max_searches self._max_searches = max_searches
@@ -74,7 +80,7 @@ class RagCapability:
) )
async def _dynamic_instructions(self) -> str: async def _dynamic_instructions(self) -> str:
collections = await self._documents.list_collections() collections = await self._documents.list_collections(self._principals)
if collections: if collections:
names = ", ".join(collections) names = ", ".join(collections)
collection_desc = f"the following knowledge base collections: {names}" collection_desc = f"the following knowledge base collections: {names}"
@@ -115,12 +121,12 @@ class RagCapability:
if self._collections: if self._collections:
all_results = [] all_results = []
for col in self._collections: for col in self._collections:
col_results = await self._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.extend(col_results)
all_results.sort(key=lambda r: r.score, reverse=True) all_results.sort(key=lambda r: r.score, reverse=True)
results = all_results[:k] results = all_results[:k]
else: else:
results = await self._documents.search(query, top_k=k) results = await self._documents.search(query, principals=self._principals, top_k=k)
if not results: if not results:
logger.info("[rag] search_knowledge query=%r -> 0 results", query) logger.info("[rag] search_knowledge query=%r -> 0 results", query)
+90 -32
View File
@@ -1,11 +1,12 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from datetime import datetime
from stirling.contracts.documents import Page, PageRange, PageText from stirling.contracts.documents import Page, PageRange, PageText
from stirling.documents.embedder import EmbeddingService from stirling.documents.embedder import EmbeddingService
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage 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__) logger = logging.getLogger(__name__)
@@ -17,7 +18,8 @@ PAGE_TEXT_CONTENT_TYPE = "page_text"
class DocumentService: class DocumentService:
"""Top-level facade for stored document content. """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``). * **Vector chunks** for RAG-style semantic retrieval (``search``).
* **Ordered pages** for whole-document reading (``read_pages``). * **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 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 pick the strategy that fits the question; they don't need to know which
storage they're hitting. 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: def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
@@ -37,20 +48,27 @@ class DocumentService:
collection: FileId, collection: FileId,
pages: list[PageText], pages: list[PageText],
source: str, source: str,
owner_id: OwnerId,
read_principals: list[PrincipalId],
expires_at: datetime | None,
) -> int: ) -> int:
"""Replace-ingest a document. Returns the number of vector chunks indexed. """Replace-ingest a document. Returns the number of vector chunks indexed.
This wipes any previously-stored content for ``collection`` and writes Wipes any previously-stored content for ``(collection, owner_id)`` and
both the vector-chunk and page-text representations from the same writes both the vector-chunk and page-text representations from the
``pages`` payload. Pages with empty/whitespace-only text are skipped same ``pages`` payload. Pages with empty/whitespace-only text are
for chunking but still written to the page store so page numbering is skipped for chunking but still written to the page store so page
preserved end-to-end. numbering is preserved end-to-end.
""" """
await self._store.delete_collection(collection) if not read_principals:
await self._store.ensure_collection(collection, source) 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] 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] = [] chunks: list[Document] = []
for page in pages: for page in pages:
@@ -70,63 +88,103 @@ class DocumentService:
if not chunks: if not chunks:
return 0 return 0
embeddings = await self._embedder.embed_documents([doc.text for doc in chunks]) 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) return len(chunks)
async def search( async def search(
self, self,
query: str, query: str,
principals: list[PrincipalId],
collection: FileId | None = None, collection: FileId | None = None,
top_k: int | None = None, top_k: int | None = None,
) -> list[SearchResult]: ) -> list[SearchResult]:
"""Embed query and search across one or all collections. """Embed query and search 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 k = top_k if top_k is not None else self._default_top_k
query_embedding = await self._embedder.embed_query(query) query_embedding = await self._embedder.embed_query(query)
if collection is not None: if collection is not None:
if not await self._store.has_collection(collection): if not await self._store.has_collection(collection, principals):
return [] 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) # Search every collection the caller can read, skipping any that error
collections = await self._store.list_collections() # (e.g. dimension mismatch).
collections = await self._store.list_collections(principals)
all_results: list[SearchResult] = [] all_results: list[SearchResult] = []
for col_name in collections: for col_name in collections:
try: 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) all_results.extend(results)
except Exception: # noqa: BLE001 - any backend error on one collection should not stop the others 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) all_results.sort(key=lambda r: r.score, reverse=True)
return all_results[:k] return all_results[:k]
async def read_pages( async def read_pages(
self, self,
collection: FileId, collection: FileId,
principals: list[PrincipalId],
page_range: PageRange | None = None, page_range: PageRange | None = None,
) -> list[Page]: ) -> 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: async def purge_owner(self, owner_id: OwnerId) -> int:
"""Remove a collection's chunks and pages.""" """Remove every collection ``owner_id`` owns, including vector chunks,
await self._store.delete_collection(collection) 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: async def reap_expired(self) -> int:
"""Check whether a collection exists.""" """Delete collections whose ``expires_at`` is set and in the past.
return await self._store.has_collection(collection) 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]: async def has_collection(self, collection: FileId, principals: list[PrincipalId]) -> bool:
"""List all available collections.""" """Check whether at least one principal can read this collection."""
return [FileId(name) for name in await self._store.list_collections()] 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: async def close(self) -> None:
"""Release the underlying store's resources.""" """Release the underlying store's resources."""
+314 -86
View File
@@ -5,19 +5,35 @@ import json
import math import math
import re import re
import sqlite3 import sqlite3
from datetime import UTC, datetime
from pathlib import Path from pathlib import Path
import sqlite_vec import sqlite_vec
from stirling.contracts.documents import Page, PageRange from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage 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): class SqliteVecStore(DocumentStore):
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server. """sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
Each collection gets its own `vec0` virtual table with a fixed embedding dimension Each ``(collection, owner_id)`` pair gets its own `vec0` virtual table with a
(detected on first insert). Document metadata lives in a regular table joined by rowid. 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: def __init__(self, db_path: str | Path) -> None:
@@ -51,17 +67,30 @@ class SqliteVecStore(DocumentStore):
self._conn.execute( self._conn.execute(
""" """
CREATE TABLE IF NOT EXISTS documents_meta ( CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY, collection TEXT NOT NULL,
source 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( self._conn.execute(
""" """
CREATE TABLE IF NOT EXISTS collections ( 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, 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 ( CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL, 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, text TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}', metadata TEXT NOT NULL DEFAULT '{}',
vec_rowid INTEGER NOT NULL, 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( self._conn.execute(
""" """
CREATE TABLE IF NOT EXISTS document_pages ( 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, page_number INTEGER NOT NULL,
text TEXT NOT NULL, text TEXT NOT NULL,
char_count INTEGER 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.execute(
self._conn.commit() "CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
)
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( self._conn.execute(
""" """
INSERT INTO documents_meta(collection, source) VALUES (?, ?) CREATE TABLE IF NOT EXISTS document_acl (
ON CONFLICT(collection) DO UPDATE SET source = excluded.source collection TEXT NOT NULL,
""", owner_id TEXT NOT NULL,
(collection, source), 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() 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 @staticmethod
def _sanitize_table_name(collection: str) -> str: def _sanitize_table_name(collection: str, owner_id: OwnerId) -> str:
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection) safe_col = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
return f"vec_{safe}" safe_owner = re.sub(r"[^a-zA-Z0-9_]", "_", owner_id)
return f"vec_{safe_owner}_{safe_col}"
@staticmethod @staticmethod
def _normalize(vector: list[float]) -> list[float]: def _normalize(vector: list[float]) -> list[float]:
@@ -123,6 +259,7 @@ class SqliteVecStore(DocumentStore):
collection: str, collection: str,
documents: list[Document], documents: list[Document],
embeddings: list[list[float]], embeddings: list[list[float]],
owner_id: OwnerId,
) -> None: ) -> None:
if len(documents) != len(embeddings): if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings") raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
@@ -130,34 +267,40 @@ class SqliteVecStore(DocumentStore):
return return
async with self._lock: 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( def _sync_add(
self, self,
collection: str, collection: str,
documents: list[Document], documents: list[Document],
embeddings: list[list[float]], embeddings: list[list[float]],
owner_id: OwnerId,
) -> None: ) -> None:
dim = len(embeddings[0]) 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: 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(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])")
self._conn.execute( self._conn.execute(
"INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)", "INSERT INTO collections(collection, owner_id, dim, table_name) VALUES (?, ?, ?, ?)",
(collection, dim, table_name), (collection, owner_id, dim, table_name),
) )
else: else:
existing_dim, table_name = row existing_dim, table_name = row
if existing_dim != dim: 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 # Upsert: delete existing docs with matching IDs first
ids = [doc.id for doc in documents] ids = [doc.id for doc in documents]
placeholders = ",".join("?" * len(ids)) placeholders = ",".join("?" * len(ids))
existing = self._conn.execute( existing = self._conn.execute(
f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})", f"SELECT vec_rowid FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, *ids), (collection, owner_id, *ids),
).fetchall() ).fetchall()
if existing: if existing:
vec_rowids = [r[0] for r in existing] vec_rowids = [r[0] for r in existing]
@@ -167,8 +310,8 @@ class SqliteVecStore(DocumentStore):
vec_rowids, vec_rowids,
) )
self._conn.execute( self._conn.execute(
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})", f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, *ids), (collection, owner_id, *ids),
) )
for doc, emb in zip(documents, embeddings): for doc, emb in zip(documents, embeddings):
@@ -179,27 +322,101 @@ class SqliteVecStore(DocumentStore):
) )
vec_rowid = cursor.lastrowid vec_rowid = cursor.lastrowid
self._conn.execute( self._conn.execute(
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)", "INSERT INTO documents(id, collection, owner_id, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?, ?)",
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid), (doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), vec_rowid),
) )
self._conn.commit() 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( async def search(
self, self,
collection: str, collection: str,
query_embedding: list[float], query_embedding: list[float],
top_k: int = 5, top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]: ) -> list[SearchResult]:
async with self._lock: 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( def _sync_search(
self, self,
collection: str, collection: str,
query_embedding: list[float], query_embedding: list[float],
top_k: int, top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]: ) -> 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: if row is None:
return [] return []
table_name, dim = row table_name, dim = row
@@ -213,11 +430,14 @@ class SqliteVecStore(DocumentStore):
f""" f"""
SELECT d.id, d.text, d.metadata, v.distance SELECT d.id, d.text, d.metadata, v.distance
FROM {table_name} v 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 = ? WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance ORDER BY v.distance
""", """,
(collection, query_blob, top_k), (collection, owner_id, query_blob, top_k),
).fetchall() ).fetchall()
return [ return [
@@ -233,77 +453,85 @@ class SqliteVecStore(DocumentStore):
for r in results for r in results
] ]
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: def _readable_owner_for(self, collection: str, principals: list[PrincipalId]) -> str | None:
async with self._lock: """Resolve which owner_id this caller is reading. ``None`` means no access."""
await asyncio.to_thread(self._sync_add_pages, collection, pages) placeholders = ",".join("?" * len(principals))
row = self._conn.execute(
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None: f"""
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,)) SELECT owner_id FROM document_acl
if pages: WHERE collection = ?
self._conn.executemany( AND permission = ?
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)", AND principal_id IN ({placeholders})
[(collection, p.page_number, p.text, p.char_count) for p in pages], ORDER BY owner_id
) LIMIT 1
self._conn.commit() """,
(collection, _READ_PERMISSION, *principals),
).fetchone()
return row[0] if row else None
async def read_pages( async def read_pages(
self, self,
collection: str, collection: str,
page_range: PageRange | None = None, page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]: ) -> list[Page]:
async with self._lock: 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( def _sync_read_pages(
self, self,
collection: str, collection: str,
page_range: PageRange | None, page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]: ) -> 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: if page_range is None:
rows = self._conn.execute( rows = self._conn.execute(
"SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number", "SELECT page_number, text, char_count FROM document_pages "
(collection,), "WHERE collection = ? AND owner_id = ? ORDER BY page_number",
(collection, owner_id),
).fetchall() ).fetchall()
else: else:
rows = self._conn.execute( rows = self._conn.execute(
"SELECT page_number, text, char_count FROM document_pages " "SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number", "WHERE collection = ? AND owner_id = ? AND page_number BETWEEN ? AND ? "
(collection, page_range.start, page_range.end), "ORDER BY page_number",
(collection, owner_id, page_range.start, page_range.end),
).fetchall() ).fetchall()
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows] 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: 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: def _sync_has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables if not principals:
# (collections, documents, document_pages) when documents_meta is deleted. return False
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone() return self._readable_owner_for(collection, principals) is not None
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()
async def list_collections(self) -> list[str]: async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
async with self._lock: 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]: def _sync_list_collections(self, principals: list[PrincipalId]) -> list[str]:
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall() 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] 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 def close(self) -> None:
async with self._lock: async with self._lock:
await asyncio.to_thread(self._sync_close) await asyncio.to_thread(self._sync_close)
+75 -33
View File
@@ -2,8 +2,10 @@ from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime
from stirling.contracts.documents import Page, PageRange from stirling.contracts.documents import Page, PageRange
from stirling.models import OwnerId, PrincipalId
@dataclass @dataclass
@@ -40,73 +42,113 @@ class DocumentStore(ABC):
* **Vector chunks** - small, embedded chunks used for RAG search. * **Vector chunks** - small, embedded chunks used for RAG search.
* **Ordered pages** - the original page text retained in document order, * **Ordered pages** - the original page text retained in document order,
used for whole-document reading. 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 @abstractmethod
async def ensure_collection(self, collection: str, source: str) -> None: async def ensure_collection(
"""Upsert the top-level ``documents_meta`` row for this 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 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 write into child tables that hold a foreign key to the parent row, so
row, so it must exist first. 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 @abstractmethod
async def add_documents( async def add_documents(
self, self,
collection: str, collection: str,
documents: list[Document], documents: list[Document],
embeddings: list[list[float]], embeddings: list[list[float]],
owner_id: OwnerId,
) -> None: ) -> 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 @abstractmethod
async def search( async def search(
self, self,
collection: str, collection: str,
query_embedding: list[float], query_embedding: list[float],
top_k: int = 5, top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]: ) -> 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 Returns an empty list when no principal in ``principals`` has a read
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None: ACL row for this collection (regardless of which owner backs it).
"""Replace the stored pages for ``collection`` with the supplied pages.
Implementations must remove any previously-stored pages for the
collection before writing, so callers can re-ingest by calling this
method again.
""" """
@abstractmethod @abstractmethod
async def read_pages( async def read_pages(
self, self,
collection: str, collection: str,
page_range: PageRange | None = None, page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]: ) -> list[Page]:
"""Return ordered pages for ``collection``. """Return ordered pages for ``collection``. Empty list when no read ACL match."""
If ``page_range`` is ``None`` all pages are returned. Otherwise only
pages whose ``page_number`` falls within the inclusive range are
returned. Pages are always ordered by ``page_number`` ascending.
"""
@abstractmethod @abstractmethod
async def delete_collection(self, collection: str) -> None: async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
"""Remove a collection's chunks and pages.""" """Check whether any principal in ``principals`` can read this collection."""
@abstractmethod @abstractmethod
async def list_collections(self) -> list[str]: async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
"""Return names of all existing collections.""" """Return collection names readable by at least one of ``principals``."""
@abstractmethod # ── lifecycle ──────────────────────────────────────────────────────────
async def has_collection(self, collection: str) -> bool:
"""Check whether a collection exists."""
@abstractmethod @abstractmethod
async def close(self) -> None: async def close(self) -> None:
+4 -1
View File
@@ -1,12 +1,15 @@
from . import tool_models from . import tool_models
from .base import ApiModel, FileId from .base import ApiModel, FileId, OwnerId, PrincipalId, UserId
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
__all__ = [ __all__ = [
"ApiModel", "ApiModel",
"FileId", "FileId",
"OPERATIONS", "OPERATIONS",
"OwnerId",
"ParamToolModel", "ParamToolModel",
"PrincipalId",
"ToolEndpoint", "ToolEndpoint",
"UserId",
"tool_models", "tool_models",
] ]
+10 -3
View File
@@ -5,11 +5,18 @@ from typing import NewType
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel from pydantic.alias_generators import to_camel
# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's # Stable, opaque identifier for a file supplied by the caller
# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key
# throughout the engine.
FileId = NewType("FileId", str) 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): class ApiModel(BaseModel):
"""Base for every contract model crossing a service boundary.""" """Base for every contract model crossing a service boundary."""
+3 -1
View File
@@ -7,14 +7,16 @@ from .progress import (
set_progress_emitter, set_progress_emitter,
) )
from .runtime import AppRuntime, build_model_settings, build_runtime from .runtime import AppRuntime, build_model_settings, build_runtime
from .tracking import setup_posthog_tracking from .tracking import current_user_id, require_current_user_id, setup_posthog_tracking
__all__ = [ __all__ = [
"AppRuntime", "AppRuntime",
"ProgressEmitter", "ProgressEmitter",
"build_model_settings", "build_model_settings",
"build_runtime", "build_runtime",
"current_user_id",
"emit_progress", "emit_progress",
"require_current_user_id",
"reset_progress_emitter", "reset_progress_emitter",
"set_progress_emitter", "set_progress_emitter",
"setup_posthog_tracking", "setup_posthog_tracking",
+4 -11
View File
@@ -16,7 +16,6 @@ from stirling.documents import (
DocumentStore, DocumentStore,
EmbeddingService, EmbeddingService,
PgVectorStore, PgVectorStore,
RagCapability,
SqliteVecStore, SqliteVecStore,
) )
@@ -47,7 +46,6 @@ class AppRuntime:
fast_model: Model fast_model: Model
smart_model: Model smart_model: Model
documents: DocumentService documents: DocumentService
rag_capability: RagCapability
@property @property
def fast_model_settings(self) -> ModelSettings: def fast_model_settings(self) -> ModelSettings:
@@ -89,8 +87,8 @@ def _build_document_store(settings: AppSettings) -> DocumentStore:
assert_never(settings.rag_backend) assert_never(settings.rag_backend)
def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]: def _build_documents(settings: AppSettings) -> DocumentService:
"""Build the document service and the RAG-search capability that wraps it.""" """Build the document service used by per-request RAG capabilities."""
logger.info("Documents: embedding_model=%s", settings.rag_embedding_model) logger.info("Documents: embedding_model=%s", settings.rag_embedding_model)
embedder = EmbeddingService( embedder = EmbeddingService(
model_name=settings.rag_embedding_model, model_name=settings.rag_embedding_model,
@@ -98,9 +96,7 @@ def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapabil
chunk_overlap=settings.rag_chunk_overlap, chunk_overlap=settings.rag_chunk_overlap,
) )
store = _build_document_store(settings) store = _build_document_store(settings)
service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k) return 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
def build_runtime(settings: AppSettings) -> AppRuntime: 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(fast_model, settings.fast_model_name)
validate_structured_output_support(smart_model, settings.smart_model_name) validate_structured_output_support(smart_model, settings.smart_model_name)
documents, rag_capability = _build_documents(settings)
return AppRuntime( return AppRuntime(
settings=settings, settings=settings,
fast_model=fast_model, fast_model=fast_model,
smart_model=smart_model, smart_model=smart_model,
documents=documents, documents=_build_documents(settings),
rag_capability=rag_capability,
) )
+16 -1
View File
@@ -27,10 +27,25 @@ from opentelemetry.trace import Span
from posthog.client import Client as PostHogClient from posthog.client import Client as PostHogClient
from stirling.config import AppSettings from stirling.config import AppSettings
from stirling.models import UserId
# Per-request user ID, set by middleware from the X-User-Id header. # 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. # 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: class LRUSet:
+22 -2
View File
@@ -15,9 +15,12 @@ import pytest
from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability from stirling.agents.shared import ChunkedReasoner, ChunkNotes, WholeDocReaderCapability
from stirling.contracts import AiFile, PageText from stirling.contracts import AiFile, PageText
from stirling.documents import Document, DocumentService, SqliteVecStore 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 from stirling.services.runtime import AppRuntime
OWNER = OwnerId("test-user")
PRINCIPALS = [PrincipalId("test-user")]
class StubEmbedder: class StubEmbedder:
"""Deterministic embeddings so tests don't need a real provider.""" """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=1, text="Chapter one prose."),
PageText(page_number=2, text="Chapter two 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) reasoner = ChunkedReasoner(runtime_with_stub_docs)
canned_notes = [ChunkNotes(pages=[1, 2], summary="overview", facts=["fact-A"])] 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( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("doc-id", "doc.pdf")], files=[_ai_file("doc-id", "doc.pdf")],
principals=PRINCIPALS,
reasoner=reasoner, reasoner=reasoner,
) )
result = await capability._read_full_document("what is in the document?") 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), FileId(cid),
[PageText(page_number=1, text=f"contents of {cid}")], [PageText(page_number=1, text=f"contents of {cid}")],
source=source, source=source,
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
) )
reasoner = ChunkedReasoner(runtime_with_stub_docs) 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( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")], files=[_ai_file("doc-a", "a.pdf"), _ai_file("doc-b", "b.pdf")],
principals=PRINCIPALS,
reasoner=reasoner, reasoner=reasoner,
) )
result = await capability._read_full_document("compare them") 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"), FileId("present"),
[PageText(page_number=1, text="real content")], [PageText(page_number=1, text="real content")],
source="present.pdf", source="present.pdf",
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
) )
# 'missing' is never ingested -> read_pages returns []. # '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( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")], files=[_ai_file("missing", "missing.pdf"), _ai_file("present", "present.pdf")],
principals=PRINCIPALS,
reasoner=reasoner, reasoner=reasoner,
) )
result = await capability._read_full_document("anything") 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( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("nope", "nope.pdf")], files=[_ai_file("nope", "nope.pdf")],
principals=PRINCIPALS,
reasoner=reasoner, reasoner=reasoner,
) )
result = await capability._read_full_document("anything") 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"), FileId("doc-id"),
[PageText(page_number=1, text="content")], [PageText(page_number=1, text="content")],
source="doc.pdf", source="doc.pdf",
owner_id=OWNER,
read_principals=PRINCIPALS,
expires_at=None,
) )
reasoner = ChunkedReasoner(runtime_with_stub_docs) reasoner = ChunkedReasoner(runtime_with_stub_docs)
with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])): with patch.object(reasoner, "gather_notes", AsyncMock(return_value=[ChunkNotes(pages=[1], summary="s")])):
capability = WholeDocReaderCapability( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("doc-id", "doc.pdf")], files=[_ai_file("doc-id", "doc.pdf")],
principals=PRINCIPALS,
reasoner=reasoner, reasoner=reasoner,
max_reads=1, max_reads=1,
) )
@@ -212,6 +231,7 @@ async def test_instructions_mention_attached_files(runtime_with_stub_docs: AppRu
capability = WholeDocReaderCapability( capability = WholeDocReaderCapability(
runtime=runtime_with_stub_docs, runtime=runtime_with_stub_docs,
files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")], files=[_ai_file("doc-a", "alpha.pdf"), _ai_file("doc-b", "beta.pdf")],
principals=PRINCIPALS,
) )
text = capability.instructions text = capability.instructions
@@ -17,9 +17,11 @@ from stirling.contracts.contradiction import (
ContradictionReport, ContradictionReport,
ContradictionSeverity, ContradictionSeverity,
) )
from stirling.models import FileId from stirling.models import FileId, PrincipalId
from stirling.services.runtime import AppRuntime from stirling.services.runtime import AppRuntime
PRINCIPALS = [PrincipalId("test-user")]
def _file(file_id: str, name: str) -> AiFile: def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name) 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() canned = _canned_report()
detector.detect = AsyncMock(return_value=canned) 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?") result = await capability._find_contradictions("are there inconsistent deadlines?")
detector.detect.assert_awaited_once() detector.detect.assert_awaited_once()
@@ -84,6 +86,7 @@ async def test_budget_gate_hides_tool_after_first_audit(runtime: AppRuntime) ->
capability = ContradictionCapability( capability = ContradictionCapability(
detector=detector, detector=detector,
files=[_file("doc-a", "a.pdf")], files=[_file("doc-a", "a.pdf")],
principals=PRINCIPALS,
max_audits=1, max_audits=1,
) )
# A real, minimal ToolDefinition — the prepare callback returns this # 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: async def test_find_contradictions_with_no_files_returns_message(runtime: AppRuntime) -> None:
detector = ContradictionDetector(runtime) detector = ContradictionDetector(runtime)
detector.detect = AsyncMock(return_value=_canned_report()) 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") result = await capability._find_contradictions("anything")
@@ -120,6 +123,7 @@ def test_instructions_mention_attached_files(runtime: AppRuntime) -> None:
capability = ContradictionCapability( capability = ContradictionCapability(
detector=detector, detector=detector,
files=[_file("doc-a", "alpha.pdf"), _file("doc-b", "beta.pdf")], files=[_file("doc-a", "alpha.pdf"), _file("doc-b", "beta.pdf")],
principals=PRINCIPALS,
) )
text = capability.instructions text = capability.instructions
@@ -149,6 +153,7 @@ def test_instructions_escape_filename_injection_attempt(runtime: AppRuntime) ->
capability = ContradictionCapability( capability = ContradictionCapability(
detector=detector, detector=detector,
files=[_file("doc-evil", evil_name)], files=[_file("doc-evil", evil_name)],
principals=PRINCIPALS,
) )
text = capability.instructions text = capability.instructions
+24 -17
View File
@@ -27,7 +27,7 @@ from stirling.agents.shared.chunked_mapper import ChunkOutput
from stirling.contracts import AiFile from stirling.contracts import AiFile
from stirling.contracts.contradiction import ContradictionSeverity from stirling.contracts.contradiction import ContradictionSeverity
from stirling.contracts.documents import Page, PageRange from stirling.contracts.documents import Page, PageRange
from stirling.models import FileId from stirling.models import FileId, PrincipalId
from stirling.services.runtime import AppRuntime 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: 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.""" """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, []) return pages_by_id.get(collection, [])
# AppRuntime is frozen; monkey-patch the documents service. # 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: []}) _install_documents_stub(runtime, {file_a.id: []})
detector = ContradictionDetector(runtime) detector = ContradictionDetector(runtime)
report = await detector.detect([file_a]) report = await detector.detect([file_a], principals=PRINCIPALS)
assert report.contradictions == [] assert report.contradictions == []
assert report.pages_examined == [] 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.")) 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 assert len(report.contradictions) == 1
c = report.contradictions[0] 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")] return_value=[ChunkOutput(pages=[1, 2], output=_ExtractedClaims(claims=[]), label="pages=1-2")]
) )
# Stubbing the summary agent is unavoidable (the production code calls # Stubbing the summary agent is unavoidable (the production code calls
# it on every detect()); we just don't assert on what it returns # 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 # Asserting on the canned value here would only re-prove that AsyncMock
# works. # works.
detector._summary_agent.run = AsyncMock(return_value=_stub_result("any text")) 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.contradictions == []
assert report.clean is True 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")) 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 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")) 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. # Lexical key collapses both subjects so the bucket still forms.
assert len(report.contradictions) == 1 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")) 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 len(report.contradictions) == 1
assert report.contradictions[0].severity == ContradictionSeverity.ERROR 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")) 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 == [] 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) 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 "No contradictions" in report.summary
assert report.clean is True 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._pair_detector.run = AsyncMock(side_effect=TimeoutError("simulated"))
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) 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 # Detector timed out so no pairs come back. Crucially: the pipeline
# reached the summary stage rather than crashing earlier, so # 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")) detector._summary_agent.run = AsyncMock(return_value=_stub_result("ok"))
with caplog.at_level(logging.WARNING, logger="stirling.agents.contradiction.detector"): 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( assert any(
"produced 0 claims" in record.getMessage() and "pages=1" in record.getMessage() for record in caplog.records "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._pair_detector.run = AsyncMock(return_value=_stub_result(_BucketContradictions(pairs=[])))
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) 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 # Every page the extractor ran against is reported, even page 2
# (which produced no claim). # (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._pair_detector.run = _stub_detector # type: ignore[method-assign]
detector._summary_agent.run = AsyncMock(return_value=_stub_result("done")) 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) # Both windows produced one valid pair each; dedup by global (i, j)
# leaves exactly two contradictions. # 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")) 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 # Both claims validated as verbatim — each against the right file's
# page text. A collision bug would have produced "paraphrased" for at # page text. A collision bug would have produced "paraphrased" for at
@@ -9,6 +9,7 @@ detector when invoked.
from __future__ import annotations from __future__ import annotations
from collections.abc import Iterator
from dataclasses import replace from dataclasses import replace
import pytest import pytest
@@ -22,10 +23,24 @@ from stirling.contracts import (
) )
from stirling.contracts.contradiction import Claim from stirling.contracts.contradiction import Claim
from stirling.documents import DocumentService, SqliteVecStore 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 stirling.services.runtime import AppRuntime
from tests.test_pdf_question_agent import StubEmbedder 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: def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name) 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, file.id,
[PageText(page_number=1, text="content")], [PageText(page_number=1, text="content")],
source=file.name, source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = PdfQuestionAgent(runtime_with_stub_docs) agent = PdfQuestionAgent(runtime_with_stub_docs)
@@ -8,6 +8,7 @@ contradiction and the right cross-references and anchor handling.
from __future__ import annotations from __future__ import annotations
import json import json
from collections.abc import Iterator
from dataclasses import replace from dataclasses import replace
from typing import Literal from typing import Literal
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
@@ -27,11 +28,25 @@ from stirling.contracts import (
) )
from stirling.contracts.contradiction import Claim from stirling.contracts.contradiction import Claim
from stirling.documents import DocumentService, SqliteVecStore 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.models.tool_models import AddCommentsParams
from stirling.services import current_user_id
from stirling.services.runtime import AppRuntime from stirling.services.runtime import AppRuntime
from tests.test_pdf_question_agent import StubEmbedder 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: def _file(file_id: str, name: str) -> AiFile:
return AiFile(id=FileId(file_id), name=name) return AiFile(id=FileId(file_id), name=name)
@@ -88,6 +103,9 @@ async def test_localiser_prompt_escapes_verdict_tag_injection(
file.id, file.id,
[PageText(page_number=1, text="x")], [PageText(page_number=1, text="x")],
source=file.name, source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = PdfReviewAgent(runtime_with_stub_docs) agent = PdfReviewAgent(runtime_with_stub_docs)
@@ -159,6 +177,9 @@ async def test_contradiction_intent_emits_add_comments_plan(
file.id, file.id,
[PageText(page_number=1, text="ignored"), PageText(page_number=5, text="ignored")], [PageText(page_number=1, text="ignored"), PageText(page_number=5, text="ignored")],
source=file.name, source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = PdfReviewAgent(runtime_with_stub_docs) agent = PdfReviewAgent(runtime_with_stub_docs)
@@ -264,6 +285,9 @@ async def test_contradiction_takes_precedence_over_math(
file.id, file.id,
[PageText(page_number=1, text="x")], [PageText(page_number=1, text="x")],
source=file.name, source=file.name,
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = PdfReviewAgent(runtime_with_stub_docs) agent = PdfReviewAgent(runtime_with_stub_docs)
+343 -53
View File
@@ -8,7 +8,17 @@ from stirling.documents.rag_capability import RagCapability
from stirling.documents.service import DocumentService from stirling.documents.service import DocumentService
from stirling.documents.sqlite_vec_store import SqliteVecStore from stirling.documents.sqlite_vec_store import SqliteVecStore
from stirling.documents.store import Document, SearchResult 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 # chunk_text
@@ -56,7 +66,8 @@ class TestSqliteVecStore:
@pytest.mark.anyio @pytest.mark.anyio
async def test_add_and_search(self) -> None: async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral() store = SqliteVecStore.ephemeral()
await store.ensure_collection("test-col", "test.pdf") await store.ensure_collection("test-col", "test.pdf", OWNER, None)
await store.grant_read("test-col", OWNER, OWNER_PRINCIPALS)
docs = [ docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}), Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}), Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
@@ -67,9 +78,9 @@ class TestSqliteVecStore:
[0.9, 0.1, 0.0], [0.9, 0.1, 0.0],
[0.0, 0.0, 1.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 len(results) == 2
assert isinstance(results[0], SearchResult) assert isinstance(results[0], SearchResult)
assert results[0].document.id == "1" assert results[0].document.id == "1"
@@ -78,33 +89,36 @@ class TestSqliteVecStore:
@pytest.mark.anyio @pytest.mark.anyio
async def test_list_and_has_collection(self) -> None: async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral() store = SqliteVecStore.ephemeral()
await store.ensure_collection("my-collection", "test.pdf") 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={})] 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 "my-collection" in collections
assert await store.has_collection("my-collection") is True assert await store.has_collection("my-collection", OWNER_PRINCIPALS) is True
assert await store.has_collection("nonexistent") is False assert await store.has_collection("nonexistent", OWNER_PRINCIPALS) is False
@pytest.mark.anyio @pytest.mark.anyio
async def test_delete_collection(self) -> None: async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral() store = SqliteVecStore.ephemeral()
await store.ensure_collection("to-delete", "test.pdf") 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={})] 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 assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is True
await store.delete_collection("to-delete") await store.delete_collection("to-delete", OWNER)
assert await store.has_collection("to-delete") is False assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is False
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_empty_collection(self) -> None: async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral() store = SqliteVecStore.ephemeral()
await store.ensure_collection("empty-test", "test.pdf") 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={})] docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]]) await store.add_documents("empty-test", docs, [[1.0, 0.0]], OWNER)
results = await store.search("empty-test", [1.0, 0.0], top_k=5) results = await store.search("empty-test", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS)
assert len(results) == 1 assert len(results) == 1
@pytest.mark.anyio @pytest.mark.anyio
@@ -112,7 +126,118 @@ class TestSqliteVecStore:
store = SqliteVecStore.ephemeral() store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})] docs = [Document(id="1", text="test", metadata={})]
with pytest.raises(ValueError, match="documents.*embeddings"): 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) # DocumentService (with stub embedder)
@@ -163,39 +288,131 @@ class TestDocumentService:
@pytest.mark.anyio @pytest.mark.anyio
async def test_ingest_and_search(self, documents: DocumentService) -> None: async def test_ingest_and_search(self, documents: DocumentService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy." text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await 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 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 len(results) > 0
assert results[0].document.text assert results[0].document.text
@pytest.mark.anyio @pytest.mark.anyio
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None: 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 assert count == 0
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None: 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 == [] assert results == []
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_all_collections(self, documents: DocumentService) -> None: 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") await documents.ingest(
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf") 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 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 @pytest.mark.anyio
async def test_delete_collection(self, documents: DocumentService) -> None: async def test_delete_collection(self, documents: DocumentService) -> None:
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf") await documents.ingest(
collections = await documents.list_collections() 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 assert "temp" in collections
await documents.delete_collection(FileId("temp")) await documents.delete_collection(FileId("temp"), owner_id=OWNER)
collections = await documents.list_collections() collections = await documents.list_collections(OWNER_PRINCIPALS)
assert "temp" not in collections assert "temp" not in collections
@pytest.mark.anyio @pytest.mark.anyio
@@ -205,9 +422,16 @@ class TestDocumentService:
PageText(page_number=2, text="Second page text."), PageText(page_number=2, text="Second page text."),
PageText(page_number=3, text="Third 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 [p.page_number for p in stored] == [1, 2, 3]
assert stored[0].text == "First page text." assert stored[0].text == "First page text."
assert stored[0].char_count == len("First page text.") assert stored[0].char_count == len("First page text.")
@@ -217,9 +441,13 @@ class TestDocumentService:
from stirling.contracts import PageRange from stirling.contracts import PageRange
pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)] 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] assert [p.page_number for p in subset] == [2, 3, 4]
@pytest.mark.anyio @pytest.mark.anyio
@@ -228,14 +456,20 @@ class TestDocumentService:
FileId("doc"), FileId("doc"),
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")], [PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
source="v1.pdf", source="v1.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
await documents.ingest( await documents.ingest(
FileId("doc"), FileId("doc"),
[PageText(page_number=1, text="new")], [PageText(page_number=1, text="new")],
source="v2.pdf", 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 [p.page_number for p in stored] == [1]
assert stored[0].text == "new" assert stored[0].text == "new"
@@ -248,9 +482,16 @@ class TestDocumentService:
PageText(page_number=2, text=" "), PageText(page_number=2, text=" "),
PageText(page_number=3, text="Real text on page 3."), 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 [p.page_number for p in stored] == [1, 2, 3]
assert stored[1].text.strip() == "" assert stored[1].text.strip() == ""
@@ -270,22 +511,36 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
class TestRagCapability: class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None: 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 instructions = cap.instructions
assert isinstance(instructions, str) assert isinstance(instructions, str)
assert "docs, manuals" in instructions assert "docs, manuals" in instructions
assert "search_knowledge" in instructions assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None: def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
cap = RagCapability(documents) cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions = cap.instructions instructions = cap.instructions
assert callable(instructions) assert callable(instructions)
@pytest.mark.anyio @pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None: 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(
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf") FileId("col-a"),
cap = RagCapability(documents) _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 instructions_fn = cap.instructions
assert callable(instructions_fn) assert callable(instructions_fn)
text = await instructions_fn() text = await instructions_fn()
@@ -294,7 +549,7 @@ class TestRagCapability:
@pytest.mark.anyio @pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None: 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 instructions_fn = cap.instructions
assert callable(instructions_fn) assert callable(instructions_fn)
text = await instructions_fn() text = await instructions_fn()
@@ -302,14 +557,21 @@ class TestRagCapability:
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None: 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") output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base." assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None: 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") await documents.ingest(
cap = RagCapability(documents) 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") output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output assert "[Result 1" in output
assert "source: guide.pdf" in output assert "source: guide.pdf" in output
@@ -318,10 +580,24 @@ class TestRagCapability:
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None: 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(
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf") 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") output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output assert "pinned.pdf" in output
assert "other.pdf" not in output assert "other.pdf" not in output
@@ -329,9 +605,16 @@ class TestRagCapability:
@pytest.mark.anyio @pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None: async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10)) paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await 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) output = await _invoke_search_knowledge(cap, "topic", max_results=2)
assert "[Result 1" in output assert "[Result 1" in output
assert "[Result 2" 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: async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None:
"""The prepare callback must return None once max_searches has been reached """The prepare callback must return None once max_searches has been reached
so the agent can no longer call the tool on subsequent turns.""" so the agent can no longer call the tool on subsequent turns."""
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf") await documents.ingest(
cap = RagCapability(documents, max_searches=2) 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() tool_def = _dummy_tool_def()
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type] assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]
+174 -9
View File
@@ -8,7 +8,11 @@ from fastapi.testclient import TestClient
from stirling.api import app from stirling.api import app
from stirling.api.dependencies import get_document_service from stirling.api.dependencies import get_document_service
from stirling.documents import Document, DocumentService, SqliteVecStore 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: 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": 1, "text": "The introduction covers the main topic."},
{"pageNumber": 2, "text": "The conclusion summarises the findings."}, {"pageNumber": 2, "text": "The conclusion summarises the findings."},
], ],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
}, },
headers=HEADERS,
) )
assert response.status_code == 200 assert response.status_code == 200
body = response.json() body = response.json()
@@ -94,7 +102,11 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
"documentId": "replace-me", "documentId": "replace-me",
"source": "replace-me.pdf", "source": "replace-me.pdf",
"pageText": [{"pageNumber": 1, "text": "Original content that existed before."}], "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 # Second ingest with different content should replace the first entirely
response = client.post( response = client.post(
@@ -103,11 +115,15 @@ async def test_ingest_document_replaces_existing_content(client: TestClient, ser
"documentId": "replace-me", "documentId": "replace-me",
"source": "replace-me.pdf", "source": "replace-me.pdf",
"pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}], "pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
}, },
headers=HEADERS,
) )
assert response.status_code == 200 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] texts = [r.document.text for r in results]
assert any("New content" in t for t in texts) assert any("New content" in t for t in texts)
assert not any("Original 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": 1, "text": " "},
{"pageNumber": 2, "text": "Real content on page 2."}, {"pageNumber": 2, "text": "Real content on page 2."},
], ],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
}, },
headers=HEADERS,
) )
assert response.status_code == 200 assert response.status_code == 200
assert response.json()["chunksIndexed"] >= 1 assert response.json()["chunksIndexed"] >= 1
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None: def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
response = client.post("/api/v1/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.status_code == 200
assert response.json()["chunksIndexed"] == 0 assert response.json()["chunksIndexed"] == 0
@@ -139,6 +169,7 @@ def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
response = client.post( response = client.post(
"/api/v1/documents", "/api/v1/documents",
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]}, json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
) )
assert response.status_code == 422 assert response.status_code == 422
@@ -147,6 +178,7 @@ def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
response = client.post( response = client.post(
"/api/v1/documents", "/api/v1/documents",
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]}, json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
) )
assert response.status_code == 422 assert response.status_code == 422
@@ -155,6 +187,7 @@ def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
response = client.post( response = client.post(
"/api/v1/documents", "/api/v1/documents",
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]}, json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
headers=HEADERS,
) )
assert response.status_code == 422 assert response.status_code == 422
@@ -166,11 +199,61 @@ def test_ingest_document_rejects_non_positive_page_number(client: TestClient) ->
"documentId": "bad-page", "documentId": "bad-page",
"source": "bad-page.pdf", "source": "bad-page.pdf",
"pageText": [{"pageNumber": 0, "text": "something"}], "pageText": [{"pageNumber": 0, "text": "something"}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
}, },
headers=HEADERS,
) )
assert response.status_code == 422 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} ────────────────────────────────────────────── # ── DELETE /documents/{id} ──────────────────────────────────────────────
@@ -181,15 +264,19 @@ def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -
"documentId": "to-delete", "documentId": "to-delete",
"source": "to-delete.pdf", "source": "to-delete.pdf",
"pageText": [{"pageNumber": 1, "text": "Text."}], "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.status_code == 200
assert response.json() == {"documentId": "to-delete", "deleted": True} assert response.json() == {"documentId": "to-delete", "deleted": True}
def test_delete_document_is_idempotent(client: TestClient) -> None: def test_delete_document_is_idempotent(client: TestClient) -> None:
response = client.delete("/api/v1/documents/never-existed") response = client.delete("/api/v1/documents/by-id/never-existed", headers=HEADERS)
assert response.status_code == 200 assert response.status_code == 200
assert response.json() == {"documentId": "never-existed", "deleted": False} assert response.json() == {"documentId": "never-existed", "deleted": False}
@@ -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: async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
client.post( client.post(
"/api/v1/documents", "/api/v1/documents",
json={"documentId": "gone", "source": "gone.pdf", "pageText": [{"pageNumber": 1, "text": "Text."}]}, json={
"documentId": "gone",
"source": "gone.pdf",
"pageText": [{"pageNumber": 1, "text": "Text."}],
"ownerId": USER,
"readPrincipals": [USER],
"expiresAt": None,
},
headers=HEADERS,
) )
assert await service.has_collection(FileId("gone")) assert await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS)
client.delete("/api/v1/documents/gone") client.delete("/api/v1/documents/by-id/gone", headers=HEADERS)
assert not await service.has_collection(FileId("gone")) 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}
+27 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Iterator
from dataclasses import replace from dataclasses import replace
import pytest import pytest
@@ -19,9 +20,25 @@ from stirling.contracts import (
SupportedCapability, SupportedCapability,
) )
from stirling.documents import Document, DocumentService, SqliteVecStore 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 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: class StubEmbedder:
"""Deterministic embeddings so RAG lookups work in tests without network.""" """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"), FileId("present-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")], [PageText(page_number=1, text="Invoice total: 120.00.")],
source="present.pdf", source="present.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = PdfQuestionAgent(runtime_with_stub_rag) 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"), FileId("invoice-id"),
[PageText(page_number=1, text="Invoice total: 120.00.")], [PageText(page_number=1, text="Invoice total: 120.00.")],
source="invoice.pdf", source="invoice.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = StubPdfQuestionAgent( agent = StubPdfQuestionAgent(
runtime_with_stub_rag, 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"), FileId("shipping-id"),
[PageText(page_number=1, text="This page contains only a shipping address.")], [PageText(page_number=1, text="This page contains only a shipping address.")],
source="shipping.pdf", source="shipping.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
) )
agent = StubPdfQuestionAgent( agent = StubPdfQuestionAgent(
runtime_with_stub_rag, runtime_with_stub_rag,