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
+343 -53
View File
@@ -8,7 +8,17 @@ from stirling.documents.rag_capability import RagCapability
from stirling.documents.service import DocumentService
from stirling.documents.sqlite_vec_store import SqliteVecStore
from stirling.documents.store import Document, SearchResult
from stirling.models import FileId
from stirling.models import FileId, OwnerId, PrincipalId
# Personal-doc tests reuse the same opaque string in all three roles — keeps the
# defaults of ingest() exercised and tests honest about what the same id means
# (caller, tenant, single ACL grantee). Org-doc tests construct ad-hoc owners
# and principal sets inline.
OWNER = OwnerId("test-user")
OWNER_PRINCIPALS = [PrincipalId("test-user")]
OTHER_OWNER = OwnerId("other-user")
OTHER_OWNER_PRINCIPALS = [PrincipalId("other-user")]
# chunk_text
@@ -56,7 +66,8 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_add_and_search(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("test-col", "test.pdf")
await store.ensure_collection("test-col", "test.pdf", OWNER, None)
await store.grant_read("test-col", OWNER, OWNER_PRINCIPALS)
docs = [
Document(id="1", text="Python is a programming language", metadata={"source": "test"}),
Document(id="2", text="Java is another programming language", metadata={"source": "test"}),
@@ -67,9 +78,9 @@ class TestSqliteVecStore:
[0.9, 0.1, 0.0],
[0.0, 0.0, 1.0],
]
await store.add_documents("test-col", docs, embeddings)
await store.add_documents("test-col", docs, embeddings, OWNER)
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2)
results = await store.search("test-col", [1.0, 0.05, 0.0], top_k=2, principals=OWNER_PRINCIPALS)
assert len(results) == 2
assert isinstance(results[0], SearchResult)
assert results[0].document.id == "1"
@@ -78,33 +89,36 @@ class TestSqliteVecStore:
@pytest.mark.anyio
async def test_list_and_has_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("my-collection", "test.pdf")
await store.ensure_collection("my-collection", "test.pdf", OWNER, None)
await store.grant_read("my-collection", OWNER, OWNER_PRINCIPALS)
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("my-collection", docs, [[1.0, 0.0]])
await store.add_documents("my-collection", docs, [[1.0, 0.0]], OWNER)
collections = await store.list_collections()
collections = await store.list_collections(OWNER_PRINCIPALS)
assert "my-collection" in collections
assert await store.has_collection("my-collection") is True
assert await store.has_collection("nonexistent") is False
assert await store.has_collection("my-collection", OWNER_PRINCIPALS) is True
assert await store.has_collection("nonexistent", OWNER_PRINCIPALS) is False
@pytest.mark.anyio
async def test_delete_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("to-delete", "test.pdf")
await store.ensure_collection("to-delete", "test.pdf", OWNER, None)
await store.grant_read("to-delete", OWNER, OWNER_PRINCIPALS)
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("to-delete", docs, [[1.0]])
await store.add_documents("to-delete", docs, [[1.0]], OWNER)
assert await store.has_collection("to-delete") is True
await store.delete_collection("to-delete")
assert await store.has_collection("to-delete") is False
assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is True
await store.delete_collection("to-delete", OWNER)
assert await store.has_collection("to-delete", OWNER_PRINCIPALS) is False
@pytest.mark.anyio
async def test_search_empty_collection(self) -> None:
store = SqliteVecStore.ephemeral()
await store.ensure_collection("empty-test", "test.pdf")
await store.ensure_collection("empty-test", "test.pdf", OWNER, None)
await store.grant_read("empty-test", OWNER, OWNER_PRINCIPALS)
docs = [Document(id="1", text="test", metadata={})]
await store.add_documents("empty-test", docs, [[1.0, 0.0]])
results = await store.search("empty-test", [1.0, 0.0], top_k=5)
await store.add_documents("empty-test", docs, [[1.0, 0.0]], OWNER)
results = await store.search("empty-test", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS)
assert len(results) == 1
@pytest.mark.anyio
@@ -112,7 +126,118 @@ class TestSqliteVecStore:
store = SqliteVecStore.ephemeral()
docs = [Document(id="1", text="test", metadata={})]
with pytest.raises(ValueError, match="documents.*embeddings"):
await store.add_documents("bad", docs, [[1.0], [2.0]])
await store.add_documents("bad", docs, [[1.0], [2.0]], OWNER)
@pytest.mark.anyio
async def test_collections_isolated_by_owner(self) -> None:
"""Two owners can store the same collection id; reads stay scoped to ACL."""
store = SqliteVecStore.ephemeral()
await store.ensure_collection("shared-id", "alice.pdf", OWNER, None)
await store.grant_read("shared-id", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("shared-id", "bob.pdf", OTHER_OWNER, None)
await store.grant_read("shared-id", OTHER_OWNER, OTHER_OWNER_PRINCIPALS)
await store.add_documents(
"shared-id",
[Document(id="1", text="alice content", metadata={})],
[[1.0, 0.0]],
OWNER,
)
await store.add_documents(
"shared-id",
[Document(id="1", text="bob content", metadata={})],
[[1.0, 0.0]],
OTHER_OWNER,
)
alice_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OWNER_PRINCIPALS)
bob_results = await store.search("shared-id", [1.0, 0.0], top_k=5, principals=OTHER_OWNER_PRINCIPALS)
assert [r.document.text for r in alice_results] == ["alice content"]
assert [r.document.text for r in bob_results] == ["bob content"]
assert await store.list_collections(OWNER_PRINCIPALS) == ["shared-id"]
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["shared-id"]
await store.delete_collection("shared-id", OWNER)
assert await store.has_collection("shared-id", OWNER_PRINCIPALS) is False
assert await store.has_collection("shared-id", OTHER_OWNER_PRINCIPALS) is True
@pytest.mark.anyio
async def test_purge_owner_removes_only_that_owners_collections(self) -> None:
"""Logout path: one owner's purge must not touch another owner's docs."""
store = SqliteVecStore.ephemeral()
for owner, principals, name in (
(OWNER, OWNER_PRINCIPALS, "a.pdf"),
(OWNER, OWNER_PRINCIPALS, "b.pdf"),
(OTHER_OWNER, OTHER_OWNER_PRINCIPALS, "c.pdf"),
):
await store.ensure_collection(name, name, owner, None)
await store.grant_read(name, owner, principals)
await store.add_documents(name, [Document(id="1", text="x", metadata={})], [[1.0, 0.0]], owner)
deleted = await store.purge_owner(OWNER)
assert deleted == 2
assert await store.list_collections(OWNER_PRINCIPALS) == []
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["c.pdf"]
@pytest.mark.anyio
async def test_reap_expired_drops_collections_past_expires_at(self) -> None:
"""TTL backstop: rows with ``expires_at`` in the past go away on reap."""
from datetime import UTC, datetime, timedelta
store = SqliteVecStore.ephemeral()
now = datetime.now(UTC)
await store.ensure_collection("fresh", "fresh.pdf", OWNER, now + timedelta(hours=1))
await store.grant_read("fresh", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("stale", "stale.pdf", OWNER, now - timedelta(seconds=1))
await store.grant_read("stale", OWNER, OWNER_PRINCIPALS)
deleted = await store.reap_expired()
assert deleted == 1
assert await store.list_collections(OWNER_PRINCIPALS) == ["fresh"]
@pytest.mark.anyio
async def test_reap_expired_keeps_persistent_and_unexpired(self) -> None:
"""Rows with ``expires_at`` in the future stay. Rows with null
``expires_at`` (org docs) are persistent and never touched, regardless
of age."""
from datetime import UTC, datetime, timedelta
store = SqliteVecStore.ephemeral()
await store.ensure_collection("session", "s.pdf", OWNER, datetime.now(UTC) + timedelta(hours=1))
await store.grant_read("session", OWNER, OWNER_PRINCIPALS)
await store.ensure_collection("persistent", "p.pdf", OTHER_OWNER, None)
await store.grant_read("persistent", OTHER_OWNER, OTHER_OWNER_PRINCIPALS)
deleted = await store.reap_expired()
assert deleted == 0
assert await store.list_collections(OWNER_PRINCIPALS) == ["session"]
assert await store.list_collections(OTHER_OWNER_PRINCIPALS) == ["persistent"]
@pytest.mark.anyio
async def test_acl_grants_read_to_extra_principal(self) -> None:
"""A principal without an ACL row can't read; once granted, they can."""
store = SqliteVecStore.ephemeral()
team_principal = PrincipalId("group:engineering")
await store.ensure_collection("doc", "engineering-runbook.pdf", OWNER, None)
await store.grant_read("doc", OWNER, OWNER_PRINCIPALS)
await store.add_documents(
"doc",
[Document(id="1", text="how to deploy", metadata={})],
[[1.0, 0.0]],
OWNER,
)
# Engineering can't see it yet.
assert await store.has_collection("doc", [team_principal]) is False
# Owner grants engineering read access.
await store.grant_read("doc", OWNER, [team_principal])
assert await store.has_collection("doc", [team_principal]) is True
# Revoke kills it.
await store.revoke("doc", OWNER, team_principal)
assert await store.has_collection("doc", [team_principal]) is False
# Owner still can.
assert await store.has_collection("doc", OWNER_PRINCIPALS) is True
# DocumentService (with stub embedder)
@@ -163,39 +288,131 @@ class TestDocumentService:
@pytest.mark.anyio
async def test_ingest_and_search(self, documents: DocumentService) -> None:
text = "Python is great for data science. It has many libraries like pandas and numpy."
count = await documents.ingest(FileId("docs"), _pages(text), source="guide.pdf")
count = await documents.ingest(
FileId("docs"),
_pages(text),
source="guide.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
assert count > 0
results = await documents.search("Python libraries", collection=FileId("docs"))
results = await documents.search("Python libraries", principals=OWNER_PRINCIPALS, collection=FileId("docs"))
assert len(results) > 0
assert results[0].document.text
@pytest.mark.anyio
async def test_ingest_empty_text_returns_zero_chunks(self, documents: DocumentService) -> None:
count = await documents.ingest(FileId("docs"), _pages(""), source="empty.pdf")
count = await documents.ingest(
FileId("docs"),
_pages(""),
source="empty.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
assert count == 0
@pytest.mark.anyio
async def test_search_nonexistent_collection_returns_empty(self, documents: DocumentService) -> None:
results = await documents.search("anything", collection=FileId("nonexistent"))
results = await documents.search("anything", principals=OWNER_PRINCIPALS, collection=FileId("nonexistent"))
assert results == []
@pytest.mark.anyio
async def test_search_all_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Machine learning overview."), source="ml.pdf")
await documents.ingest(FileId("col-b"), _pages("Deep learning with neural networks."), source="dl.pdf")
async def test_search_all_collections_for_principal(self, documents: DocumentService) -> None:
await documents.ingest(
FileId("col-a"),
_pages("Machine learning overview."),
source="ml.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Deep learning with neural networks."),
source="dl.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
results = await documents.search("neural networks")
results = await documents.search("neural networks", principals=OWNER_PRINCIPALS)
assert len(results) > 0
@pytest.mark.anyio
async def test_search_does_not_cross_owner_boundary(self, documents: DocumentService) -> None:
"""A search by one principal set never returns docs owned by an unrelated owner."""
await documents.ingest(
FileId("col-a"),
_pages("Alice's private notes."),
source="alice.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Bob's private notes."),
source="bob.pdf",
owner_id=OTHER_OWNER,
read_principals=OTHER_OWNER_PRINCIPALS,
expires_at=None,
)
alice_results = await documents.search("notes", principals=OWNER_PRINCIPALS)
bob_results = await documents.search("notes", principals=OTHER_OWNER_PRINCIPALS)
alice_texts = [r.document.text for r in alice_results]
bob_texts = [r.document.text for r in bob_results]
assert any("Alice" in t for t in alice_texts)
assert not any("Bob" in t for t in alice_texts)
assert any("Bob" in t for t in bob_texts)
assert not any("Alice" in t for t in bob_texts)
@pytest.mark.anyio
async def test_org_doc_visible_only_to_granted_group(self, documents: DocumentService) -> None:
"""Org-owned doc with explicit ACL: only granted principals see it."""
org_owner = OwnerId("org:acme")
eng_group = PrincipalId("group:engineering")
hr_group = PrincipalId("group:hr")
await documents.ingest(
FileId("runbook"),
_pages("Production deploy steps."),
source="runbook.pdf",
owner_id=org_owner,
read_principals=[eng_group],
expires_at=None,
)
# Engineering can read.
eng_results = await documents.search("deploy", principals=[eng_group])
assert len(eng_results) > 0
# HR cannot.
hr_results = await documents.search("deploy", principals=[hr_group])
assert hr_results == []
# Caller with both memberships (or org-wide principal) still sees it via eng.
multi_results = await documents.search("deploy", principals=[hr_group, eng_group])
assert len(multi_results) > 0
@pytest.mark.anyio
async def test_delete_collection(self, documents: DocumentService) -> None:
await documents.ingest(FileId("temp"), _pages("Temporary data."), source="tmp.pdf")
collections = await documents.list_collections()
await documents.ingest(
FileId("temp"),
_pages("Temporary data."),
source="tmp.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
collections = await documents.list_collections(OWNER_PRINCIPALS)
assert "temp" in collections
await documents.delete_collection(FileId("temp"))
collections = await documents.list_collections()
await documents.delete_collection(FileId("temp"), owner_id=OWNER)
collections = await documents.list_collections(OWNER_PRINCIPALS)
assert "temp" not in collections
@pytest.mark.anyio
@@ -205,9 +422,16 @@ class TestDocumentService:
PageText(page_number=2, text="Second page text."),
PageText(page_number=3, text="Third page text."),
]
await documents.ingest(FileId("ordered"), pages, source="ordered.pdf")
await documents.ingest(
FileId("ordered"),
pages,
source="ordered.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("ordered"))
stored = await documents.read_pages(FileId("ordered"), principals=OWNER_PRINCIPALS)
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[0].text == "First page text."
assert stored[0].char_count == len("First page text.")
@@ -217,9 +441,13 @@ class TestDocumentService:
from stirling.contracts import PageRange
pages = [PageText(page_number=i, text=f"page {i}") for i in range(1, 6)]
await documents.ingest(FileId("ranged"), pages, source="r.pdf")
await documents.ingest(
FileId("ranged"), pages, source="r.pdf", owner_id=OWNER, read_principals=OWNER_PRINCIPALS, expires_at=None
)
subset = await documents.read_pages(FileId("ranged"), PageRange(start=2, end=4))
subset = await documents.read_pages(
FileId("ranged"), principals=OWNER_PRINCIPALS, page_range=PageRange(start=2, end=4)
)
assert [p.page_number for p in subset] == [2, 3, 4]
@pytest.mark.anyio
@@ -228,14 +456,20 @@ class TestDocumentService:
FileId("doc"),
[PageText(page_number=1, text="old"), PageText(page_number=2, text="old2")],
source="v1.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("doc"),
[PageText(page_number=1, text="new")],
source="v2.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("doc"))
stored = await documents.read_pages(FileId("doc"), principals=OWNER_PRINCIPALS)
assert [p.page_number for p in stored] == [1]
assert stored[0].text == "new"
@@ -248,9 +482,16 @@ class TestDocumentService:
PageText(page_number=2, text=" "),
PageText(page_number=3, text="Real text on page 3."),
]
await documents.ingest(FileId("with-blanks"), pages, source="blanks.pdf")
await documents.ingest(
FileId("with-blanks"),
pages,
source="blanks.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
stored = await documents.read_pages(FileId("with-blanks"))
stored = await documents.read_pages(FileId("with-blanks"), principals=OWNER_PRINCIPALS)
assert [p.page_number for p in stored] == [1, 2, 3]
assert stored[1].text.strip() == ""
@@ -270,22 +511,36 @@ async def _invoke_search_knowledge(capability: RagCapability, query: str, max_re
class TestRagCapability:
def test_instructions_static_when_collections_pinned(self, documents: DocumentService) -> None:
cap = RagCapability(documents, collections=[FileId("docs"), FileId("manuals")])
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("docs"), FileId("manuals")])
instructions = cap.instructions
assert isinstance(instructions, str)
assert "docs, manuals" in instructions
assert "search_knowledge" in instructions
def test_instructions_dynamic_when_no_collections(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions = cap.instructions
assert callable(instructions)
@pytest.mark.anyio
async def test_dynamic_instructions_list_available_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("col-a"), _pages("Alpha content."), source="a.pdf")
await documents.ingest(FileId("col-b"), _pages("Beta content."), source="b.pdf")
cap = RagCapability(documents)
await documents.ingest(
FileId("col-a"),
_pages("Alpha content."),
source="a.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("col-b"),
_pages("Beta content."),
source="b.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -294,7 +549,7 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_dynamic_instructions_when_store_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
instructions_fn = cap.instructions
assert callable(instructions_fn)
text = await instructions_fn()
@@ -302,14 +557,21 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_returns_no_results_message_when_empty(self, documents: DocumentService) -> None:
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "anything")
assert output == "No relevant results found in the knowledge base."
@pytest.mark.anyio
async def test_search_knowledge_formats_results_with_source_and_score(self, documents: DocumentService) -> None:
await documents.ingest(FileId("docs"), _pages("Python is a programming language."), source="guide.pdf")
cap = RagCapability(documents)
await documents.ingest(
FileId("docs"),
_pages("Python is a programming language."),
source="guide.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "Python")
assert "[Result 1" in output
assert "source: guide.pdf" in output
@@ -318,10 +580,24 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_restricts_to_pinned_collections(self, documents: DocumentService) -> None:
await documents.ingest(FileId("pinned"), _pages("Pinned collection content."), source="pinned.pdf")
await documents.ingest(FileId("other"), _pages("Content in another collection."), source="other.pdf")
await documents.ingest(
FileId("pinned"),
_pages("Pinned collection content."),
source="pinned.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
await documents.ingest(
FileId("other"),
_pages("Content in another collection."),
source="other.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, collections=[FileId("pinned")])
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, collections=[FileId("pinned")])
output = await _invoke_search_knowledge(cap, "content")
assert "pinned.pdf" in output
assert "other.pdf" not in output
@@ -329,9 +605,16 @@ class TestRagCapability:
@pytest.mark.anyio
async def test_search_knowledge_respects_max_results(self, documents: DocumentService) -> None:
paragraphs = "\n\n".join(f"Paragraph {i} about topic." for i in range(10))
await documents.ingest(FileId("bulk"), _pages(paragraphs), source="bulk.pdf")
await documents.ingest(
FileId("bulk"),
_pages(paragraphs),
source="bulk.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS)
output = await _invoke_search_knowledge(cap, "topic", max_results=2)
assert "[Result 1" in output
assert "[Result 2" in output
@@ -341,8 +624,15 @@ class TestRagCapability:
async def test_search_knowledge_tool_is_hidden_after_budget_exhausted(self, documents: DocumentService) -> None:
"""The prepare callback must return None once max_searches has been reached
so the agent can no longer call the tool on subsequent turns."""
await documents.ingest(FileId("docs"), _pages("Some content."), source="x.pdf")
cap = RagCapability(documents, max_searches=2)
await documents.ingest(
FileId("docs"),
_pages("Some content."),
source="x.pdf",
owner_id=OWNER,
read_principals=OWNER_PRINCIPALS,
expires_at=None,
)
cap = RagCapability(documents, principals=OWNER_PRINCIPALS, max_searches=2)
tool_def = _dummy_tool_def()
assert await cap._prepare_search_knowledge(None, tool_def) is tool_def # type: ignore[arg-type]