ZacharyZhang-NY
83dfc14daa
A holder is now somebody who can prove it. Schema v11 stores a sha256 of a per-holder token beside the directory row, so a copied database cannot replay it and the plaintext is shown exactly once. The operator token mints and revokes those tokens; a holder token reads that holder's own memories and nothing else; naming somebody else without the operator token is refused with both ids in the message. The visibility rule was leaking before this: browse had none at all, and the search rules short-circuited to the whole library whenever no holder was named. The rule is now written out in five copies with one parameter carried four times, and a test compares the copies so they cannot drift — semgrep rejects an interpolated constant, so the copies are not a choice. Eyes carry three states on purpose. NULL means the whole library, which is what a process that owns the store is entitled to; a holder id means that holder; the empty string means nobody, which is what a tokenless HTTP caller is. Making NULL mean nobody instead would have forced every internal caller — CLI, proofs, the compiler, the bench harnesses — to name eyes, contradicting the rule that a local process reads as its own trusted principal. The capability card answers "can this store answer my kind of question" before anything is sent. It is derived from what the store holds, so it cannot promise a modality the store lacks, and it carries no memory text, so it is safe to read before trusting the store. HTTP, CLI and one MCP tool all serve it without a token.
290 lines
12 KiB
Python
290 lines
12 KiB
Python
"""A cartridge is a retrieval scope: mount it and only it answers, append and the bundle follows."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Callable
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pyarrow.parquet as pq
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from typer.testing import CliRunner
|
|
|
|
import ely_eye.api as api_module
|
|
import ely_eye.cli as cli_module
|
|
from ely_eye import cartridge_store
|
|
from ely_eye.cartridge import CartridgeService
|
|
from ely_eye.cli import app
|
|
from ely_eye.config import Settings
|
|
from ely_eye.db import Database
|
|
from ely_eye.memory import MemoryService
|
|
from ely_eye.migrations import SCHEMA_VERSION
|
|
from ely_eye.retrieval import RetrievalService
|
|
from ely_eye.schemas import MemoryOrigin, SearchFilters
|
|
|
|
HOME_KEYS = "我们的钥匙放在玄关抽屉里"
|
|
WORK_KEYS = "备用钥匙在办公室第二个抽屉里"
|
|
|
|
|
|
class Store:
|
|
"""One SQLite file, the four services that read and write it."""
|
|
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
self.memory = MemoryService(settings=settings)
|
|
self.db = self.memory.db
|
|
self.cartridges = CartridgeService(settings=settings, database=self.db)
|
|
self.search = RetrievalService(settings=settings, database=self.db)
|
|
|
|
def remember(self, text: str, scenario: str) -> str:
|
|
atom = self.memory.remember(
|
|
text, holder_id="wife", origin=MemoryOrigin.self_, scenario=scenario
|
|
)
|
|
return atom.atom_id
|
|
|
|
def cartridge(self, name: str, atom_ids: list[str]) -> str:
|
|
atoms = self.db.get_atoms(atom_ids)
|
|
sources = self.db.get_sources(sorted({atom.source_id for atom in atoms}))
|
|
return self.cartridges.materialize(name, atoms, sources).cartridge_id
|
|
|
|
def texts(self, query: str = "钥匙", **filters: Any) -> list[str]:
|
|
eyes: dict[str, Any] = {"as_holder": "wife", **filters}
|
|
result = self.search.search(query, filters=SearchFilters(**eyes))
|
|
return [hit.atom.text for hit in result.hits]
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path: Path, model_dir: Path) -> Store:
|
|
return Store(Settings(home=tmp_path / "home", model_dir=model_dir, default_holder="wife"))
|
|
|
|
|
|
def test_mounting_one_cartridge_stops_the_other_from_answering(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
work = store.cartridge("work", [store.remember(WORK_KEYS, "work")])
|
|
assert sorted(store.texts()) == sorted([HOME_KEYS, WORK_KEYS])
|
|
|
|
store.cartridges.switch(home)
|
|
assert store.texts() == [HOME_KEYS]
|
|
store.cartridges.switch(work)
|
|
assert store.texts() == [WORK_KEYS]
|
|
|
|
|
|
def test_an_explicit_scope_overrides_what_is_mounted(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
work = store.cartridge("work", [store.remember(WORK_KEYS, "work")])
|
|
store.cartridges.switch(home)
|
|
assert store.texts(cartridge_ids=[work]) == [WORK_KEYS]
|
|
|
|
|
|
def test_nothing_mounted_answers_from_the_whole_library(store: Store) -> None:
|
|
store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
state = store.cartridges.mount_state()
|
|
assert state.mounted == []
|
|
assert state.scoped is False
|
|
assert state.atoms_in_scope == 0
|
|
assert store.texts() == [HOME_KEYS]
|
|
|
|
|
|
def test_mounting_two_cartridges_answers_from_both(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
work = store.cartridge("work", [store.remember(WORK_KEYS, "work")])
|
|
store.cartridges.mount(home)
|
|
store.cartridges.mount(work)
|
|
state = store.cartridges.mount_state()
|
|
assert sorted(state.mounted) == sorted([home, work])
|
|
assert state.atoms_in_scope == 2
|
|
assert sorted(store.texts()) == sorted([HOME_KEYS, WORK_KEYS])
|
|
|
|
|
|
def test_unmounting_reports_the_state_and_refuses_a_stranger(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
store.cartridges.mount(home)
|
|
state = store.cartridges.unmount(home)
|
|
assert state.mounted == [] and state.scoped is False
|
|
with pytest.raises(ValueError, match=f"{home} is not mounted; nothing to unmount"):
|
|
store.cartridges.unmount(home)
|
|
|
|
|
|
def test_append_adds_a_memory_and_the_bundle_follows(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
before = next(m for m in store.cartridges.list_manifests() if m.cartridge_id == home)
|
|
later = store.remember("备用钥匙在办公室第二个抽屉里", "home")
|
|
after = store.cartridges.append(home, [later])
|
|
|
|
assert after.dna != before.dna
|
|
assert after.atom_count == before.atom_count + 1
|
|
assert cartridge_store.membership_count(store.db, home) == 2
|
|
rows = pq.read_table(store.settings.cartridge_dir / home / "atoms.parquet").num_rows
|
|
assert rows == before.atom_count + 1
|
|
|
|
store.cartridges.switch(home)
|
|
assert sorted(store.texts()) == sorted([HOME_KEYS, WORK_KEYS])
|
|
|
|
|
|
def test_append_refuses_atoms_that_are_not_in_the_store(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
with pytest.raises(ValueError, match="not in the store: mem-missing:text"):
|
|
store.cartridges.append(home, ["mem-missing:text"])
|
|
|
|
|
|
def test_append_needs_something_to_append(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
with pytest.raises(ValueError, match="append needs at least one atom"):
|
|
store.cartridges.append(home, [])
|
|
|
|
|
|
def test_a_cartridge_without_a_fingerprint_cannot_be_mounted(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
row = next(r for r in store.db.list_cartridges() if r["cartridge_id"] == home)
|
|
manifest = json.loads(row["manifest_json"])
|
|
manifest["kv_fingerprint"] = None
|
|
store.db.write_cartridge(
|
|
cartridge_id=home,
|
|
name=row["name"],
|
|
root_path=Path(row["root_path"]),
|
|
manifest_json=json.dumps(manifest),
|
|
dna=row["dna"],
|
|
token_equivalent=row["token_equivalent"],
|
|
created_at=datetime.fromisoformat(row["created_at"]),
|
|
)
|
|
with pytest.raises(ValueError, match="carries no kv fingerprint; rebuild it before mounting"):
|
|
store.cartridges.mount(home)
|
|
|
|
|
|
def test_a_mount_refuses_a_different_rope_and_names_the_field(
|
|
store: Store, model_dir: Path, write_model_fixture: Callable[..., Path]
|
|
) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
write_model_fixture(model_dir, rope_parameters={"rope_type": "yarn", "factor": 2.0})
|
|
with pytest.raises(ValueError, match="rope: factor: cartridge 4.0 vs local 2.0"):
|
|
store.cartridges.mount(home)
|
|
|
|
|
|
def test_a_mount_refuses_when_the_machine_cannot_prove_the_model(store: Store) -> None:
|
|
home = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
blind = CartridgeService(
|
|
settings=Settings(home=store.settings.home, model_dir=None), database=store.db
|
|
)
|
|
with pytest.raises(ValueError, match="model_hash, tokenizer_hash, chat_template_hash, rope"):
|
|
blind.mount(home)
|
|
|
|
|
|
def test_forgetting_a_memory_takes_it_out_of_the_cartridge(store: Store) -> None:
|
|
atom_id = store.remember(HOME_KEYS, "home")
|
|
home = store.cartridge("home", [atom_id])
|
|
store.cartridges.mount(home)
|
|
store.memory.forget(atom_id=atom_id, reason="moved")
|
|
assert cartridge_store.membership_count(store.db, home) == 0
|
|
assert store.texts() == []
|
|
|
|
|
|
def upload_home(client: TestClient) -> str:
|
|
answer = client.post(
|
|
"/api/ingest/upload?cartridge_name=home",
|
|
files={"files": ("pantry.txt", b"the pantry keeps flour", "text/plain")},
|
|
)
|
|
assert answer.status_code == 200, answer.text
|
|
return answer.json()["cartridge_id"]
|
|
|
|
|
|
def test_the_api_mounts_and_reports_the_scope(
|
|
tmp_path: Path, model_dir: Path, monkeypatch: pytest.MonkeyPatch, write_token: str
|
|
) -> None:
|
|
settings = Settings(
|
|
home=tmp_path / "home", model_dir=model_dir, default_holder="wife", auth_token=write_token
|
|
)
|
|
monkeypatch.setattr(api_module, "get_settings", lambda: settings)
|
|
client = TestClient(api_module.create_app(), headers={"Authorization": f"Bearer {write_token}"})
|
|
|
|
atom = client.post("/api/memory", json={"text": HOME_KEYS, "scenario": "home"}).json()
|
|
assert client.get("/api/mounts").json() == {"mounted": [], "scoped": False, "atoms_in_scope": 0}
|
|
cartridge_id = upload_home(client)
|
|
|
|
answer = client.post(f"/api/cartridges/{cartridge_id}/mount")
|
|
assert answer.status_code == 200
|
|
assert answer.json()["kv_fingerprint"]["rope"]["factor"] == 4.0
|
|
assert client.get("/api/mounts").json() == {
|
|
"mounted": [cartridge_id],
|
|
"scoped": True,
|
|
"atoms_in_scope": 2,
|
|
}
|
|
hits = client.post("/api/search", json={"query": "钥匙"}).json()["hits"]
|
|
assert [hit["atom"]["atom_id"] for hit in hits] == [atom["atom_id"]]
|
|
|
|
stranger = client.post("/api/cartridges/cartridge-does-not-exist/mount")
|
|
assert stranger.status_code == 400 and "Unknown cartridge" in stranger.json()["detail"]
|
|
assert client.post(f"/api/cartridges/{cartridge_id}/unmount").status_code == 200
|
|
assert client.post(f"/api/cartridges/{cartridge_id}/unmount").status_code == 400
|
|
|
|
|
|
def test_the_api_appends_and_refuses_an_unknown_atom(
|
|
tmp_path: Path, model_dir: Path, monkeypatch: pytest.MonkeyPatch, write_token: str
|
|
) -> None:
|
|
settings = Settings(
|
|
home=tmp_path / "home", model_dir=model_dir, default_holder="wife", auth_token=write_token
|
|
)
|
|
monkeypatch.setattr(api_module, "get_settings", lambda: settings)
|
|
client = TestClient(api_module.create_app(), headers={"Authorization": f"Bearer {write_token}"})
|
|
client.post("/api/memory", json={"text": HOME_KEYS, "scenario": "home"})
|
|
cartridge_id = upload_home(client)
|
|
|
|
answer = client.post(
|
|
f"/api/cartridges/{cartridge_id}/append", json={"atom_ids": ["mem-nope:text"]}
|
|
)
|
|
assert answer.status_code == 400 and "not in the store" in answer.json()["detail"]
|
|
|
|
extra = client.post("/api/memory", json={"text": WORK_KEYS, "scenario": "home"}).json()
|
|
grown = client.post(
|
|
f"/api/cartridges/{cartridge_id}/append", json={"atom_ids": [extra["atom_id"]]}
|
|
)
|
|
assert grown.status_code == 200 and grown.json()["atom_count"] == 3
|
|
|
|
|
|
def test_the_cli_mounts_and_answers_in_json(
|
|
tmp_path: Path, model_dir: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
settings = Settings(home=tmp_path / "home", model_dir=model_dir, default_holder="wife")
|
|
monkeypatch.setattr(cli_module, "get_settings", lambda: settings)
|
|
runner = CliRunner()
|
|
store = Store(settings)
|
|
cartridge_id = store.cartridge("home", [store.remember(HOME_KEYS, "home")])
|
|
|
|
mounted = runner.invoke(app, ["mount", cartridge_id])
|
|
assert mounted.exit_code == 0, mounted.output
|
|
assert json.loads(mounted.output)["cartridge_id"] == cartridge_id
|
|
|
|
state = json.loads(runner.invoke(app, ["mounts"]).output)
|
|
assert state == {"mounted": [cartridge_id], "scoped": True, "atoms_in_scope": 1}
|
|
|
|
refused = runner.invoke(app, ["mount", "cartridge-does-not-exist"])
|
|
assert isinstance(refused.exception, ValueError) and "Unknown cartridge" in str(
|
|
refused.exception
|
|
)
|
|
|
|
|
|
def test_a_third_version_store_gains_the_scope_tables_on_upgrade(
|
|
tmp_path: Path, model_dir: Path
|
|
) -> None:
|
|
settings = Settings(home=tmp_path / "home", model_dir=model_dir)
|
|
first = Database(settings)
|
|
with first.connect() as connection:
|
|
connection.execute("DROP TABLE cartridge_atoms")
|
|
connection.execute("DROP TABLE mounts")
|
|
connection.execute("PRAGMA user_version = 3")
|
|
|
|
reopened = Database(settings)
|
|
with reopened.connect() as connection:
|
|
tables = {
|
|
row[0]
|
|
for row in connection.execute(
|
|
"SELECT name FROM sqlite_master WHERE type = 'table'"
|
|
).fetchall()
|
|
}
|
|
version = connection.execute("PRAGMA user_version").fetchone()[0]
|
|
assert {"cartridge_atoms", "mounts"} <= tables
|
|
assert version == SCHEMA_VERSION
|
|
assert cartridge_store.mounted_ids(reopened) == []
|