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

Set up document management for Stirling Engine (#6476)

Browse Source 
# Description of Changes
Change Stirling Engine to support deleting documents automatically. This
happens both on user logout and after an amount of time specified by the
Java when ingesting a document (allowing for personal documents to have
short lifetimes but org documents to be left in the db with no expiry
date). Also sets up an [ACL
policy](https://en.wikipedia.org/wiki/Access-control_list) for the
documents so the database knows which users have access to which
documents. This is not fully implemented in the Java, so currently all
docs are treated as having a single owner, the uploader, but
theoretically when we need to support org storage, we shouldn't need to
change the db schema.
This commit is contained in:
James Brunton
2026-06-03 11:52:11 +00:00
committed by GitHub
parent 71633861d0
commit 1264f4cfed
48 changed files with 2039 additions and 411 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/stirling/agents/contradiction/capability.py
+4 -1
View File
@@ -21,6 +21,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.contradiction.detector import ContradictionDetector
from stirling.contracts import AiFile
from stirling.contracts.contradiction import Claim, ContradictionReport
from stirling.models import PrincipalId
logger = logging.getLogger(__name__)
@@ -52,6 +53,7 @@ class ContradictionCapability:
self,
detector: ContradictionDetector,
files: list[AiFile],
principals: list[PrincipalId],
*,
max_audits: int = DEFAULT_MAX_AUDITS,
) -> None:
@@ -59,6 +61,7 @@ class ContradictionCapability:
raise ValueError("max_audits must be >= 1")
self._detector = detector
self._files = files
self._principals = principals
self._max_audits = max_audits
self._audit_count = 0
toolset: FunctionToolset[None] = FunctionToolset()
@@ -121,7 +124,7 @@ class ContradictionCapability:
if not self._files:
return "No documents attached to audit."
report = await self._detector.detect(self._files, query=query)
report = await self._detector.detect(self._files, principals=self._principals, query=query)
formatted = self.format_report(report)
logger.info(
"[contradiction-capability] audit query=%r files=%d -> %d findings, %d chars",
engine/src/stirling/agents/contradiction/detector.py
+11 -4
View File
@@ -41,6 +41,7 @@ from stirling.contracts.contradiction import (
ContradictionSeverity,
)
from stirling.contracts.documents import Page
from stirling.models import PrincipalId
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
@@ -205,8 +206,13 @@ class ContradictionDetector:
# Public entry point
# ------------------------------------------------------------------
async def detect(self, files: list[AiFile], query: str | None = None) -> ContradictionReport:
"""Run the full pipeline over the supplied files.
async def detect(
self,
files: list[AiFile],
principals: list[PrincipalId],
query: str | None = None,
) -> ContradictionReport:
"""Run the full pipeline over the supplied files for ``principals``.
``files`` must have already been ingested (the caller is
responsible for the ``has_collection`` precheck — the question
@@ -234,7 +240,7 @@ class ContradictionDetector:
effective_query = query or "extract claims"
per_file_results = await asyncio.gather(
*(self._extract_claims_for_file(file, effective_query) for file in files),
*(self._extract_claims_for_file(file, effective_query, principals) for file in files),
return_exceptions=True,
)
@@ -312,6 +318,7 @@ class ContradictionDetector:
self,
file: AiFile,
query: str,
principals: list[PrincipalId],
) -> _FileExtractionResult:
"""Run the per-chunk extractor over one file's pages.
@@ -324,7 +331,7 @@ class ContradictionDetector:
``asyncio.gather`` and the mapper's internal semaphore — this
helper itself awaits each step sequentially within one file.
"""
file_pages = await self._runtime.documents.read_pages(file.id)
file_pages = await self._runtime.documents.read_pages(file.id, principals=principals)
if not file_pages:
logger.info(
"[contradiction] no stored pages for %s (id=%s); skipping",
engine/src/stirling/agents/pdf_questions.py
+8 -2
View File
@@ -27,8 +27,9 @@ from stirling.contracts import (
format_file_names,
)
from stirling.documents import RagCapability
from stirling.models import PrincipalId
from stirling.models.agent_tool_models import AgentToolId, MathAuditorAgentParams
from stirling.services import AppRuntime
from stirling.services import AppRuntime, require_current_user_id
logger = logging.getLogger(__name__)
@@ -161,9 +162,10 @@ class PdfQuestionAgent:
)
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = []
for file in files:
if not await self.runtime.documents.has_collection(file.id):
if not await self.runtime.documents.has_collection(file.id, principals=principals):
missing.append(file)
return missing
@@ -175,8 +177,10 @@ class PdfQuestionAgent:
upstream classifier keeps that judgement in the same call that writes
the answer, and lets the agent mix tools when the question warrants it.
"""
principals = [PrincipalId(require_current_user_id())]
rag = RagCapability(
documents=self.runtime.documents,
principals=principals,
collections=[file.id for file in request.files],
top_k=self.runtime.settings.rag_default_top_k,
max_searches=self.runtime.settings.rag_max_searches,
@@ -184,11 +188,13 @@ class PdfQuestionAgent:
whole_doc = WholeDocReaderCapability(
runtime=self.runtime,
files=request.files,
principals=principals,
reasoner=self._chunked_reasoner,
)
contradiction = ContradictionCapability(
detector=self._contradiction_detector,
files=request.files,
principals=principals,
)
agent = Agent(
model=self.runtime.smart_model,
engine/src/stirling/agents/pdf_review.py
+9 -4
View File
@@ -49,14 +49,14 @@ from stirling.contracts import (
Verdict,
)
from stirling.contracts.ledger import Discrepancy
from stirling.models import ApiModel, ToolEndpoint
from stirling.models import ApiModel, PrincipalId, ToolEndpoint
from stirling.models.agent_tool_models import (
AgentToolId,
MathAuditorAgentParams,
PdfCommentAgentParams,
)
from stirling.models.tool_models import AddCommentsParams
from stirling.services import AppRuntime
from stirling.services import AppRuntime, require_current_user_id
# Fallback right-margin placement used when a finding has no usable
# anchor text. A4/Letter portrait assumed.
@@ -158,7 +158,11 @@ class PdfReviewAgent:
files_to_ingest=missing,
content_types=[PdfContentType.PAGE_TEXT],
)
report = await self._contradiction_detector.detect(request.files, query=request.user_message)
report = await self._contradiction_detector.detect(
request.files,
principals=[PrincipalId(require_current_user_id())],
query=request.user_message,
)
comments_json = await self._build_contradiction_comments_payload(request.user_message, report)
return EditPlanResponse(
summary="",
@@ -193,9 +197,10 @@ class PdfReviewAgent:
)
async def _find_missing_files(self, files: list[AiFile]) -> list[AiFile]:
principals = [PrincipalId(require_current_user_id())]
missing: list[AiFile] = []
for file in files:
if not await self.runtime.documents.has_collection(file.id):
if not await self.runtime.documents.has_collection(file.id, principals=principals):
missing.append(file)
return missing
engine/src/stirling/agents/shared/whole_doc_reader.py
+4 -1
View File
@@ -19,6 +19,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.agents.shared.chunked_reasoner import ChunkedReasoner
from stirling.contracts import AiFile
from stirling.models import PrincipalId
from stirling.services import AppRuntime
logger = logging.getLogger(__name__)
@@ -47,12 +48,14 @@ class WholeDocReaderCapability:
self,
runtime: AppRuntime,
files: list[AiFile],
principals: list[PrincipalId],
*,
reasoner: ChunkedReasoner | None = None,
max_reads: int = DEFAULT_MAX_READS,
) -> None:
self._runtime = runtime
self._files = files
self._principals = principals
self._reasoner = reasoner if reasoner is not None else ChunkedReasoner(runtime)
self._max_reads = max_reads
self._read_count = 0
@@ -115,7 +118,7 @@ class WholeDocReaderCapability:
sections: list[str] = []
for file in self._files:
pages = await self._runtime.documents.read_pages(file.id)
pages = await self._runtime.documents.read_pages(file.id, principals=self._principals)
if not pages:
logger.info(
"[whole-doc-reader] no stored pages for %s (id=%s); skipping",
engine/src/stirling/api/app.py
+51
View File
@@ -1,5 +1,7 @@
from __future__ import annotations
import asyncio
import logging
from contextlib import asynccontextmanager
from typing import Annotated
@@ -29,8 +31,45 @@ from stirling.api.routes import (
)
from stirling.config import AppSettings, load_settings
from stirling.contracts import HealthResponse
from stirling.documents import DocumentService
from stirling.services import build_runtime, setup_posthog_tracking
logger = logging.getLogger(__name__)
async def _run_expired_doc_reaper(
documents: DocumentService,
interval_seconds: int,
) -> None:
"""Periodically delete documents whose ``expires_at`` has passed.
A reaped collection drops everything rooted at that document. Backstop
for the explicit logout purge: catches sessions that ended without a
clean logout (tab close, JWT expiry, engine restart). Persistent rows
(``expires_at`` null, the shape we use for org-shared docs) are never
touched. Runs until cancelled by the lifespan teardown.
"""
await _reap(documents)
while True:
await asyncio.sleep(interval_seconds)
await _reap(documents)
async def _reap(documents: DocumentService) -> None:
"""One reaper iteration. Logs the deleted count on success and the full
exception with traceback on failure; never re-raises non-cancel errors so
a bad iteration doesn't kill the loop. ``asyncio.CancelledError`` is
re-raised so the lifespan teardown can cancel the task cleanly.
"""
try:
deleted = await documents.reap_expired()
if deleted:
logger.info("Reaped %d expired document collection(s)", deleted)
except asyncio.CancelledError:
raise
except Exception:
logger.exception("Document reaper iteration failed; will retry on next interval")
def _load_startup_settings(fast_api: FastAPI) -> AppSettings:
override = fast_api.dependency_overrides.get(load_settings)
@@ -56,7 +95,19 @@ async def lifespan(fast_api: FastAPI):
tracer_provider = setup_posthog_tracking(settings)
if tracer_provider:
Agent.instrument_all(InstrumentationSettings(tracer_provider=tracer_provider))
reaper_task = asyncio.create_task(
_run_expired_doc_reaper(
runtime.documents,
interval_seconds=settings.documents_reaper_interval_seconds,
),
name="expired-document-reaper",
)
yield
reaper_task.cancel()
try:
await reaper_task
except asyncio.CancelledError:
pass
await runtime.documents.close()
if tracer_provider:
tracer_provider.shutdown()
engine/src/stirling/api/dependencies.py
+20 -2
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from fastapi import Request
from fastapi import HTTPException, Request, status
from stirling.agents import (
ExecutionPlanningAgent,
@@ -12,7 +12,8 @@ from stirling.agents import (
from stirling.agents.ledger import MathAuditorAgent
from stirling.agents.pdf_comment import PdfCommentAgent
from stirling.documents import DocumentService
from stirling.services import AppRuntime
from stirling.models import UserId
from stirling.services import AppRuntime, current_user_id
def get_runtime(request: Request) -> AppRuntime:
@@ -49,3 +50,20 @@ def get_math_auditor_agent(request: Request) -> MathAuditorAgent:
def get_pdf_comment_agent(request: Request) -> PdfCommentAgent:
return request.app.state.pdf_comment_agent
def require_user_id() -> UserId:
"""FastAPI dependency for routes that touch per-user storage.
Reads ``X-User-Id`` (already extracted into a ContextVar by ``UserIdMiddleware``)
and returns it. Returns HTTP 401 if the caller didn't supply the header. Apply
to any route that ingests, searches, reads, or deletes document content so
the tenancy gate is enforced at the API boundary.
"""
user_id = current_user_id.get()
if user_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="X-User-Id header is required",
)
return user_id
engine/src/stirling/api/middleware.py
+2 -1
View File
@@ -4,6 +4,7 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin
from starlette.requests import Request
from starlette.responses import Response
from stirling.models import UserId
from stirling.services.tracking import current_user_id
_USER_ID_HEADER = "X-User-Id"
@@ -15,7 +16,7 @@ class UserIdMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
user_id = request.headers.get(_USER_ID_HEADER)
if user_id:
token = current_user_id.set(user_id)
token = current_user_id.set(UserId(user_id))
try:
return await call_next(request)
finally:
engine/src/stirling/api/routes/documents.py
+47 -13
View File
@@ -1,49 +1,83 @@
from __future__ import annotations
import logging
from typing import Annotated
from fastapi import APIRouter, Depends
from stirling.api.dependencies import get_document_service
from stirling.api.dependencies import get_document_service, require_user_id
from stirling.contracts import (
DeleteDocumentResponse,
IngestDocumentRequest,
IngestDocumentResponse,
)
from stirling.contracts.documents import PurgeOwnerResponse
from stirling.documents import DocumentService
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId, UserId
router = APIRouter(prefix="/api/v1/documents", tags=["documents"])
logger = logging.getLogger(__name__)
@router.post("", response_model=IngestDocumentResponse)
async def ingest_document(
request: IngestDocumentRequest,
documents: Annotated[DocumentService, Depends(get_document_service)],
_user_id: Annotated[UserId, Depends(require_user_id)],
) -> IngestDocumentResponse:
"""Replace-ingest a document's content under ``document_id``.
"""Replace-ingest a document's content under ``(owner_id, read_principals)``.
Stores both representations in one shot:
* embedded chunks for RAG search,
* ordered page text for whole-document reading.
Any previously-stored content for this document is removed first.
``owner_id`` and ``read_principals`` are required on the request body and
are stored verbatim.
``X-User-Id`` is still required so the caller is authenticated and
PostHog tracks the right user, but the value is not used to constrain
the body.
"""
pages = request.page_text or []
chunks_indexed = await documents.ingest(
collection=request.document_id,
pages=pages,
pages=request.page_text or [],
source=request.source,
owner_id=request.owner_id,
read_principals=request.read_principals,
expires_at=request.expires_at,
)
return IngestDocumentResponse(document_id=request.document_id, chunks_indexed=chunks_indexed)
@router.delete("/{document_id}", response_model=DeleteDocumentResponse)
@router.delete("/by-id/{document_id}", response_model=DeleteDocumentResponse)
async def delete_document(
document_id: FileId,
documents: Annotated[DocumentService, Depends(get_document_service)],
user_id: Annotated[UserId, Depends(require_user_id)],
) -> DeleteDocumentResponse:
"""Remove a document's content. Idempotent."""
existed = await documents.has_collection(document_id)
"""Remove the caller's copy of this document.
Owner is inferred from the caller - only personal-doc deletes go through
here. Org-doc deletes will need an explicit owner_id once we add the
admin endpoints for them; for now this route can't reach docs owned by
a different principal.
"""
owner_id = OwnerId(user_id)
principals = [PrincipalId(user_id)]
existed = await documents.has_collection(document_id, principals=principals)
if existed:
await documents.delete_collection(document_id)
await documents.delete_collection(document_id, owner_id=owner_id)
return DeleteDocumentResponse(document_id=document_id, deleted=existed)
@router.delete("/by-owner", response_model=PurgeOwnerResponse)
async def purge_caller_documents(
documents: Annotated[DocumentService, Depends(get_document_service)],
user_id: Annotated[UserId, Depends(require_user_id)],
) -> PurgeOwnerResponse:
"""Delete every personal-doc collection owned by the caller.
Called by Java on logout so a user's document content disappears as soon as
the session ends. Org-owned docs (where the caller is a reader but not
the owner) are not touched - only collections whose ``owner_id`` matches
the calling user are removed.
"""
deleted = await documents.purge_owner(OwnerId(user_id))
logger.info("Purged %d collection(s) for owner=%s", deleted, user_id)
return PurgeOwnerResponse(owner_id=OwnerId(user_id), deleted=deleted)
engine/src/stirling/config/settings.py
+4
View File
@@ -37,6 +37,10 @@ class AppSettings(BaseSettings):
rag_chunk_overlap: int = Field(validation_alias="STIRLING_RAG_CHUNK_OVERLAP")
rag_default_top_k: int = Field(validation_alias="STIRLING_RAG_TOP_K")
rag_max_searches: int = Field(validation_alias="STIRLING_RAG_MAX_SEARCHES")
documents_reaper_interval_seconds: int = Field(
default=900,
validation_alias="STIRLING_DOCUMENTS_REAPER_INTERVAL_SECONDS",
)
# Chunked reasoner settings (whole-document map-reduce).
chunked_reasoner_chars_per_slice: int = Field(validation_alias="STIRLING_CHUNKED_REASONER_CHARS_PER_SLICE")
engine/src/stirling/contracts/__init__.py
+2
View File
@@ -42,6 +42,7 @@ from .documents import (
Page,
PageRange,
PageText,
PurgeOwnerResponse,
)
from .execution import (
AgentExecutionRequest,
@@ -137,6 +138,7 @@ __all__ = [
"ContradictionSeverity",
"ConversationMessage",
"DeleteDocumentResponse",
"PurgeOwnerResponse",
"PdfToMarkdownCannotDoResponse",
"PdfToMarkdownOrchestrateResponse",
"PdfToMarkdownRequest",
engine/src/stirling/contracts/documents.py
+19 -3
View File
@@ -1,8 +1,10 @@
from __future__ import annotations
from datetime import datetime
from pydantic import Field
from stirling.models import ApiModel
from stirling.models import ApiModel, OwnerId, PrincipalId
from .common import FileId
@@ -35,20 +37,27 @@ class PageRange(ApiModel):
class IngestDocumentRequest(ApiModel):
"""Replace-ingest a document's content under the given ``document_id``.
"""Replace-ingest a document's content under ``(document_id, owner_id)``.
Each call wipes any previously-stored content for the document and writes
Each call wipes any previously-stored content for the pair and writes
both the vector-chunk and ordered-page representations from the supplied
pages.
``source`` is a human-readable label (typically the original filename)
that flows into chunk metadata so search results are readable when
``document_id`` is a hash.
``owner_id`` and ``read_principals`` are required: the engine never
defaults them. Callers must declare ownership and access explicitly.
"""
document_id: FileId = Field(min_length=1)
source: str = Field(min_length=1)
page_text: list[PageText] | None = None
owner_id: OwnerId = Field(min_length=1)
read_principals: list[PrincipalId] = Field(min_length=1)
# When to delete this doc. ``None`` means "persistent" (keep until an explicit delete)
expires_at: datetime | None
class IngestDocumentResponse(ApiModel):
@@ -59,3 +68,10 @@ class IngestDocumentResponse(ApiModel):
class DeleteDocumentResponse(ApiModel):
document_id: FileId
deleted: bool
class PurgeOwnerResponse(ApiModel):
"""Returned by ``DELETE /api/v1/documents/by-owner``."""
owner_id: OwnerId
deleted: int = Field(ge=0)
engine/src/stirling/documents/pgvector_store.py
+228 -60
View File
@@ -1,24 +1,34 @@
from __future__ import annotations
import json
from datetime import datetime
import psycopg
from pgvector.psycopg import register_vector_async
from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
class PgVectorStore(DocumentStore):
"""PostgreSQL + pgvector backed store.
"""PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads.
Connects to an external Postgres instance (DSN provided via config) and uses the
`vector` extension for similarity search. The schema is created on first use.
Holds two tables under the same connection:
Owned tables:
* ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
* ``rag_documents`` - vector chunks for RAG search.
* ``document_pages`` - ordered page text for whole-document reading.
* ``document_acl`` - principals (users, groups, orgs) granted permissions
on a ``(collection, owner_id)`` pair. Read methods take the caller's
principal set and join through this table.
Writes are owner-scoped; reads are ACL-scoped.
"""
def __init__(self, dsn: str) -> None:
@@ -41,60 +51,141 @@ class PgVectorStore(DocumentStore):
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY,
source TEXT NOT NULL
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
source TEXT NOT NULL,
expires_at TIMESTAMP WITH TIME ZONE,
PRIMARY KEY (collection, owner_id)
)
"""
)
# Partial index over rows that can actually expire keeps the reaper
# scan tight even when most rows are persistent (org docs).
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_meta_expires_at "
"ON documents_meta(expires_at) WHERE expires_at IS NOT NULL"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS rag_documents (
id TEXT NOT NULL,
collection TEXT NOT NULL
REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
text TEXT NOT NULL,
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
embedding vector NOT NULL,
PRIMARY KEY (id, collection)
PRIMARY KEY (id, collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
await cur.execute("CREATE INDEX IF NOT EXISTS idx_rag_collection ON rag_documents(collection)")
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_collection_owner ON rag_documents(collection, owner_id)"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS document_pages (
collection TEXT NOT NULL
REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
page_number INTEGER NOT NULL,
text TEXT NOT NULL,
char_count INTEGER NOT NULL,
PRIMARY KEY (collection, page_number)
PRIMARY KEY (collection, owner_id, page_number),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
await cur.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
)
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS document_acl (
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
permission TEXT NOT NULL,
PRIMARY KEY (collection, owner_id, principal_id, permission),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
# Hot path: every read joins ``WHERE principal_id = ANY(...) AND permission = ?``.
await cur.execute(
"CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
)
await conn.commit()
self._initialized = True
async def ensure_collection(self, collection: str, source: str) -> None:
# ── lifecycle of the (collection, owner_id) row ────────────────────────
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"""
INSERT INTO documents_meta (collection, source)
VALUES (%s, %s)
ON CONFLICT (collection) DO UPDATE SET source = EXCLUDED.source
INSERT INTO documents_meta (collection, owner_id, source, expires_at)
VALUES (%s, %s, %s, %s)
ON CONFLICT (collection, owner_id) DO UPDATE SET
source = EXCLUDED.source,
expires_at = EXCLUDED.expires_at
""",
(collection, source),
(collection, owner_id, source, expires_at),
)
await conn.commit()
async def purge_owner(self, owner_id: OwnerId) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM documents_meta WHERE owner_id = %s",
(owner_id,),
)
deleted = cur.rowcount
await conn.commit()
return deleted
async def reap_expired(self) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()")
deleted = cur.rowcount
await conn.commit()
return deleted
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
# Cascade FKs handle rag_documents, document_pages, document_acl.
await cur.execute(
"DELETE FROM documents_meta WHERE collection = %s AND owner_id = %s",
(collection, owner_id),
)
deleted = cur.rowcount > 0
await conn.commit()
return deleted
# ── write paths ────────────────────────────────────────────────────────
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
@@ -107,36 +198,100 @@ class PgVectorStore(DocumentStore):
for doc, emb in zip(documents, embeddings):
await cur.execute(
"""
INSERT INTO rag_documents (id, collection, text, metadata, embedding)
VALUES (%s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection)
INSERT INTO rag_documents (id, collection, owner_id, text, metadata, embedding)
VALUES (%s, %s, %s, %s, %s::jsonb, %s)
ON CONFLICT (id, collection, owner_id)
DO UPDATE SET
text = EXCLUDED.text,
metadata = EXCLUDED.metadata,
embedding = EXCLUDED.embedding
""",
(doc.id, collection, doc.text, json.dumps(doc.metadata), emb),
(doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), emb),
)
await conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM document_pages WHERE collection = %s AND owner_id = %s",
(collection, owner_id),
)
if pages:
await cur.executemany(
"""
INSERT INTO document_pages (collection, owner_id, page_number, text, char_count)
VALUES (%s, %s, %s, %s, %s)
""",
[(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages],
)
await conn.commit()
# ── ACL management ─────────────────────────────────────────────────────
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
if not principals:
return
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.executemany(
"""
INSERT INTO document_acl (collection, owner_id, principal_id, permission)
VALUES (%s, %s, %s, %s)
ON CONFLICT (collection, owner_id, principal_id, permission) DO NOTHING
""",
[(collection, owner_id, p, _READ_PERMISSION) for p in principals],
)
await conn.commit()
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s",
(collection, owner_id, principal),
)
await conn.commit()
# ── read paths (ACL-gated) ─────────────────────────────────────────────
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
await cur.execute(
"""
SELECT id, text, metadata, 1 - (embedding <=> %s) AS score
FROM rag_documents
WHERE collection = %s
WHERE collection = %s AND owner_id = %s
ORDER BY embedding <=> %s
LIMIT %s
""",
(query_embedding, collection, query_embedding, top_k),
(query_embedding, collection, owner_id, query_embedding, top_k),
)
rows = await cur.fetchall()
@@ -148,71 +303,84 @@ class PgVectorStore(DocumentStore):
for r in rows
]
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("DELETE FROM document_pages WHERE collection = %s", (collection,))
if pages:
await cur.executemany(
"""
INSERT INTO document_pages (collection, page_number, text, char_count)
VALUES (%s, %s, %s, %s)
""",
[(collection, p.page_number, p.text, p.char_count) for p in pages],
)
await conn.commit()
@staticmethod
async def _readable_owner_for(
cur: psycopg.AsyncCursor,
collection: str,
principals: list[PrincipalId],
) -> str | None:
"""Resolve which owner_id this caller is reading. ``None`` means no access."""
await cur.execute(
"""
SELECT owner_id FROM document_acl
WHERE collection = %s
AND permission = %s
AND principal_id = ANY(%s)
ORDER BY owner_id
LIMIT 1
""",
(collection, _READ_PERMISSION, list(principals)),
)
row = await cur.fetchone()
return row[0] if row else None
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
return []
if page_range is None:
await cur.execute(
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = %s ORDER BY page_number",
(collection,),
"WHERE collection = %s AND owner_id = %s ORDER BY page_number",
(collection, owner_id),
)
else:
await cur.execute(
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = %s AND page_number BETWEEN %s AND %s "
"WHERE collection = %s AND owner_id = %s "
"AND page_number BETWEEN %s AND %s "
"ORDER BY page_number",
(collection, page_range.start, page_range.end),
(collection, owner_id, page_range.start, page_range.end),
)
rows = await cur.fetchall()
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
async def delete_collection(self, collection: str) -> None:
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals:
return False
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
# Cascade FKs handle rag_documents and document_pages.
await cur.execute("DELETE FROM documents_meta WHERE collection = %s", (collection,))
await conn.commit()
owner_id = await self._readable_owner_for(cur, collection, principals)
return owner_id is not None
async def list_collections(self) -> list[str]:
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute("SELECT collection FROM documents_meta ORDER BY collection")
rows = await cur.fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
async with conn.cursor() as cur:
await cur.execute(
"SELECT 1 FROM documents_meta WHERE collection = %s",
(collection,),
"""
SELECT DISTINCT collection FROM document_acl
WHERE permission = %s
AND principal_id = ANY(%s)
ORDER BY collection
""",
(_READ_PERMISSION, list(principals)),
)
row = await cur.fetchone()
return row is not None
rows = await cur.fetchall()
return [r[0] for r in rows]
async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release.
engine/src/stirling/documents/rag_capability.py
+14 -8
View File
@@ -8,7 +8,7 @@ from pydantic_ai.toolsets import AbstractToolset
from stirling.documents.service import DocumentService
from stirling.documents.store import SearchResult
from stirling.models import FileId
from stirling.models import FileId, PrincipalId
logger = logging.getLogger(__name__)
@@ -16,9 +16,13 @@ logger = logging.getLogger(__name__)
class RagCapability:
"""Bundles RAG instructions and the ``search_knowledge`` toolset for agent injection.
Agents consume this as::
Agents construct one per run, scoped to the caller's principal set, as::
rag = runtime.rag_capability
rag = RagCapability(
documents=self.runtime.documents,
principals=current_principals(),
collections=[file.id for file in request.files],
)
Agent(
...,
instructions=rag.instructions,
@@ -26,20 +30,22 @@ class RagCapability:
)
When no collections are pinned, the instructions are generated dynamically at
run time so the agent sees the current list of collections in the store.
run time so the agent sees the collections the caller can read.
Lifecycle: a ``RagCapability`` instance is intended to live for the duration of a
single agent run.
single agent run and binds to one caller's principal set.
"""
def __init__(
self,
documents: DocumentService,
principals: list[PrincipalId],
collections: list[FileId] | None = None,
top_k: int = 5,
max_searches: int = 5,
) -> None:
self._documents = documents
self._principals = principals
self._collections = collections
self._top_k = top_k
self._max_searches = max_searches
@@ -74,7 +80,7 @@ class RagCapability:
)
async def _dynamic_instructions(self) -> str:
collections = await self._documents.list_collections()
collections = await self._documents.list_collections(self._principals)
if collections:
names = ", ".join(collections)
collection_desc = f"the following knowledge base collections: {names}"
@@ -115,12 +121,12 @@ class RagCapability:
if self._collections:
all_results = []
for col in self._collections:
col_results = await self._documents.search(query, collection=col, top_k=k)
col_results = await self._documents.search(query, principals=self._principals, collection=col, top_k=k)
all_results.extend(col_results)
all_results.sort(key=lambda r: r.score, reverse=True)
results = all_results[:k]
else:
results = await self._documents.search(query, top_k=k)
results = await self._documents.search(query, principals=self._principals, top_k=k)
if not results:
logger.info("[rag] search_knowledge query=%r -> 0 results", query)
engine/src/stirling/documents/service.py
+90 -32
View File
@@ -1,11 +1,12 @@
from __future__ import annotations
import logging
from datetime import datetime
from stirling.contracts.documents import Page, PageRange, PageText
from stirling.documents.embedder import EmbeddingService
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId
logger = logging.getLogger(__name__)
@@ -17,7 +18,8 @@ PAGE_TEXT_CONTENT_TYPE = "page_text"
class DocumentService:
"""Top-level facade for stored document content.
Holds two representations of every document under a single ``collection``:
Holds two representations of every document under a single
``(collection, owner_id)`` pair:
* **Vector chunks** for RAG-style semantic retrieval (``search``).
* **Ordered pages** for whole-document reading (``read_pages``).
@@ -25,6 +27,15 @@ class DocumentService:
Both are populated by :meth:`ingest` from a single ``pages`` payload. Agents
pick the strategy that fits the question; they don't need to know which
storage they're hitting.
**Owner vs principal.** ``owner_id`` is the tenant (a user or an org).
``principals`` is the caller's accessible principal set, matched against
the ACL on every read. Personal-doc behaviour: ingest with
``owner_id=user:bob`` and the route grants ``user:bob`` read access; Bob's
later searches pass ``principals=[user:bob]`` and find his docs.
Org-doc behaviour: ingest with ``owner_id=org:acme`` and grant read to
whichever groups should see it (``group:engineering``, etc.); members'
principal sets pick the doc up automatically.
"""
def __init__(self, embedder: EmbeddingService, store: DocumentStore, default_top_k: int = 5) -> None:
@@ -37,20 +48,27 @@ class DocumentService:
collection: FileId,
pages: list[PageText],
source: str,
owner_id: OwnerId,
read_principals: list[PrincipalId],
expires_at: datetime | None,
) -> int:
"""Replace-ingest a document. Returns the number of vector chunks indexed.
This wipes any previously-stored content for ``collection`` and writes
both the vector-chunk and page-text representations from the same
``pages`` payload. Pages with empty/whitespace-only text are skipped
for chunking but still written to the page store so page numbering is
preserved end-to-end.
Wipes any previously-stored content for ``(collection, owner_id)`` and
writes both the vector-chunk and page-text representations from the
same ``pages`` payload. Pages with empty/whitespace-only text are
skipped for chunking but still written to the page store so page
numbering is preserved end-to-end.
"""
await self._store.delete_collection(collection)
await self._store.ensure_collection(collection, source)
if not read_principals:
raise ValueError("read_principals must not be empty - every doc needs at least one reader")
await self._store.delete_collection(collection, owner_id)
await self._store.ensure_collection(collection, source, owner_id, expires_at)
stored_pages = [StoredPage(page_number=p.page_number, text=p.text, char_count=len(p.text)) for p in pages]
await self._store.add_pages(collection, stored_pages)
await self._store.add_pages(collection, stored_pages, owner_id)
await self._store.grant_read(collection, owner_id, read_principals)
chunks: list[Document] = []
for page in pages:
@@ -70,63 +88,103 @@ class DocumentService:
if not chunks:
return 0
embeddings = await self._embedder.embed_documents([doc.text for doc in chunks])
await self._store.add_documents(collection, chunks, embeddings)
await self._store.add_documents(collection, chunks, embeddings, owner_id)
return len(chunks)
async def search(
self,
query: str,
principals: list[PrincipalId],
collection: FileId | None = None,
top_k: int | None = None,
) -> list[SearchResult]:
"""Embed query and search across one or all collections.
"""Embed query and search collections readable by ``principals``.
If collection is None, searches all available collections and merges results.
If ``collection`` is supplied, search only that collection (and only
if at least one principal can read it). If ``None``, search every
collection any principal in ``principals`` can read, merge, and
return the top-k.
"""
k = top_k if top_k is not None else self._default_top_k
query_embedding = await self._embedder.embed_query(query)
if collection is not None:
if not await self._store.has_collection(collection):
if not await self._store.has_collection(collection, principals):
return []
return await self._store.search(collection, query_embedding, k)
return await self._store.search(collection, query_embedding, k, principals)
# Search all collections, skipping any that error (e.g. dimension mismatch)
collections = await self._store.list_collections()
# Search every collection the caller can read, skipping any that error
# (e.g. dimension mismatch).
collections = await self._store.list_collections(principals)
all_results: list[SearchResult] = []
for col_name in collections:
try:
results = await self._store.search(col_name, query_embedding, k)
results = await self._store.search(col_name, query_embedding, k, principals)
all_results.extend(results)
except Exception: # noqa: BLE001 - any backend error on one collection should not stop the others
logger.warning("Skipping collection %s during cross-collection search", col_name, exc_info=True)
logger.warning(
"Skipping collection %s during cross-collection search",
col_name,
exc_info=True,
)
# Sort by score descending, return top_k across all collections
# Sort by score descending, return top_k across all the caller's collections
all_results.sort(key=lambda r: r.score, reverse=True)
return all_results[:k]
async def read_pages(
self,
collection: FileId,
principals: list[PrincipalId],
page_range: PageRange | None = None,
) -> list[Page]:
"""Return ordered page text for ``collection``.
"""Return ordered page text for ``collection`` if any principal can read it."""
return await self._store.read_pages(collection, page_range, principals)
Empty list if the collection has no stored pages.
async def delete_collection(self, collection: FileId, owner_id: OwnerId) -> bool:
"""Remove a collection (chunks, pages, ACL).
Returns ``True`` if the collection was found and deleted, ``False`` if
no row matched ``(collection, owner_id)``.
"""
return await self._store.read_pages(collection, page_range)
return await self._store.delete_collection(collection, owner_id)
async def delete_collection(self, collection: FileId) -> None:
"""Remove a collection's chunks and pages."""
await self._store.delete_collection(collection)
async def purge_owner(self, owner_id: OwnerId) -> int:
"""Remove every collection ``owner_id`` owns, including vector chunks,
page text, and ACL rows. Returns the number of collections purged."""
return await self._store.purge_owner(owner_id)
async def has_collection(self, collection: FileId) -> bool:
"""Check whether a collection exists."""
return await self._store.has_collection(collection)
async def reap_expired(self) -> int:
"""Delete collections whose ``expires_at`` is set and in the past.
Persistent collections (``expires_at=null``) are never touched.
Returns the number of collections deleted."""
return await self._store.reap_expired()
async def list_collections(self) -> list[FileId]:
"""List all available collections."""
return [FileId(name) for name in await self._store.list_collections()]
async def has_collection(self, collection: FileId, principals: list[PrincipalId]) -> bool:
"""Check whether at least one principal can read this collection."""
return await self._store.has_collection(collection, principals)
async def list_collections(self, principals: list[PrincipalId]) -> list[FileId]:
"""List collections readable by at least one of ``principals``."""
return [FileId(name) for name in await self._store.list_collections(principals)]
async def grant_read(
self,
collection: FileId,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
"""Grant read access to additional principals on an existing doc."""
await self._store.grant_read(collection, owner_id, principals)
async def revoke(
self,
collection: FileId,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
"""Revoke a principal's access on an existing doc."""
await self._store.revoke(collection, owner_id, principal)
async def close(self) -> None:
"""Release the underlying store's resources."""
engine/src/stirling/documents/sqlite_vec_store.py
+314 -86
View File
@@ -5,19 +5,35 @@ import json
import math
import re
import sqlite3
from datetime import UTC, datetime
from pathlib import Path
import sqlite_vec
from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
# sqlite stores TIMESTAMP as TEXT. We normalise to UTC ISO 8601 ``YYYY-MM-DD HH:MM:SS``
# so lexicographic comparison against ``datetime('now')`` matches chronological order.
_SQLITE_DATETIME_FMT = "%Y-%m-%d %H:%M:%S"
def _to_sqlite_utc(dt: datetime | None) -> str | None:
if dt is None:
return None
if dt.tzinfo is not None:
dt = dt.astimezone(UTC).replace(tzinfo=None)
return dt.strftime(_SQLITE_DATETIME_FMT)
class SqliteVecStore(DocumentStore):
"""sqlite-vec backed vector store. Single-file SQLite database, embedded, no server.
Each collection gets its own `vec0` virtual table with a fixed embedding dimension
(detected on first insert). Document metadata lives in a regular table joined by rowid.
Each ``(collection, owner_id)`` pair gets its own `vec0` virtual table with a
fixed embedding dimension (detected on first insert). Document metadata lives
in a regular table joined by rowid.
"""
def __init__(self, db_path: str | Path) -> None:
@@ -51,17 +67,30 @@ class SqliteVecStore(DocumentStore):
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS documents_meta (
collection TEXT PRIMARY KEY,
source TEXT NOT NULL
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
source TEXT NOT NULL,
expires_at TIMESTAMP,
PRIMARY KEY (collection, owner_id)
)
"""
)
# The reaper filters on ``expires_at IS NOT NULL AND expires_at < now`` so
# a partial index over non-null rows keeps the scan tight even when most
# rows are persistent (org docs).
self._conn.execute(
"CREATE INDEX IF NOT EXISTS idx_meta_expires_at ON documents_meta(expires_at) WHERE expires_at IS NOT NULL"
)
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS collections (
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
dim INTEGER NOT NULL,
table_name TEXT NOT NULL
table_name TEXT NOT NULL,
PRIMARY KEY (collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
@@ -69,47 +98,154 @@ class SqliteVecStore(DocumentStore):
"""
CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL,
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
text TEXT NOT NULL,
metadata TEXT NOT NULL DEFAULT '{}',
vec_rowid INTEGER NOT NULL,
PRIMARY KEY (id, collection)
PRIMARY KEY (id, collection, owner_id),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection ON documents(collection)")
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection_owner ON documents(collection, owner_id)")
self._conn.execute(
"""
CREATE TABLE IF NOT EXISTS document_pages (
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
page_number INTEGER NOT NULL,
text TEXT NOT NULL,
char_count INTEGER NOT NULL,
PRIMARY KEY (collection, page_number)
PRIMARY KEY (collection, owner_id, page_number),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_collection ON document_pages(collection)")
self._conn.commit()
async def ensure_collection(self, collection: str, source: str) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_ensure_collection, collection, source)
def _sync_ensure_collection(self, collection: str, source: str) -> None:
self._conn.execute(
"CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
)
self._conn.execute(
"""
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
""",
(collection, source),
CREATE TABLE IF NOT EXISTS document_acl (
collection TEXT NOT NULL,
owner_id TEXT NOT NULL,
principal_id TEXT NOT NULL,
permission TEXT NOT NULL,
PRIMARY KEY (collection, owner_id, principal_id, permission),
FOREIGN KEY (collection, owner_id)
REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE
)
"""
)
# Lookup by principal is the hot path for search/list (every read
# joins through this index). Composite ordering matches the WHERE.
self._conn.execute(
"CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
)
self._conn.commit()
# ── lifecycle of the (collection, owner_id) row ────────────────────────
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_ensure_collection, collection, source, owner_id, expires_at)
def _sync_ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
self._conn.execute(
"""
INSERT INTO documents_meta(collection, owner_id, source, expires_at)
VALUES (?, ?, ?, ?)
ON CONFLICT(collection, owner_id) DO UPDATE SET
source = excluded.source,
expires_at = excluded.expires_at
""",
(collection, owner_id, source, _to_sqlite_utc(expires_at)),
)
self._conn.commit()
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
async with self._lock:
return await asyncio.to_thread(self._sync_delete_collection, collection, owner_id)
def _sync_delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
# (collections, documents, document_pages, document_acl) when documents_meta is deleted.
row = self._conn.execute(
"SELECT table_name FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is not None:
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
cursor = self._conn.execute(
"DELETE FROM documents_meta WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
)
self._conn.commit()
return cursor.rowcount > 0
async def purge_owner(self, owner_id: OwnerId) -> int:
async with self._lock:
return await asyncio.to_thread(self._sync_purge_owner, owner_id)
def _sync_purge_owner(self, owner_id: OwnerId) -> int:
# Drop all vec0 virtual tables for this owner first (FK cascade can't reach them).
vec_tables = [
r[0]
for r in self._conn.execute("SELECT table_name FROM collections WHERE owner_id = ?", (owner_id,)).fetchall()
]
for name in vec_tables:
self._conn.execute(f"DROP TABLE IF EXISTS {name}")
cursor = self._conn.execute("DELETE FROM documents_meta WHERE owner_id = ?", (owner_id,))
self._conn.commit()
return cursor.rowcount
async def reap_expired(self) -> int:
async with self._lock:
return await asyncio.to_thread(self._sync_reap_expired)
def _sync_reap_expired(self) -> int:
# Drop vec0 virtual tables for expired collections first (FK cascade can't reach them).
vec_tables = [
r[0]
for r in self._conn.execute(
"""
SELECT c.table_name FROM collections c
JOIN documents_meta m
ON m.collection = c.collection AND m.owner_id = c.owner_id
WHERE m.expires_at IS NOT NULL AND m.expires_at < datetime('now')
"""
).fetchall()
]
for name in vec_tables:
self._conn.execute(f"DROP TABLE IF EXISTS {name}")
cursor = self._conn.execute(
"DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < datetime('now')"
)
self._conn.commit()
return cursor.rowcount
# ── write paths ────────────────────────────────────────────────────────
@staticmethod
def _sanitize_table_name(collection: str) -> str:
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
return f"vec_{safe}"
def _sanitize_table_name(collection: str, owner_id: OwnerId) -> str:
safe_col = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
safe_owner = re.sub(r"[^a-zA-Z0-9_]", "_", owner_id)
return f"vec_{safe_owner}_{safe_col}"
@staticmethod
def _normalize(vector: list[float]) -> list[float]:
@@ -123,6 +259,7 @@ class SqliteVecStore(DocumentStore):
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
if len(documents) != len(embeddings):
raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings")
@@ -130,34 +267,40 @@ class SqliteVecStore(DocumentStore):
return
async with self._lock:
await asyncio.to_thread(self._sync_add, collection, documents, embeddings)
await asyncio.to_thread(self._sync_add, collection, documents, embeddings, owner_id)
def _sync_add(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
dim = len(embeddings[0])
row = self._conn.execute("SELECT dim, table_name FROM collections WHERE name = ?", (collection,)).fetchone()
row = self._conn.execute(
"SELECT dim, table_name FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is None:
table_name = self._sanitize_table_name(collection)
table_name = self._sanitize_table_name(collection, owner_id)
self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])")
self._conn.execute(
"INSERT INTO collections(name, dim, table_name) VALUES (?, ?, ?)",
(collection, dim, table_name),
"INSERT INTO collections(collection, owner_id, dim, table_name) VALUES (?, ?, ?, ?)",
(collection, owner_id, dim, table_name),
)
else:
existing_dim, table_name = row
if existing_dim != dim:
raise ValueError(f"Collection {collection} has dim {existing_dim}, got embedding of dim {dim}")
raise ValueError(
f"Collection {collection} for owner {owner_id} has dim {existing_dim}, got embedding of dim {dim}"
)
# Upsert: delete existing docs with matching IDs first
ids = [doc.id for doc in documents]
placeholders = ",".join("?" * len(ids))
existing = self._conn.execute(
f"SELECT vec_rowid FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
f"SELECT vec_rowid FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, owner_id, *ids),
).fetchall()
if existing:
vec_rowids = [r[0] for r in existing]
@@ -167,8 +310,8 @@ class SqliteVecStore(DocumentStore):
vec_rowids,
)
self._conn.execute(
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})",
(collection, *ids),
f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
(collection, owner_id, *ids),
)
for doc, emb in zip(documents, embeddings):
@@ -179,27 +322,101 @@ class SqliteVecStore(DocumentStore):
)
vec_rowid = cursor.lastrowid
self._conn.execute(
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)",
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid),
"INSERT INTO documents(id, collection, owner_id, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?, ?)",
(doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), vec_rowid),
)
self._conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_add_pages, collection, pages, owner_id)
def _sync_add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
self._conn.execute(
"DELETE FROM document_pages WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
)
if pages:
self._conn.executemany(
"INSERT INTO document_pages(collection, owner_id, page_number, text, char_count) "
"VALUES (?, ?, ?, ?, ?)",
[(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages],
)
self._conn.commit()
# ── ACL management ─────────────────────────────────────────────────────
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
if not principals:
return
async with self._lock:
await asyncio.to_thread(self._sync_grant_read, collection, owner_id, principals)
def _sync_grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
self._conn.executemany(
"""
INSERT INTO document_acl(collection, owner_id, principal_id, permission)
VALUES (?, ?, ?, ?)
ON CONFLICT(collection, owner_id, principal_id, permission) DO NOTHING
""",
[(collection, owner_id, p, _READ_PERMISSION) for p in principals],
)
self._conn.commit()
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_revoke, collection, owner_id, principal)
def _sync_revoke(self, collection: str, owner_id: OwnerId, principal: PrincipalId) -> None:
self._conn.execute(
"DELETE FROM document_acl WHERE collection = ? AND owner_id = ? AND principal_id = ?",
(collection, owner_id, principal),
)
self._conn.commit()
# ── read paths (ACL-gated) ─────────────────────────────────────────────
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
async with self._lock:
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k)
return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k, principals)
def _sync_search(
self,
collection: str,
query_embedding: list[float],
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone()
if not principals:
return []
owner_id = self._readable_owner_for(collection, principals)
if owner_id is None:
return []
row = self._conn.execute(
"SELECT table_name, dim FROM collections WHERE collection = ? AND owner_id = ?",
(collection, owner_id),
).fetchone()
if row is None:
return []
table_name, dim = row
@@ -213,11 +430,14 @@ class SqliteVecStore(DocumentStore):
f"""
SELECT d.id, d.text, d.metadata, v.distance
FROM {table_name} v
JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ?
JOIN documents d
ON d.vec_rowid = v.rowid
AND d.collection = ?
AND d.owner_id = ?
WHERE v.embedding MATCH ? AND k = ?
ORDER BY v.distance
""",
(collection, query_blob, top_k),
(collection, owner_id, query_blob, top_k),
).fetchall()
return [
@@ -233,77 +453,85 @@ class SqliteVecStore(DocumentStore):
for r in results
]
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_add_pages, collection, pages)
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None:
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,))
if pages:
self._conn.executemany(
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)",
[(collection, p.page_number, p.text, p.char_count) for p in pages],
)
self._conn.commit()
def _readable_owner_for(self, collection: str, principals: list[PrincipalId]) -> str | None:
"""Resolve which owner_id this caller is reading. ``None`` means no access."""
placeholders = ",".join("?" * len(principals))
row = self._conn.execute(
f"""
SELECT owner_id FROM document_acl
WHERE collection = ?
AND permission = ?
AND principal_id IN ({placeholders})
ORDER BY owner_id
LIMIT 1
""",
(collection, _READ_PERMISSION, *principals),
).fetchone()
return row[0] if row else None
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
async with self._lock:
return await asyncio.to_thread(self._sync_read_pages, collection, page_range)
return await asyncio.to_thread(self._sync_read_pages, collection, page_range, principals)
def _sync_read_pages(
self,
collection: str,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
if not principals:
return []
owner_id = self._readable_owner_for(collection, principals)
if owner_id is None:
return []
if page_range is None:
rows = self._conn.execute(
"SELECT page_number, text, char_count FROM document_pages WHERE collection = ? ORDER BY page_number",
(collection,),
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = ? AND owner_id = ? ORDER BY page_number",
(collection, owner_id),
).fetchall()
else:
rows = self._conn.execute(
"SELECT page_number, text, char_count FROM document_pages "
"WHERE collection = ? AND page_number BETWEEN ? AND ? ORDER BY page_number",
(collection, page_range.start, page_range.end),
"WHERE collection = ? AND owner_id = ? AND page_number BETWEEN ? AND ? "
"ORDER BY page_number",
(collection, owner_id, page_range.start, page_range.end),
).fetchall()
return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows]
async def delete_collection(self, collection: str) -> None:
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
async with self._lock:
await asyncio.to_thread(self._sync_delete_collection, collection)
return await asyncio.to_thread(self._sync_has_collection, collection, principals)
def _sync_delete_collection(self, collection: str) -> None:
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
# (collections, documents, document_pages) when documents_meta is deleted.
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
if row is not None:
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,))
self._conn.commit()
def _sync_has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals:
return False
return self._readable_owner_for(collection, principals) is not None
async def list_collections(self) -> list[str]:
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
async with self._lock:
return await asyncio.to_thread(self._sync_list_collections)
return await asyncio.to_thread(self._sync_list_collections, principals)
def _sync_list_collections(self) -> list[str]:
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall()
def _sync_list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals:
return []
placeholders = ",".join("?" * len(principals))
rows = self._conn.execute(
f"""
SELECT DISTINCT collection FROM document_acl
WHERE permission = ?
AND principal_id IN ({placeholders})
ORDER BY collection
""",
(_READ_PERMISSION, *principals),
).fetchall()
return [r[0] for r in rows]
async def has_collection(self, collection: str) -> bool:
async with self._lock:
return await asyncio.to_thread(self._sync_has_collection, collection)
def _sync_has_collection(self, collection: str) -> bool:
row = self._conn.execute(
"SELECT 1 FROM documents_meta WHERE collection = ?",
(collection,),
).fetchone()
return row is not None
async def close(self) -> None:
async with self._lock:
await asyncio.to_thread(self._sync_close)
engine/src/stirling/documents/store.py
+75 -33
View File
@@ -2,8 +2,10 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from stirling.contracts.documents import Page, PageRange
from stirling.models import OwnerId, PrincipalId
@dataclass
@@ -40,73 +42,113 @@ class DocumentStore(ABC):
* **Vector chunks** - small, embedded chunks used for RAG search.
* **Ordered pages** - the original page text retained in document order,
used for whole-document reading.
Both representations live under the same ``collection`` (file id) and are
rooted at a single parent row in ``documents_meta``. Removing that parent
row cascades to both child representations, so :meth:`delete_collection`
is one logical delete.
"""
# ── lifecycle of the (collection, owner_id) row ────────────────────────
@abstractmethod
async def ensure_collection(self, collection: str, source: str) -> None:
"""Upsert the top-level ``documents_meta`` row for this collection.
async def ensure_collection(
self,
collection: str,
source: str,
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
"""Upsert the top-level ``documents_meta`` row for ``(collection, owner_id)``.
Must be called before :meth:`add_pages` or :meth:`add_documents`. Both
of those write into child tables that hold a foreign key to the parent
row, so it must exist first.
write into child tables that hold a foreign key to the parent row, so
it must exist first.
"""
@abstractmethod
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
"""Remove a collection's chunks, pages, and ACL rows.
Returns ``True`` when a row matched and was deleted, ``False`` when
``(collection, owner_id)`` didn't exist.
"""
@abstractmethod
async def purge_owner(self, owner_id: OwnerId) -> int:
"""Remove every collection (and ACL row) belonging to ``owner_id``.
Returns the number of collections purged.
"""
@abstractmethod
async def reap_expired(self) -> int:
"""Remove collections whose ``expires_at`` is non-null and in the past."""
# ── write paths (scoped by owner) ──────────────────────────────────────
@abstractmethod
async def add_documents(
self,
collection: str,
documents: list[Document],
embeddings: list[list[float]],
owner_id: OwnerId,
) -> None:
"""Store vector chunks with their embeddings in the named collection."""
"""Store vector chunks with their embeddings under the owner's collection."""
@abstractmethod
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
"""Replace the stored pages for ``(collection, owner_id)`` with the supplied pages."""
# ── ACL management ─────────────────────────────────────────────────────
@abstractmethod
async def grant_read(
self,
collection: str,
owner_id: OwnerId,
principals: list[PrincipalId],
) -> None:
"""Grant read access on ``(collection, owner_id)`` to each principal."""
@abstractmethod
async def revoke(
self,
collection: str,
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
"""Remove every permission this principal has on ``(collection, owner_id)``."""
# ── read paths (scoped by ACL principal set) ───────────────────────────
@abstractmethod
async def search(
self,
collection: str,
query_embedding: list[float],
top_k: int = 5,
top_k: int,
principals: list[PrincipalId],
) -> list[SearchResult]:
"""Return the top_k most similar vector chunks from the collection."""
"""Top-k similar vector chunks from ``collection`` that any principal can read.
@abstractmethod
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
"""Replace the stored pages for ``collection`` with the supplied pages.
Implementations must remove any previously-stored pages for the
collection before writing, so callers can re-ingest by calling this
method again.
Returns an empty list when no principal in ``principals`` has a read
ACL row for this collection (regardless of which owner backs it).
"""
@abstractmethod
async def read_pages(
self,
collection: str,
page_range: PageRange | None = None,
page_range: PageRange | None,
principals: list[PrincipalId],
) -> list[Page]:
"""Return ordered pages for ``collection``.
If ``page_range`` is ``None`` all pages are returned. Otherwise only
pages whose ``page_number`` falls within the inclusive range are
returned. Pages are always ordered by ``page_number`` ascending.
"""
"""Return ordered pages for ``collection``. Empty list when no read ACL match."""
@abstractmethod
async def delete_collection(self, collection: str) -> None:
"""Remove a collection's chunks and pages."""
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
"""Check whether any principal in ``principals`` can read this collection."""
@abstractmethod
async def list_collections(self) -> list[str]:
"""Return names of all existing collections."""
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
"""Return collection names readable by at least one of ``principals``."""
@abstractmethod
async def has_collection(self, collection: str) -> bool:
"""Check whether a collection exists."""
# ── lifecycle ──────────────────────────────────────────────────────────
@abstractmethod
async def close(self) -> None:
engine/src/stirling/models/__init__.py
+4 -1
View File
@@ -1,12 +1,15 @@
from . import tool_models
from .base import ApiModel, FileId
from .base import ApiModel, FileId, OwnerId, PrincipalId, UserId
from .tool_models import OPERATIONS, ParamToolModel, ToolEndpoint
__all__ = [
"ApiModel",
"FileId",
"OPERATIONS",
"OwnerId",
"ParamToolModel",
"PrincipalId",
"ToolEndpoint",
"UserId",
"tool_models",
]
engine/src/stirling/models/base.py
+10 -3
View File
@@ -5,11 +5,18 @@ from typing import NewType
from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel
# Stable, opaque identifier for a file supplied by the caller. Owned by the caller's
# ID strategy (content hash, filesystem path, etc.) and used as the RAG collection key
# throughout the engine.
# Stable, opaque identifier for a file supplied by the caller
FileId = NewType("FileId", str)
# Stable, opaque identifier for the calling user
UserId = NewType("UserId", str)
# Tenant that owns a document
OwnerId = NewType("OwnerId", str)
# An entity that can hold permissions on a document
PrincipalId = NewType("PrincipalId", str)
class ApiModel(BaseModel):
"""Base for every contract model crossing a service boundary."""
engine/src/stirling/services/__init__.py
+3 -1
View File
@@ -7,14 +7,16 @@ from .progress import (
set_progress_emitter,
)
from .runtime import AppRuntime, build_model_settings, build_runtime
from .tracking import setup_posthog_tracking
from .tracking import current_user_id, require_current_user_id, setup_posthog_tracking
__all__ = [
"AppRuntime",
"ProgressEmitter",
"build_model_settings",
"build_runtime",
"current_user_id",
"emit_progress",
"require_current_user_id",
"reset_progress_emitter",
"set_progress_emitter",
"setup_posthog_tracking",
engine/src/stirling/services/runtime.py
+4 -11
View File
@@ -16,7 +16,6 @@ from stirling.documents import (
DocumentStore,
EmbeddingService,
PgVectorStore,
RagCapability,
SqliteVecStore,
)
@@ -47,7 +46,6 @@ class AppRuntime:
fast_model: Model
smart_model: Model
documents: DocumentService
rag_capability: RagCapability
@property
def fast_model_settings(self) -> ModelSettings:
@@ -89,8 +87,8 @@ def _build_document_store(settings: AppSettings) -> DocumentStore:
assert_never(settings.rag_backend)
def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapability]:
"""Build the document service and the RAG-search capability that wraps it."""
def _build_documents(settings: AppSettings) -> DocumentService:
"""Build the document service used by per-request RAG capabilities."""
logger.info("Documents: embedding_model=%s", settings.rag_embedding_model)
embedder = EmbeddingService(
model_name=settings.rag_embedding_model,
@@ -98,9 +96,7 @@ def _build_documents(settings: AppSettings) -> tuple[DocumentService, RagCapabil
chunk_overlap=settings.rag_chunk_overlap,
)
store = _build_document_store(settings)
service = DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
capability = RagCapability(documents=service, top_k=settings.rag_default_top_k)
return service, capability
return DocumentService(embedder=embedder, store=store, default_top_k=settings.rag_default_top_k)
def build_runtime(settings: AppSettings) -> AppRuntime:
@@ -109,14 +105,11 @@ def build_runtime(settings: AppSettings) -> AppRuntime:
validate_structured_output_support(fast_model, settings.fast_model_name)
validate_structured_output_support(smart_model, settings.smart_model_name)
documents, rag_capability = _build_documents(settings)
return AppRuntime(
settings=settings,
fast_model=fast_model,
smart_model=smart_model,
documents=documents,
rag_capability=rag_capability,
documents=_build_documents(settings),
)
engine/src/stirling/services/tracking.py
+16 -1
View File
@@ -27,10 +27,25 @@ from opentelemetry.trace import Span
from posthog.client import Client as PostHogClient
from stirling.config import AppSettings
from stirling.models import UserId
# Per-request user ID, set by middleware from the X-User-Id header.
# When not set, PostHog generates a random ID and marks the event as personless.
current_user_id: ContextVar[str | None] = ContextVar("current_user_id", default=None)
current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None)
def require_current_user_id() -> UserId:
"""Return the request's user ID or raise if the X-User-Id header was missing.
Use at the boundary of any code path that touches per-user document
storage (vector chunks, page text, ACL rows). Routes prefer the FastAPI dependency form
(``Depends(require_user_id)``); agent internals that don't have a request
object in scope call this helper directly to fail closed.
"""
user_id = current_user_id.get()
if user_id is None:
raise RuntimeError("X-User-Id is required for this operation but was not set on the request")
return user_id
class LRUSet: