Engram-protocol/engram

1"""Tests for kvcos.engram.knowledge_index — HNSW knowledge search."""2 3import json4from pathlib import Path5 6import pytest7import torch8 9from kvcos.engram.embedder import get_fingerprint10from kvcos.engram.format import EigramEncoder11from kvcos.engram.knowledge_index import KnowledgeIndex12 13 14@pytest.fixture15def knowledge_dir(tmp_path):16    """Create a temporary knowledge directory with test .eng files."""17    encoder = EigramEncoder()18    project_dir = tmp_path / "test_project"19    project_dir.mkdir()20 21    docs = [22        ("doc_ml", "Machine learning model training and optimization"),23        ("doc_db", "PostgreSQL database schema migration tools"),24        ("doc_api", "REST API endpoint authentication and authorization"),25        ("doc_test", "Unit testing with pytest fixtures and mocking"),26        ("doc_deploy", "Docker container deployment to Kubernetes cluster"),27    ]28 29    for doc_id, text in docs:30        fp, source = get_fingerprint(text)31        dim = fp.shape[0]32 33        blob = encoder.encode(34            vec_perdoc=torch.zeros(116),35            vec_fcdb=torch.zeros(116),36            joint_center=torch.zeros(128),37            corpus_hash="test" * 8,38            model_id=source[:16],39            basis_rank=116,40            n_corpus=0,41            layer_range=(0, 0),42            context_len=len(text),43            l2_norm=float(torch.norm(fp).item()),44            scs=0.0,45            margin_proof=0.0,46            task_description=text[:256],47            cache_id=doc_id,48            vec_fourier=fp if dim == 2048 else None,49            vec_fourier_v2=fp,50            confusion_flag=False,51        )52 53        eng_path = project_dir / f"{doc_id}.eng"54        eng_path.write_bytes(blob)55 56        meta = {57            "cache_id": doc_id,58            "task_description": text,59            "source_path": f"/test/{doc_id}.md",60            "project": "test_project",61            "fp_source": source,62            "chunk_index": 0,63            "chunk_total": 1,64            "headers": [],65        }66        meta_path = Path(str(eng_path) + ".meta.json")67        meta_path.write_text(json.dumps(meta))68 69    return tmp_path70 71 72class TestKnowledgeIndexBuild:73    def test_build_from_directory(self, knowledge_dir):74        kidx = KnowledgeIndex.build_from_knowledge_dir(75            knowledge_dir, verbose=False76        )77        assert len(kidx) == 578 79    def test_build_empty_directory(self, tmp_path):80        with pytest.raises(ValueError, match="No .eng files"):81            KnowledgeIndex.build_from_knowledge_dir(tmp_path, verbose=False)82 83 84class TestKnowledgeIndexSearch:85    def test_search_returns_results(self, knowledge_dir):86        kidx = KnowledgeIndex.build_from_knowledge_dir(87            knowledge_dir, verbose=False88        )89        results = kidx.search("database query optimization", k=3)90        assert len(results) == 391        assert all(r.score > 0 for r in results)92 93    def test_search_result_fields(self, knowledge_dir):94        kidx = KnowledgeIndex.build_from_knowledge_dir(95            knowledge_dir, verbose=False96        )97        results = kidx.search("testing", k=1)98        r = results[0]99        assert r.doc_id100        assert isinstance(r.score, float)101        assert r.rank == 0102        assert r.project == "test_project"103 104    def test_search_with_tensor(self, knowledge_dir):105        kidx = KnowledgeIndex.build_from_knowledge_dir(106            knowledge_dir, verbose=False107        )108        query_fp, _ = get_fingerprint("unit tests")109        results = kidx.search(query_fp, k=2)110        assert len(results) == 2111 112    def test_search_margin(self, knowledge_dir):113        kidx = KnowledgeIndex.build_from_knowledge_dir(114            knowledge_dir, verbose=False115        )116        results = kidx.search("testing", k=3)117        # Top result should have a margin118        assert results[0].margin >= 0119 120 121class TestKnowledgeIndexPersistence:122    def test_save_and_load(self, knowledge_dir, tmp_path):123        kidx = KnowledgeIndex.build_from_knowledge_dir(124            knowledge_dir, verbose=False125        )126        index_dir = tmp_path / "index"127        kidx.save(index_dir)128 129        loaded = KnowledgeIndex.load(index_dir)130        assert len(loaded) == len(kidx)131 132        # Search should work on loaded index133        results = loaded.search("database", k=2)134        assert len(results) == 2135 136    def test_load_nonexistent(self, tmp_path):137        with pytest.raises(FileNotFoundError):138            KnowledgeIndex.load(tmp_path / "nonexistent")139