James BruntonandGitHub
1264f4cfed
# 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.
550 lines
21 KiB
Python
550 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
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, 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:
|
|
is_memory = str(db_path) == ":memory:"
|
|
self._db_path: Path | None = None if is_memory else Path(db_path)
|
|
|
|
if self._db_path is not None:
|
|
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(str(self._db_path), check_same_thread=False)
|
|
else:
|
|
conn = sqlite3.connect(":memory:", check_same_thread=False)
|
|
|
|
conn.enable_load_extension(True)
|
|
sqlite_vec.load(conn)
|
|
conn.enable_load_extension(False)
|
|
# Required so cascade deletes from documents_meta clean up child tables.
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
if self._db_path is not None:
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
|
|
self._conn = conn
|
|
self._lock = asyncio.Lock()
|
|
self._init_schema()
|
|
|
|
@classmethod
|
|
def ephemeral(cls) -> SqliteVecStore:
|
|
"""In-memory store for testing."""
|
|
return cls(":memory:")
|
|
|
|
def _init_schema(self) -> None:
|
|
self._conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS documents_meta (
|
|
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 (
|
|
collection TEXT NOT NULL,
|
|
owner_id TEXT NOT NULL,
|
|
dim INTEGER 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
|
|
)
|
|
"""
|
|
)
|
|
self._conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS documents (
|
|
id TEXT NOT NULL,
|
|
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, 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_owner ON documents(collection, owner_id)")
|
|
self._conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS document_pages (
|
|
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, 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_owner ON document_pages(collection, owner_id)"
|
|
)
|
|
self._conn.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
|
|
)
|
|
"""
|
|
)
|
|
# 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, 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]:
|
|
norm = math.sqrt(sum(x * x for x in vector))
|
|
if norm == 0:
|
|
return list(vector)
|
|
return [x / norm for x in vector]
|
|
|
|
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")
|
|
if not documents:
|
|
return
|
|
|
|
async with self._lock:
|
|
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 collection = ? AND owner_id = ?",
|
|
(collection, owner_id),
|
|
).fetchone()
|
|
if row is None:
|
|
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(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} 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 owner_id = ? AND id IN ({placeholders})",
|
|
(collection, owner_id, *ids),
|
|
).fetchall()
|
|
if existing:
|
|
vec_rowids = [r[0] for r in existing]
|
|
row_placeholders = ",".join("?" * len(vec_rowids))
|
|
self._conn.execute(
|
|
f"DELETE FROM {table_name} WHERE rowid IN ({row_placeholders})",
|
|
vec_rowids,
|
|
)
|
|
self._conn.execute(
|
|
f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
|
|
(collection, owner_id, *ids),
|
|
)
|
|
|
|
for doc, emb in zip(documents, embeddings):
|
|
normalized = self._normalize(list(emb))
|
|
cursor = self._conn.execute(
|
|
f"INSERT INTO {table_name}(embedding) VALUES (?)",
|
|
(sqlite_vec.serialize_float32(normalized),),
|
|
)
|
|
vec_rowid = cursor.lastrowid
|
|
self._conn.execute(
|
|
"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,
|
|
principals: list[PrincipalId],
|
|
) -> list[SearchResult]:
|
|
async with self._lock:
|
|
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]:
|
|
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
|
|
if len(query_embedding) != dim:
|
|
raise ValueError(f"Query embedding dim {len(query_embedding)} does not match collection dim {dim}")
|
|
|
|
normalized = self._normalize(list(query_embedding))
|
|
query_blob = sqlite_vec.serialize_float32(normalized)
|
|
|
|
results = self._conn.execute(
|
|
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 = ?
|
|
AND d.owner_id = ?
|
|
WHERE v.embedding MATCH ? AND k = ?
|
|
ORDER BY v.distance
|
|
""",
|
|
(collection, owner_id, query_blob, top_k),
|
|
).fetchall()
|
|
|
|
return [
|
|
SearchResult(
|
|
document=Document(
|
|
id=r[0],
|
|
text=r[1],
|
|
metadata=json.loads(r[2]) if r[2] else {},
|
|
),
|
|
# For normalized vectors: cosine_sim = 1 - (L2^2 / 2)
|
|
score=max(0.0, 1.0 - (r[3] ** 2) / 2.0),
|
|
)
|
|
for r in results
|
|
]
|
|
|
|
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,
|
|
principals: list[PrincipalId],
|
|
) -> list[Page]:
|
|
async with self._lock:
|
|
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 = ? 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 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 has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
|
|
async with self._lock:
|
|
return await asyncio.to_thread(self._sync_has_collection, collection, principals)
|
|
|
|
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, principals: list[PrincipalId]) -> list[str]:
|
|
async with self._lock:
|
|
return await asyncio.to_thread(self._sync_list_collections, principals)
|
|
|
|
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 close(self) -> None:
|
|
async with self._lock:
|
|
await asyncio.to_thread(self._sync_close)
|
|
|
|
def _sync_close(self) -> None:
|
|
"""Checkpoint the WAL into the main database file and close the connection so
|
|
the .db-shm and .db-wal files are cleaned up on graceful shutdown."""
|
|
if self._db_path is not None:
|
|
try:
|
|
self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
|
self._conn.commit()
|
|
except sqlite3.Error:
|
|
# Best effort: if checkpointing fails we still want to close the connection.
|
|
pass
|
|
self._conn.close()
|