| """ |
| ENGRAM Protocol — Retriever Tests |
| Tests for EGRRetriever: store → index → query → retrieve pipeline. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import torch |
|
|
| from kvcos.core.cache_spec import LLAMA_3_1_8B |
| from kvcos.core.serializer import EngramSerializer |
| from kvcos.core.types import CompressionMethod, StateExtractionMode |
| from kvcos.core.manifold_index import ManifoldIndex |
| from kvcos.core.retriever import EGRRetriever, RetrievalResponse |
| from kvcos.core.state_extractor import MARStateExtractor |
| from kvcos.storage.local import LocalStorageBackend |
| from tests.conftest import make_synthetic_kv |
|
|
|
|
| def _build_retriever( |
| data_dir: Path, mode: StateExtractionMode = StateExtractionMode.MEAN_POOL, |
| ) -> EGRRetriever: |
| ext = MARStateExtractor(mode=mode, rank=128) |
| dim = ext.output_dim(LLAMA_3_1_8B) |
| idx = ManifoldIndex(dim=dim) |
| storage = LocalStorageBackend(data_dir=data_dir) |
| return EGRRetriever(ext, idx, storage) |
|
|
|
|
| class TestIndexAndRetrieve: |
| """Full store → search → load pipeline.""" |
|
|
| def test_index_returns_cache_id(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| retriever = _build_retriever(tmp_data_dir) |
|
|
| cid = retriever.index_engram( |
| keys=keys, values=values, spec=LLAMA_3_1_8B, |
| agent_id="test", task_description="test engram", |
| model_id=LLAMA_3_1_8B["model_id"], |
| output_dir=tmp_data_dir, |
| ) |
| assert isinstance(cid, str) |
| assert len(cid) > 0 |
|
|
| def test_retrieve_finds_stored(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| retriever = _build_retriever(tmp_data_dir) |
|
|
| retriever.index_engram( |
| keys=keys, values=values, spec=LLAMA_3_1_8B, |
| agent_id="test", task_description="findable engram", |
| model_id=LLAMA_3_1_8B["model_id"], |
| output_dir=tmp_data_dir, |
| ) |
|
|
| query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64, seed=99) |
| response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=1) |
|
|
| assert isinstance(response, RetrievalResponse) |
| assert len(response.results) == 1 |
| assert response.results[0].keys.shape == keys.shape |
|
|
| def test_retrieve_empty_index(self, tmp_data_dir: Path) -> None: |
| retriever = _build_retriever(tmp_data_dir) |
| query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=5) |
| assert len(response.results) == 0 |
|
|
| def test_delete_removes(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| retriever = _build_retriever(tmp_data_dir) |
|
|
| cid = retriever.index_engram( |
| keys=keys, values=values, spec=LLAMA_3_1_8B, |
| agent_id="test", task_description="deletable", |
| model_id=LLAMA_3_1_8B["model_id"], |
| output_dir=tmp_data_dir, |
| ) |
| assert retriever.delete_engram(cid) |
|
|
| query_keys, _ = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| response = retriever.retrieve(query_keys, LLAMA_3_1_8B, top_k=5) |
| assert len(response.results) == 0 |
|
|
