diff --git a/pyproject.toml b/pyproject.toml index a96007b..c8f93ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "ratatoskr" -version = "0.17.2" +version = "0.17.3" description = "Worldtree Conversation API debug TUI — multi-pane observability dashboard" readme = "README.md" requires-python = ">=3.12" @@ -32,6 +32,7 @@ provider = [ "ratatoskr[web]", # reuse the starlette + uvicorn ASGI stack "bifrost>=0.6.1", # consumer engines + library (0.6.0 yanked: circular import) "jsonschema>=4", # bifrost runtime dep — envelope validation + "sqlite-vec>=0.1.6", # vector index for the memory plane (vec0 virtual table) ] dev = [ "pytest>=8", @@ -45,9 +46,10 @@ dev = [ ] [project.scripts] -ratatoskr = "ratatoskr.cli:main" -ratatoskr-web = "ratatoskr.web.entrypoint:main" -ratatoskr-provider = "ratatoskr.provider.serve:main" +ratatoskr = "ratatoskr.cli:main" +ratatoskr-web = "ratatoskr.web.entrypoint:main" +ratatoskr-provider = "ratatoskr.provider.serve:main" +ratatoskr-memory-provider = "ratatoskr.provider.serve_memory:main" [project.urls] Repository = "https://gitea.phasefinal.com/vh/ratatoskr" diff --git a/src/ratatoskr/provider/memory_store.py b/src/ratatoskr/provider/memory_store.py new file mode 100644 index 0000000..bc75ad7 --- /dev/null +++ b/src/ratatoskr/provider/memory_store.py @@ -0,0 +1,300 @@ +"""SQLite + sqlite-vec durable memory store (Bifrost memory plane, v1 basic plane). + +Contract: docs/contracts/bifrost_memory_provider.contract.md + +Worldtree writes Tier-3 agent memory chunks here (`upsert_many`) and recalls them +by vector similarity (`search`), with point reads (`get`/`get_many`) and deletes +(`delete_many`). We persist each chunk verbatim and read only its structural +surface — embedding (rank), scope (isolation), id + revision (optimistic lock), +origin/injection_source (consistency rule). Semantic content is never interpreted. +We implement bifrost's OWN MemoryDataStore Protocol; conformance is #195 parity +vs InMemoryMemoryStore. +""" +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import time +from typing import Any + +import sqlite_vec +from bifrost.consumer import ConsumerRegistration, build_memory_app +from bifrost.memory import ( + IdempotencyConflict, + InvalidArguments, + RevisionMismatch, + StoreCapabilities, +) +from bifrost.reference_server import JwtVerifier + +_SHORT_RETRY_TTL_SECONDS = 300 +_DURABLE_JOB_TTL_SECONDS = 24 * 60 * 60 + + +def _ctx_actor(ctx: Any) -> str: + """Reference `_ctx_actor`: actor = job_id | jwt_sub | session_id (never the record).""" + return str( + getattr(ctx, "job_id", None) + or getattr(ctx, "jwt_sub", None) + or getattr(ctx, "session_id", "") + ) + + +def _payload_digest(value: Any) -> str: + """Reference digest: sha256 of canonical JSON (sorted keys, compact, str-coerced).""" + blob = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(blob.encode("utf-8")).hexdigest() + + +def _idempotency_ttl_seconds(idempotency_class: str | None) -> int: + if idempotency_class == "durable-job": + return _DURABLE_JOB_TTL_SECONDS + return _SHORT_RETRY_TTL_SECONDS + + +def _chunk_id(record: dict) -> str: + """Reference `_chunk_id`: first non-empty of id / chunk_id / memory_id.""" + for key in ("id", "chunk_id", "memory_id"): + value = record.get(key) + if isinstance(value, str) and value: + return value + raise InvalidArguments("record missing id") + + +def _record_vector(record: dict) -> list[float]: + """Reference `_record_vector`: embedding, falling back to vector, else [].""" + value = record.get("embedding", record.get("vector", [])) + return [float(v) for v in value] if isinstance(value, list) else [] + + +def _scope_matches(record_scope: Any, scope_filter: dict) -> bool: + """INV-005: record is in-scope iff every scope_filter axis matches record["scope"].""" + if not isinstance(record_scope, dict): + return False + return all(record_scope.get(axis) == value for axis, value in scope_filter.items()) + + +def _validate_injection(record: dict) -> None: + """INV-007: injected_context requires injection_source; non-injected forbids it.""" + origin = record.get("origin") + injection_source = record.get("injection_source") + if origin == "injected_context" and not injection_source: + raise InvalidArguments("injected_context record requires injection_source") + if injection_source and origin != "injected_context": + raise InvalidArguments("injection_source only valid for injected_context origin") + + +class RatatoskrMemoryStore: + """The MemoryDataStore-shaped store handed to bifrost's build_memory_app.""" + + def __init__(self, conn: sqlite3.Connection, embedding_dim: int): + self._conn = conn + self._dim = embedding_dim + + def describe_store(self) -> dict: + # INV-006: advertise ONLY the v1 basic-plane capabilities (advertise-=>-implement). + return StoreCapabilities( + relational_edges_supported=False, + optimistic_locking_supported=True, + atomic_supersede_supported=False, + transaction_supported=False, + filterable_metadata_fields=[], + ).to_dict() + + async def upsert_many( + self, + records: list[dict], + *, + idempotency_key: str, + ctx: Any, + expected_revisions: dict | None = None, + idempotency_class: str | None = None, + ) -> dict: + if not (isinstance(idempotency_key, str) and idempotency_key): # PRE-001 + raise InvalidArguments("idempotency_key required") + # INV-002: idempotency_id = ("default", verb, actor-from-ctx, key); digest over payload. + digest = _payload_digest({"records": records, "expected_revisions": expected_revisions}) + idempotency_id = "|".join(("default", "upsert_many", _ctx_actor(ctx), idempotency_key)) + cached = self._conn.execute( + "SELECT digest FROM memory_idempotency WHERE idempotency_id = ?", + (idempotency_id,), + ).fetchone() + if cached is not None: + if cached[0] == digest: + return {"upserted": len(records), "replayed": True} + raise IdempotencyConflict("idempotency key reused with different payload") + for record in records: # INV-007 / PRE-002: validate before any write + _validate_injection(record) + # INV-004: all chunk rows + vec rows + the idempotency record in one transaction; + # a RevisionMismatch raised here rolls the whole batch back. + with self._conn: + if expected_revisions: # INV-003: optimistic lock, checked before any write + for record in records: + cid = _chunk_id(record) + if cid in expected_revisions: + row = self._conn.execute( + "SELECT revision FROM memory_chunks WHERE chunk_id = ?", (cid,) + ).fetchone() + current = row[0] if row is not None else 0 + if current != expected_revisions[cid]: + raise RevisionMismatch( + f"stale expected revision for {cid}: " + f"{expected_revisions[cid]} != {current}" + ) + for record in records: + chunk_id = _chunk_id(record) + # INV-003: first insert -> revision 1; re-upsert -> revision + 1. + self._conn.execute( + "INSERT INTO memory_chunks " + "(chunk_id, record_json, revision, scope_json, origin) " + "VALUES (?, ?, 1, ?, ?) " + "ON CONFLICT(chunk_id) DO UPDATE SET " + "record_json=excluded.record_json, revision=memory_chunks.revision + 1, " + "scope_json=excluded.scope_json, origin=excluded.origin", + ( + chunk_id, + json.dumps(record), + json.dumps(record.get("scope")), + record.get("origin"), + ), + ) + self._conn.execute("DELETE FROM memory_vec WHERE chunk_id = ?", (chunk_id,)) + self._conn.execute( + "INSERT INTO memory_vec(chunk_id, embedding) VALUES (?, ?)", + (chunk_id, sqlite_vec.serialize_float32(_record_vector(record))), + ) + self._conn.execute( + "INSERT INTO memory_idempotency (idempotency_id, digest, expires_at) " + "VALUES (?, ?, ?)", + (idempotency_id, digest, time.time() + _idempotency_ttl_seconds(idempotency_class)), + ) + return {"upserted": len(records), "replayed": False} + + async def search( + self, + vector: list[float], + *, + top_k: int, + scope_filter: dict | None = None, + metadata_filter: dict | None = None, + include: dict | None = None, + fidelity_target: Any = None, + ) -> list[dict]: + if len(vector) != self._dim: # PRE-001 + raise InvalidArguments(f"vector length {len(vector)} != embedding_dim {self._dim}") + if metadata_filter: # PRE-002: v1 advertises no filterable metadata fields + raise InvalidArguments("metadata_filter is unsupported in v1") + if scope_filter is not None and not isinstance(scope_filter, dict): # STEP 1 + raise InvalidArguments("scope_filter must be a flat {axis: value} dict") + if top_k <= 0: # POST-001: at most top_k + return [] + total = self._conn.execute("SELECT COUNT(*) FROM memory_vec").fetchone()[0] + if total == 0: + return [] + # Over-fetch every candidate ranked by cosine distance, then scope-filter and + # take top_k — so top_k counts IN-SCOPE hits (INV-005), not pre-filter hits. + rows = self._conn.execute( + "SELECT v.chunk_id, v.distance, c.record_json, c.revision " + "FROM memory_vec v JOIN memory_chunks c ON c.chunk_id = v.chunk_id " + f"WHERE v.embedding MATCH ? AND k = {total} ORDER BY v.distance", + (sqlite_vec.serialize_float32([float(x) for x in vector]),), + ).fetchall() + results: list[dict] = [] + for chunk_id, distance, record_json, revision in rows: + record = json.loads(record_json) + if scope_filter and not _scope_matches(record.get("scope"), scope_filter): + continue + results.append( + { + "chunk": record, + "chunk_id": chunk_id, + "score": 1.0 - distance, # vec0 cosine distance -> similarity + "recalled_view": record.get("distillate", record), + "revision": revision, + } + ) + if len(results) >= top_k: + break + return results + + async def get(self, chunk_id: str) -> dict | None: + # INV-001: verbatim round-trip + an attached revision key, or None. + row = self._conn.execute( + "SELECT record_json, revision FROM memory_chunks WHERE chunk_id = ?", + (chunk_id,), + ).fetchone() + if row is None: + return None + record = json.loads(row[0]) + record["revision"] = row[1] + return record + + async def get_many(self, ids: list[str]) -> list[dict]: + # List form of get: found records only (absent ids are skipped). + found = [] + for chunk_id in ids: + record = await self.get(chunk_id) + if record is not None: + found.append(record) + return found + + async def delete_many(self, ids: list[str]) -> dict: + # One transaction: chunk row + its vec row leave together (no orphan vec rows). + deleted = 0 + with self._conn: + for chunk_id in ids: + cur = self._conn.execute( + "DELETE FROM memory_chunks WHERE chunk_id = ?", (chunk_id,) + ) + if cur.rowcount > 0: + deleted += 1 + self._conn.execute("DELETE FROM memory_vec WHERE chunk_id = ?", (chunk_id,)) + return {"deleted": deleted} + + +def open_memory_store(db_path: str, *, embedding_dim: int) -> RatatoskrMemoryStore: + """Open the SQLite+sqlite-vec memory store, creating schema + the vec index on first use.""" + if not (isinstance(embedding_dim, int) and embedding_dim > 0): # PRE-002 + raise ValueError("embedding_dim must be a positive int") + conn = sqlite3.connect(db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + if db_path != ":memory:": + conn.execute("PRAGMA journal_mode=WAL") + conn.execute( + "CREATE TABLE IF NOT EXISTS memory_chunks (" + "chunk_id TEXT PRIMARY KEY, record_json TEXT NOT NULL, " + "revision INTEGER NOT NULL, scope_json TEXT, origin TEXT)" + ) + conn.execute( + "CREATE TABLE IF NOT EXISTS memory_idempotency (" + "idempotency_id TEXT PRIMARY KEY, digest TEXT NOT NULL, expires_at REAL)" + ) + # INV-005 / cosine recall: vec0 index at the fixed PINNED_EMBEDDER_DIM, cosine metric. + conn.execute( + "CREATE VIRTUAL TABLE IF NOT EXISTS memory_vec USING vec0(" + f"chunk_id TEXT PRIMARY KEY, embedding float[{embedding_dim}] distance_metric=cosine)" + ) + conn.commit() + return RatatoskrMemoryStore(conn, embedding_dim) + + +def build_memory_provider_app( + store: RatatoskrMemoryStore, + heimdall_key: bytes, + consumer_id: str = "ratatoskr", +): + """Wire the JWT verifier + registration and hand the store to bifrost. + + Returns a Starlette ASGI app exposing POST /bifrost/handshake and + POST /bifrost/memory-call. The library owns the wire; this is the thin glue. + """ + if not isinstance(store.describe_store(), dict): # PRE-001 / INV-008 + raise ValueError("store must advertise capabilities via describe_store()") + if not (isinstance(heimdall_key, bytes) and heimdall_key): # PRE-002 + raise ValueError("heimdall_key must be non-empty bytes") + verifier = JwtVerifier(algorithm="HS256", key_bytes=heimdall_key) + registration = ConsumerRegistration(consumer_id=consumer_id) + return build_memory_app(store=store, verifier=verifier, registration=registration) diff --git a/src/ratatoskr/provider/serve_memory.py b/src/ratatoskr/provider/serve_memory.py new file mode 100644 index 0000000..44c6dd1 --- /dev/null +++ b/src/ratatoskr/provider/serve_memory.py @@ -0,0 +1,60 @@ +"""Runnable entrypoint: serve the memory provider as an ASGI app. + +For the live negotiation smoke against a Worldtree instance. Config from env: +- RATATOSKR_HEIMDALL_KEY (required): HS256 shared key for the consumer, utf-8. +- RATATOSKR_MEMORY_EMBEDDING_DIM (required): the pinned embedder dim; the + sqlite-vec index is created at this fixed dim, so a wrong value silently + breaks search — no default. +- RATATOSKR_MEMORY_DB (default "memory.db"): SQLite path; ":memory:" = ephemeral. +- RATATOSKR_CONSUMER_ID (default "ratatoskr"). +- RATATOSKR_PROVIDER_HOST (default "0.0.0.0"), + RATATOSKR_MEMORY_PROVIDER_PORT (default 8391 — distinct from affect's 8390 so + both planes can run side-by-side as separate apps, per the v1 contract). +""" +from __future__ import annotations + +import os +from collections.abc import Mapping + +from ratatoskr.provider.memory_store import build_memory_provider_app, open_memory_store + + +def build_memory_app_from_env(env: Mapping[str, str] | None = None): + """Build the memory ASGI app from environment config (testable seam).""" + env = os.environ if env is None else env + key = env.get("RATATOSKR_HEIMDALL_KEY") + if not key: + raise RuntimeError( + "RATATOSKR_HEIMDALL_KEY is required to serve the memory provider" + ) + raw_dim = env.get("RATATOSKR_MEMORY_EMBEDDING_DIM") + if not raw_dim: + raise RuntimeError( + "RATATOSKR_MEMORY_EMBEDDING_DIM is required (Worldtree's PINNED_EMBEDDER_DIM)" + ) + try: + embedding_dim = int(raw_dim) + except ValueError as exc: + raise RuntimeError( + f"RATATOSKR_MEMORY_EMBEDDING_DIM must be an int, got {raw_dim!r}" + ) from exc + if embedding_dim <= 0: + raise RuntimeError("RATATOSKR_MEMORY_EMBEDDING_DIM must be a positive int") + store = open_memory_store( + env.get("RATATOSKR_MEMORY_DB", "memory.db"), embedding_dim=embedding_dim + ) + return build_memory_provider_app( + store, + heimdall_key=key.encode(), + consumer_id=env.get("RATATOSKR_CONSUMER_ID", "ratatoskr"), + ) + + +def main() -> None: + import uvicorn + + uvicorn.run( + build_memory_app_from_env(), + host=os.environ.get("RATATOSKR_PROVIDER_HOST", "0.0.0.0"), + port=int(os.environ.get("RATATOSKR_MEMORY_PROVIDER_PORT", "8391")), + ) diff --git a/tests/test_provider_memory.py b/tests/test_provider_memory.py new file mode 100644 index 0000000..700d763 --- /dev/null +++ b/tests/test_provider_memory.py @@ -0,0 +1,414 @@ +"""Tests for the Tier-3 Bifrost memory provider (ratatoskr.provider.memory_store). + +Contract: docs/contracts/bifrost_memory_provider.contract.md (v1.1) +Vertical tracer-first: fresh_db -> basic_upsert (round-trip) -> replay -> conflict +-> optimistic_lock -> injection_rule -> search/scope_isolation -> get/get_many -> +delete_many -> build_memory_provider_app -> #195 parity vs InMemoryMemoryStore. +""" +from __future__ import annotations + +import types + +import pytest +from bifrost.memory import IdempotencyConflict, InvalidArguments, RevisionMismatch + +from ratatoskr.provider.memory_store import ( + build_memory_provider_app, + open_memory_store, +) + +EMBEDDING_DIM = 8 + + +def _ctx(sub: str = "sub-1"): + # Mirrors bifrost reference _ctx_actor: actor = job_id | jwt_sub | session_id. + return types.SimpleNamespace(jwt_sub=sub) + + +def _vec(*head: float) -> list[float]: + v = list(head) + [0.0] * EMBEDDING_DIM + return v[:EMBEDDING_DIM] + + +def _chunk(cid: str = "c1", *, embedding=None, scope=None, **extra) -> dict: + rec = { + "id": cid, + "embedding": embedding if embedding is not None else _vec(1.0), + "scope": scope if scope is not None else {"end_user": "u1"}, + "origin": "worldtree", + "distillate": {"summary": f"distillate-{cid}"}, + "content": f"content-{cid}", + } + rec.update(extra) + return rec + + +def _row_count(store, table: str) -> int: + return store._conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + + +# --- open_memory_store --- + +def test_fresh_db_advertises_v1_caps_and_schema(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + caps = store.describe_store() + assert caps["relational_edges_supported"] is False + assert caps["optimistic_locking_supported"] is True + assert caps["atomic_supersede_supported"] is False + assert caps["transaction_supported"] is False + assert caps["filterable_metadata_fields"] == [] + # tables + vec index queryable + store._conn.execute("SELECT * FROM memory_chunks") + store._conn.execute("SELECT * FROM memory_idempotency") + store._conn.execute("SELECT * FROM memory_vec") + + +def test_reopen_existing_file_is_idempotent(tmp_path): + db = str(tmp_path / "memory.db") + open_memory_store(db, embedding_dim=EMBEDDING_DIM) # first open creates schema + store = open_memory_store(db, embedding_dim=EMBEDDING_DIM) # reopen: IF NOT EXISTS no-op + assert isinstance(store.describe_store(), dict) + store._conn.execute("SELECT * FROM memory_chunks") + store._conn.execute("SELECT * FROM memory_vec") + + +# --- upsert_many + get (tracer round-trip) --- + +async def test_basic_upsert_round_trips_verbatim_with_revision(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + c1 = _chunk("c1", embedding=_vec(1.0)) + c2 = _chunk("c2", embedding=_vec(0.0, 1.0)) + result = await store.upsert_many([c1, c2], idempotency_key="k1", ctx=_ctx()) + assert result == {"upserted": 2, "replayed": False} + # INV-001: each chunk round-trips verbatim, with a revision key attached (first insert -> 1) + assert await store.get("c1") == {**c1, "revision": 1} + assert await store.get("c2") == {**c2, "revision": 1} + + +async def test_replay_same_key_same_payload_no_rewrite(): + # INV-002: same idempotency_key + same digest -> replay (no second write, revision frozen) + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + c1 = _chunk("c1") + assert await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) == { + "upserted": 1, + "replayed": False, + } + assert await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) == { + "upserted": 1, + "replayed": True, + } + assert (await store.get("c1"))["revision"] == 1 # replay did not re-write / re-increment + assert _row_count(store, "memory_chunks") == 1 + + +async def test_conflict_same_key_different_payload_raises_and_keeps_first(): + # INV-002: same key, different digest -> IdempotencyConflict; the first batch is intact + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + first = _chunk("c1", content="first") + await store.upsert_many([first], idempotency_key="k1", ctx=_ctx()) + with pytest.raises(IdempotencyConflict): + await store.upsert_many( + [_chunk("c1", content="second")], idempotency_key="k1", ctx=_ctx() + ) + assert await store.get("c1") == {**first, "revision": 1} # untouched + + +async def test_optimistic_lock_stale_expected_revision_raises_nothing_written(): + # INV-003: a stale expected_revisions entry rolls back the whole batch + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + c1 = _chunk("c1", content="v1") + await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) # revision 1 + with pytest.raises(RevisionMismatch): + await store.upsert_many( + [_chunk("c1", content="v2")], + idempotency_key="k2", # distinct key: not replay/conflict + ctx=_ctx(), + expected_revisions={"c1": 5}, # stale: stored revision is 1 + ) + assert await store.get("c1") == {**c1, "revision": 1} # nothing written + assert _row_count(store, "memory_chunks") == 1 + + +async def test_optimistic_lock_match_upserts_and_increments_revision(): + # INV-003: a matching expected_revisions writes and increments (1 -> 2) + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + await store.upsert_many([_chunk("c1", content="v1")], idempotency_key="k1", ctx=_ctx()) + v2 = _chunk("c1", content="v2") + assert await store.upsert_many( + [v2], idempotency_key="k2", ctx=_ctx(), expected_revisions={"c1": 1} + ) == {"upserted": 1, "replayed": False} + assert await store.get("c1") == {**v2, "revision": 2} # re-upsert increments + assert _row_count(store, "memory_chunks") == 1 + + +async def test_injection_rule_injected_without_source_raises_no_write(): + # INV-007: origin == injected_context requires injection_source + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + bad = _chunk("c1", origin="injected_context") # no injection_source + with pytest.raises(InvalidArguments): + await store.upsert_many([bad], idempotency_key="k1", ctx=_ctx()) + assert _row_count(store, "memory_chunks") == 0 + + +async def test_injection_rule_non_injected_with_source_raises_no_write(): + # INV-007: a non-injected record carrying injection_source is rejected + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + bad = _chunk("c1", origin="worldtree", injection_source="elsewhere") + with pytest.raises(InvalidArguments): + await store.upsert_many([bad], idempotency_key="k1", ctx=_ctx()) + assert _row_count(store, "memory_chunks") == 0 + + +# --- search --- + +async def test_basic_search_ranks_by_cosine_with_recalled_view(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + scope = {"end_user": "u1"} + c1 = _chunk("c1", embedding=_vec(1.0, 0.0), scope=scope) + c3 = _chunk("c3", embedding=_vec(0.9, 0.1), scope=scope) + await store.upsert_many( + [c1, _chunk("c2", embedding=_vec(0.0, 1.0), scope=scope), c3], + idempotency_key="k1", + ctx=_ctx(), + ) + results = await store.search(_vec(1.0, 0.0), top_k=2, scope_filter=scope) + assert [r["chunk_id"] for r in results] == ["c1", "c3"] # nearest to [1,0] by cosine + top = results[0] + assert top["chunk"] == c1 # verbatim chunk, no revision attached + assert top["recalled_view"] == {"summary": "distillate-c1"} # = chunk["distillate"] + assert top["revision"] == 1 + assert isinstance(top["score"], float) + + +async def test_scope_isolation_excludes_other_scope_even_if_closer(): + # INV-005: an out-of-scope chunk that scores HIGHER must not leak; only in-scope returned + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + await store.upsert_many( + [ + _chunk("u2-near", embedding=_vec(1.0, 0.0), scope={"end_user": "u2"}), # closest + _chunk("u1-far", embedding=_vec(0.0, 1.0), scope={"end_user": "u1"}), # in-scope, far + ], + idempotency_key="k1", + ctx=_ctx(), + ) + results = await store.search(_vec(1.0, 0.0), top_k=2, scope_filter={"end_user": "u1"}) + assert [r["chunk_id"] for r in results] == ["u1-far"] # u2-near excluded despite ranking first + + +async def test_search_empty_store_returns_empty(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + assert await store.search(_vec(1.0), top_k=5) == [] + + +async def test_search_non_empty_metadata_filter_rejected(): + # PRE-002: v1 advertises no filterable metadata fields + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + with pytest.raises(InvalidArguments): + await store.search(_vec(1.0), top_k=5, metadata_filter={"x": 1}) + + +async def test_search_wrong_vector_dim_rejected(): + # PRE-001: vector length must equal the pinned embedding_dim + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + with pytest.raises(InvalidArguments): + await store.search([1.0, 0.0], top_k=5) + + +async def test_search_non_dict_scope_filter_rejected(): + # search STEP 1: scope_filter must be a flat {axis: value} dict + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + with pytest.raises(InvalidArguments): + await store.search(_vec(1.0), top_k=5, scope_filter="u1") + + +async def test_search_top_k_zero_returns_empty(): + # POST-001: at most top_k — zero means zero + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + await store.upsert_many([_chunk("c1")], idempotency_key="k1", ctx=_ctx()) + assert await store.search(_vec(1.0), top_k=0, scope_filter={"end_user": "u1"}) == [] + + +async def test_scope_isolation_fills_top_k_from_in_scope_past_higher_out_of_scope(): + # INV-005: top_k counts IN-SCOPE hits. An out-of-scope chunk ranking #1 is skipped, + # and top_k is still filled from the in-scope set when enough in-scope chunks exist. + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + await store.upsert_many( + [ + _chunk("u2-nearest", embedding=_vec(1.0, 0.0), scope={"end_user": "u2"}), # ranks #1 + _chunk("u1-near", embedding=_vec(0.95, 0.05), scope={"end_user": "u1"}), + _chunk("u1-mid", embedding=_vec(0.8, 0.2), scope={"end_user": "u1"}), + _chunk("u1-far", embedding=_vec(0.0, 1.0), scope={"end_user": "u1"}), + ], + idempotency_key="k1", + ctx=_ctx(), + ) + results = await store.search(_vec(1.0, 0.0), top_k=2, scope_filter={"end_user": "u1"}) + # exactly top_k in-scope (the 2 nearest u1 chunks); the higher-ranked u2 chunk is excluded + assert [r["chunk_id"] for r in results] == ["u1-near", "u1-mid"] + + +# --- get / get_many --- + +async def test_get_absent_returns_none(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + assert await store.get("nope") is None + + +async def test_get_many_returns_found_records_only(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + c1 = _chunk("c1") + await store.upsert_many([c1], idempotency_key="k1", ctx=_ctx()) + assert await store.get_many(["c1", "absent"]) == [{**c1, "revision": 1}] + + +# --- delete_many --- + +async def test_delete_hit_removes_chunk_and_vec_row(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + await store.upsert_many([_chunk("c1"), _chunk("c2")], idempotency_key="k1", ctx=_ctx()) + assert await store.delete_many(["c1"]) == {"deleted": 1} + assert await store.get("c1") is None + assert _row_count(store, "memory_chunks") == 1 + assert _row_count(store, "memory_vec") == 1 # c1's vec row gone too (no orphan) + # delete_hit: search no longer surfaces it (vec/chunk coupling held) + hits = await store.search(_vec(1.0), top_k=5, scope_filter={"end_user": "u1"}) + assert all(r["chunk_id"] != "c1" for r in hits) + + +async def test_delete_absent_counts_zero(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + assert await store.delete_many(["nope"]) == {"deleted": 0} + + +# --- build_memory_provider_app --- + +def test_build_app_exposes_handshake_and_memory_routes(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + app = build_memory_provider_app(store, heimdall_key=b"secret-key") + routes = {getattr(r, "path", None): r for r in app.routes} + assert "/bifrost/handshake" in routes + assert "/bifrost/memory-call" in routes + assert "POST" in routes["/bifrost/memory-call"].methods + assert "POST" in routes["/bifrost/handshake"].methods # both routes are POST (incl. POST) + + +def test_build_app_rejects_empty_key(): + store = open_memory_store(":memory:", embedding_dim=EMBEDDING_DIM) + with pytest.raises(ValueError): + build_memory_provider_app(store, heimdall_key=b"") + + +# --- #195 conformance: parity vs the reference store through the real engine --- + +def _dispatch_ctx(*scopes: str, session_id: str = "actor-1"): + return types.SimpleNamespace( + scope=list(scopes), session_id=session_id, jwt_sub=session_id, job_id=None + ) + + +def _ref_record(chunk_id: str, vector: list[float], *, end_user: str = "u1") -> dict: + # Mirrors bifrost's reference `record` helper so the envelope validates. + return { + "id": chunk_id, + "embedding": vector, + "distillate": {"text": chunk_id}, + "metadata": {"worldtree.appraisal_confidence": 0.8}, + "scope": {"end_user": end_user, "tenant": "t1"}, + "origin": "worldtree", + "source_role": "assistant", + "trust_tier": "tier-3", + "provenance": {"trace": chunk_id}, + } + + +async def test_parity_upsert_many_vs_reference_through_dispatch(): + from bifrost.consumer.testing import InMemoryMemoryStore + from bifrost.memory import dispatch_memory_call + + ref = InMemoryMemoryStore() + mine = open_memory_store(":memory:", embedding_dim=2) + wctx = _dispatch_ctx("memory:write") + env = { + "operation": "upsert_many", + "args": {"records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])]}, + "idempotency_key": "k1", + } + # happy persist + replay: wire bodies must agree + assert await dispatch_memory_call(env, wctx, ref) == await dispatch_memory_call(env, wctx, mine) + assert await dispatch_memory_call(env, wctx, ref) == await dispatch_memory_call(env, wctx, mine) + + +async def test_parity_search_ranked_ids_vs_reference_through_dispatch(): + from bifrost.consumer.testing import InMemoryMemoryStore + from bifrost.memory import dispatch_memory_call + + ref = InMemoryMemoryStore() + mine = open_memory_store(":memory:", embedding_dim=2) + wctx = _dispatch_ctx("memory:write") + rctx = _dispatch_ctx("memory:read") + up = { + "operation": "upsert_many", + "args": { + "records": [ + _ref_record("a", [1.0, 0.0]), + _ref_record("b", [0.0, 1.0]), + _ref_record("c", [0.9, 0.1]), + ] + }, + "idempotency_key": "k1", + } + await dispatch_memory_call(up, wctx, ref) + await dispatch_memory_call(up, wctx, mine) + search_env = { + "operation": "search", + "args": {"vector": [1.0, 0.0], "top_k": 2, "scope_filter": {"end_user": "u1"}}, + } + rstatus, rbody = await dispatch_memory_call(search_env, rctx, ref) + mstatus, mbody = await dispatch_memory_call(search_env, rctx, mine) + assert rstatus == mstatus == 200 + # #195: same ranked chunk_ids and the same per-result shape (scores may differ in the + # last float digit between vec0's cosine and the reference's Python cosine). + assert [r["chunk_id"] for r in rbody["results"]] == [r["chunk_id"] for r in mbody["results"]] + assert set(rbody["results"][0]) == set(mbody["results"][0]) + + +async def test_parity_expected_revisions_vs_reference_through_dispatch(): + # #195: pins the partial-map optimistic-lock semantics against the reference + # (does an expected_revisions map that omits some batch records lock only the + # listed ones?). Resolves the contract's ambiguous "each record's stored revision". + from bifrost.consumer.testing import InMemoryMemoryStore + from bifrost.memory import dispatch_memory_call + + ref = InMemoryMemoryStore() + mine = open_memory_store(":memory:", embedding_dim=2) + wctx = _dispatch_ctx("memory:write") + + seed = { + "operation": "upsert_many", + "args": {"records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])]}, + "idempotency_key": "seed", + } + assert await dispatch_memory_call(seed, wctx, ref) == await dispatch_memory_call(seed, wctx, mine) + + # partial map: only "a" is locked (revision 1); "b" is omitted from expected_revisions + partial = { + "operation": "upsert_many", + "args": { + "records": [_ref_record("a", [1.0, 0.0]), _ref_record("b", [0.0, 1.0])], + "expected_revisions": {"a": 1}, + }, + "idempotency_key": "partial", + } + assert await dispatch_memory_call(partial, wctx, ref) == await dispatch_memory_call( + partial, wctx, mine + ) + + # stale lock: both map to the same RevisionMismatch wire error + stale = { + "operation": "upsert_many", + "args": {"records": [_ref_record("a", [1.0, 0.0])], "expected_revisions": {"a": 99}}, + "idempotency_key": "stale", + } + assert await dispatch_memory_call(stale, wctx, ref) == await dispatch_memory_call( + stale, wctx, mine + ) diff --git a/tests/test_provider_serve_memory.py b/tests/test_provider_serve_memory.py new file mode 100644 index 0000000..3bd53ae --- /dev/null +++ b/tests/test_provider_serve_memory.py @@ -0,0 +1,48 @@ +"""Tests for the memory-provider serve entrypoint (ratatoskr.provider.serve_memory). + +Only the env -> app seam is unit-tested; uvicorn.run is the untestable shell. +""" +from __future__ import annotations + +import pytest + +from ratatoskr.provider.serve_memory import build_memory_app_from_env + + +def test_build_memory_app_from_env_requires_heimdall_key(): + with pytest.raises(RuntimeError): + build_memory_app_from_env( + {"RATATOSKR_MEMORY_DB": ":memory:", "RATATOSKR_MEMORY_EMBEDDING_DIM": "8"} + ) + + +def test_build_memory_app_from_env_requires_embedding_dim(): + # A wrong/missing dim silently breaks vector search -> require it explicitly. + with pytest.raises(RuntimeError): + build_memory_app_from_env( + {"RATATOSKR_HEIMDALL_KEY": "k", "RATATOSKR_MEMORY_DB": ":memory:"} + ) + + +def test_build_memory_app_from_env_rejects_non_positive_dim(): + with pytest.raises(RuntimeError): + build_memory_app_from_env( + { + "RATATOSKR_HEIMDALL_KEY": "k", + "RATATOSKR_MEMORY_DB": ":memory:", + "RATATOSKR_MEMORY_EMBEDDING_DIM": "0", + } + ) + + +def test_build_memory_app_from_env_builds_app_with_routes(): + app = build_memory_app_from_env( + { + "RATATOSKR_HEIMDALL_KEY": "shared-secret", + "RATATOSKR_MEMORY_DB": ":memory:", + "RATATOSKR_MEMORY_EMBEDDING_DIM": "8", + } + ) + paths = {getattr(r, "path", None) for r in app.routes} + assert "/bifrost/handshake" in paths + assert "/bifrost/memory-call" in paths diff --git a/uv.lock b/uv.lock index 17715b6..804386d 100644 --- a/uv.lock +++ b/uv.lock @@ -1052,7 +1052,7 @@ wheels = [ [[package]] name = "ratatoskr" -version = "0.17.2" +version = "0.17.3" source = { editable = "." } dependencies = [ { name = "httpx" }, @@ -1075,6 +1075,7 @@ dev = [ provider = [ { name = "bifrost" }, { name = "jsonschema" }, + { name = "sqlite-vec" }, { name = "starlette" }, { name = "uvicorn", extra = ["standard"] }, ] @@ -1097,6 +1098,7 @@ requires-dist = [ { name = "ratatoskr", extras = ["web"], marker = "extra == 'provider'" }, { name = "respx", marker = "extra == 'dev'", specifier = ">=0.21" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.6" }, + { name = "sqlite-vec", marker = "extra == 'provider'", specifier = ">=0.1.6" }, { name = "starlette", marker = "extra == 'web'", specifier = ">=0.40" }, { name = "textual", specifier = ">=0.85" }, { name = "textual-dev", marker = "extra == 'dev'", specifier = ">=1.5" }, @@ -1278,6 +1280,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9b/36/9c015cd052fca743dae8cb2aeb16b551444787467db42ceab0fc968865af/ruff-0.15.13-py3-none-win_arm64.whl", hash = "sha256:2471da9bd1068c8c064b5fd9c0c4b6dddffd6369cb1cd68b29993b1709ff1b21", size = 11179336 }, ] +[[package]] +name = "sqlite-vec" +version = "0.1.9" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/85/9fad0045d8e7c8df3e0fa5a56c630e8e15ad6e5ca2e6106fceb666aa6638/sqlite_vec-0.1.9-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:1b62a7f0a060d9475575d4e599bbf94a13d85af896bc1ce86ee80d1b5b48e5fb", size = 131171 }, + { url = "https://files.pythonhosted.org/packages/a4/3d/3677e0cd2f92e5ebc43cd29fbf565b75582bff1ccfa0b8327c7508e1084f/sqlite_vec-0.1.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1d52e30513bae4cc9778ddbf6145610434081be4c3afe57cd877893bad9f6b6c", size = 165434 }, + { url = "https://files.pythonhosted.org/packages/00/d4/f2b936d3bdc38eadcbd2a87875815db36430fab0363182ba5d12cd8e0b51/sqlite_vec-0.1.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e921e592f24a5f9a18f590b6ddd530eb637e2d474e3b1972f9bbeb773aa3cb9", size = 160076 }, + { url = "https://files.pythonhosted.org/packages/6f/ad/6afd073b0f817b3e03f9e37ad626ae341805891f23c74b5292818f49ac63/sqlite_vec-0.1.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:1515727990b49e79bcaf75fdee2ffc7d461f8b66905013231251f1c8938e7786", size = 163388 }, + { url = "https://files.pythonhosted.org/packages/42/89/81b2907cda14e566b9bf215e2ad82fc9b349edf07d2010756ffdb902f328/sqlite_vec-0.1.9-py3-none-win_amd64.whl", hash = "sha256:4a28dc12fa4b53d7b1dced22da2488fade444e96b5d16fd2d698cd670675cf32", size = 292804 }, +] + [[package]] name = "starlette" version = "1.1.0"