ely-eye/backend/tests/test_cartridge_scope.py

"""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]:
        result = self.search.search(query, filters=SearchFilters(**filters) if filters else None)
        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) == []