Initial release: iai-mcp v0.1.0

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: XNLLLLH <XNLLLLH@users.noreply.github.com>

tests/test_pipeline.py

@@ -0,0 +1,431 @@
+"""Tests for iai_mcp.pipeline (D-13 5-stage retrieval pipeline).
+
+Uses a FakeEmbedder fixture so tests don't pull BAAI/bge-small-en-v1.5 from
+HuggingFace during every run. The Embedder contract verified separately in
+test_embed.py.
+"""
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from iai_mcp.community import CommunityAssignment
+from iai_mcp.graph import MemoryGraph
+from iai_mcp.pipeline import (
+    W_AAAK,
+    W_AGE,
+    W_COSINE,
+    W_DEGREE,
+    _aaak_overlap,
+    _community_gate,
+    _cosine,
+    _pick_seeds,
+    recall_for_response,
+)
+from iai_mcp.store import MemoryStore
+from iai_mcp.types import EMBED_DIM, MemoryRecord
+
+
+class _FakeEmbedder:
+    """Stand-in for the configured embedder so tests don't require the model.
+
+    Returns a deterministic primary-axis vector for any input. recall_for_response
+    uses embedder.embed() only for the cue, so this is sufficient.
+
+    DIM follows the default registry (bge-m3 = 1024d). Plan-02 tests
+    that hand-build vectors must use `[1.0] + [0.0] * (DIM - 1)` style via this
+    constant so the store.insert() dim-check passes.
+    """
+
+    DIM = EMBED_DIM  # 1024 under default (bge-m3)
+
+    def embed(self, text: str) -> list[float]:
+        return [1.0] + [0.0] * (EMBED_DIM - 1)
+
+    def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        return [self.embed(t) for t in texts]
+
+
+def _make(vec: list[float], text: str = "rec", aaak: str = "", detail: int = 2) -> MemoryRecord:
+    """Construct a MemoryRecord for pipeline tests.
+
+    Uses `tier="episodic"` to stay within TIER_ENUM; created_at at current UTC.
+    language="en" required.
+    """
+    now = datetime.now(timezone.utc)
+    return MemoryRecord(
+        id=uuid4(),
+        tier="episodic",
+        literal_surface=text,
+        aaak_index=aaak,
+        embedding=vec,
+        community_id=None,
+        centrality=0.0,
+        detail_level=detail,
+        pinned=False,
+        stability=0.0,
+        difficulty=0.0,
+        last_reviewed=None,
+        never_decay=False,
+        never_merge=False,
+        provenance=[],
+        created_at=now,
+        updated_at=now,
+        tags=[],
+        language="en",
+    )
+
+
+# ---------------------------------------------------------- stage-unit tests
+
+
+def test_community_gate_picks_nearest() -> None:
+    """CONN-06: top-1 gate on 3 centroids picks the one nearest the cue."""
+    c0 = uuid4()
+    c1 = uuid4()
+    c2 = uuid4()
+    centroids = {
+        c0: [1.0] + [0.0] * (EMBED_DIM - 1),
+        c1: [0.0] * 384,
+        c2: [-1.0] + [0.0] * (EMBED_DIM - 1),
+    }
+    a = CommunityAssignment(community_centroids=centroids)
+    cue = [1.0] + [0.0] * (EMBED_DIM - 1)
+    gated = _community_gate(cue, a, top_n=1)
+    assert len(gated) == 1
+    assert gated[0] == c0
+
+
+def test_community_gate_returns_top_n_in_order() -> None:
+    c0 = uuid4()
+    c1 = uuid4()
+    c2 = uuid4()
+    centroids = {
+        c0: [1.0] + [0.0] * (EMBED_DIM - 1),
+        c1: [0.5, 0.5] + [0.0] * (EMBED_DIM - 2),
+        c2: [-1.0] + [0.0] * (EMBED_DIM - 1),
+    }
+    a = CommunityAssignment(community_centroids=centroids)
+    cue = [1.0] + [0.0] * (EMBED_DIM - 1)
+    gated = _community_gate(cue, a, top_n=3)
+    assert gated == [c0, c1, c2]
+
+
+def test_aaak_overlap_basic_jaccard() -> None:
+    assert _aaak_overlap("", "anything") == 0.0
+    assert _aaak_overlap("x", "") == 0.0
+    assert _aaak_overlap("a b", "a b") == 1.0  # identical
+    # Jaccard({a,b}, {b,c}) = 1 / 3
+    assert abs(_aaak_overlap("a b", "b c") - 1 / 3) < 1e-9
+
+
+def test_aaak_overlap_slash_split_symmetric() -> None:
+    """AAAK tokens use '/' as separator; both sides must split on it."""
+    # Identical slash-delimited paths -> 1.0 (bug fix: cue side also splits).
+    assert _aaak_overlap("auth/login", "auth/login") == 1.0
+    # Partial share: {auth, login} vs {auth, logout} -> Jaccard = 1/3.
+    assert abs(_aaak_overlap("auth/login", "auth/logout") - 1 / 3) < 1e-9
+    # Case-insensitive.
+    assert _aaak_overlap("AUTH/Login", "auth/login") == 1.0
+
+
+def test_cosine_basic_properties() -> None:
+    assert _cosine([1.0, 0.0], [1.0, 0.0]) == 1.0
+    assert _cosine([1.0, 0.0], [-1.0, 0.0]) == -1.0
+    assert _cosine([1.0, 0.0], [0.0, 1.0]) == 0.0
+    assert _cosine([0.0, 0.0], [1.0, 0.0]) == 0.0  # zero vector guard
+
+
+def test_score_weight_constants_match_d13() -> None:
+    """D-13 score = cos + 0.3*aaak + 0.1*log(1+deg) − 0.05*age."""
+    assert W_COSINE == 1.0
+    assert W_AAAK == 0.3
+    assert W_DEGREE == 0.1
+    assert W_AGE == 0.05
+
+
+# ------------------------------------------------------------- end-to-end
+
+
+def test_pipeline_returns_hits_with_adjacent_suggestions(tmp_path) -> None:
+    """End-to-end: pipeline returns ranked hits with non-empty activation_trace
+    and adjacent_suggestions is a list on every hit (AUTIST-07 contract)."""
+    store = MemoryStore(path=tmp_path)
+    records = [
+        _make([1.0] + [0.0] * (EMBED_DIM - 1), text="primary match", aaak="test match"),
+        _make([0.9, 0.1] + [0.0] * (EMBED_DIM - 2), text="close match"),
+        _make([0.0, 1.0] + [0.0] * (EMBED_DIM - 2), text="orthogonal"),
+        _make([-1.0] + [0.0] * (EMBED_DIM - 1), text="opposite"),
+        _make([0.5, 0.5] + [0.0] * (EMBED_DIM - 2), text="mid"),
+    ]
+    for r in records:
+        store.insert(r)
+    graph = MemoryGraph()
+    for r in records:
+        graph.add_node(r.id, community_id=None, embedding=r.embedding)
+    for i in range(len(records) - 1):
+        graph.add_edge(records[i].id, records[i + 1].id)
+
+    community_id = uuid4()
+    assignment = CommunityAssignment(
+        node_to_community={r.id: community_id for r in records},
+        community_centroids={community_id: [1.0] + [0.0] * (EMBED_DIM - 1)},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[community_id],
+        mid_regions={community_id: [r.id for r in records]},
+    )
+
+    resp = recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="test match",
+        session_id="s1",
+    )
+    assert len(resp.hits) >= 1
+    # Primary record has aaak overlap ("test match" in both cue and aaak_index),
+    # cosine=1.0, and degree=1: score = 1.0 + 0.3*1.0 + 0.1*log(2) ≈ 1.369.
+    # Close record has cos≈0.994, no aaak, degree=2: 0.994 + 0.1*log(3) ≈ 1.104.
+    # Primary must win thanks to the AAAK overlap bonus.
+    assert resp.hits[0].literal_surface == "primary match"
+    # Opposite record must NOT appear as a top hit (negative cosine).
+    assert all(h.literal_surface != "opposite" for h in resp.hits[:2])
+    # adjacent_suggestions must be a list on every hit.
+    for h in resp.hits:
+        assert isinstance(h.adjacent_suggestions, list)
+    # activation_trace = seeds ∪ spread; must not be empty here.
+    assert len(resp.activation_trace) >= 1
+
+
+def test_pipeline_provenance_appended_to_every_hit(tmp_path) -> None:
+    """MEM-05 regression: every hit returned gets a provenance entry."""
+    store = MemoryStore(path=tmp_path)
+    r1 = _make([1.0] + [0.0] * (EMBED_DIM - 1), text="primary")
+    store.insert(r1)
+    graph = MemoryGraph()
+    graph.add_node(r1.id, community_id=None, embedding=r1.embedding)
+    community_id = uuid4()
+    assignment = CommunityAssignment(
+        node_to_community={r1.id: community_id},
+        community_centroids={community_id: [1.0] + [0.0] * (EMBED_DIM - 1)},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[community_id],
+        mid_regions={community_id: [r1.id]},
+    )
+    recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="anything",
+        session_id="session-42",
+    )
+    refreshed = store.get(r1.id)
+    assert refreshed is not None
+    assert len(refreshed.provenance) == 1
+    assert refreshed.provenance[0]["session_id"] == "session-42"
+    assert refreshed.provenance[0]["cue"] == "anything"
+
+
+def test_pipeline_budget_caps_hit_count(tmp_path) -> None:
+    """Budget enforcement: when tokens exceeded, pipeline stops adding hits."""
+    store = MemoryStore(path=tmp_path)
+    # 5 records each with ~200 chars (~50 tokens). Budget=60 -> only first fits.
+    long_text = "x" * 200
+    records = []
+    for i in range(5):
+        r = _make(
+            [1.0, float(i) * 0.001] + [0.0] * (EMBED_DIM - 2),
+            text=f"{long_text}-{i}",
+        )
+        records.append(r)
+        store.insert(r)
+    graph = MemoryGraph()
+    for r in records:
+        graph.add_node(r.id, community_id=None, embedding=r.embedding)
+    community_id = uuid4()
+    assignment = CommunityAssignment(
+        node_to_community={r.id: community_id for r in records},
+        community_centroids={community_id: [1.0] + [0.0] * (EMBED_DIM - 1)},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[community_id],
+        mid_regions={community_id: [r.id for r in records]},
+    )
+    resp = recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="c",
+        session_id="s",
+        budget_tokens=60,
+    )
+    # With 50-token records and 60-token budget, at most 1 hit fits then loop breaks.
+    # (We always admit 1 even if it exceeds budget, per the len(hits)>=1 guard.)
+    assert len(resp.hits) == 1
+
+
+def test_pipeline_anti_hits_from_contradicts_edge(tmp_path) -> None:
+    """D-13 anti-hit contract: contradicts-edge neighbours of a top hit surface."""
+    from iai_mcp.core import dispatch
+
+    store = MemoryStore(path=tmp_path)
+    r1 = _make([1.0] + [0.0] * (EMBED_DIM - 1), text="original")
+    store.insert(r1)
+    dispatch(
+        store,
+        "memory_contradict",
+        {
+            "id": str(r1.id),
+            "new_fact": "refuted version",
+            "cue_embedding": r1.embedding,
+        },
+    )
+
+    graph = MemoryGraph()
+    graph.add_node(r1.id, community_id=None, embedding=[1.0] + [0.0] * (EMBED_DIM - 1))
+    community_id = uuid4()
+    assignment = CommunityAssignment(
+        node_to_community={r1.id: community_id},
+        community_centroids={community_id: [1.0] + [0.0] * (EMBED_DIM - 1)},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[community_id],
+        mid_regions={community_id: [r1.id]},
+    )
+    resp = recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="anything",
+        session_id="s1",
+    )
+    assert len(resp.anti_hits) >= 1
+    assert "refuted" in resp.anti_hits[0].literal_surface
+
+
+def test_pipeline_activation_trace_includes_seeds(tmp_path) -> None:
+    """activation_trace = seeds ∪ spread; must contain each seed."""
+    store = MemoryStore(path=tmp_path)
+    a = _make([1.0] + [0.0] * (EMBED_DIM - 1), text="A")
+    b = _make([0.9, 0.1] + [0.0] * (EMBED_DIM - 2), text="B")
+    c = _make([0.0, 1.0] + [0.0] * (EMBED_DIM - 2), text="C")
+    for r in (a, b, c):
+        store.insert(r)
+    graph = MemoryGraph()
+    for r in (a, b, c):
+        graph.add_node(r.id, community_id=None, embedding=r.embedding)
+    graph.add_edge(a.id, b.id)
+    graph.add_edge(b.id, c.id)
+    community_id = uuid4()
+    assignment = CommunityAssignment(
+        node_to_community={r.id: community_id for r in (a, b, c)},
+        community_centroids={community_id: [1.0] + [0.0] * (EMBED_DIM - 1)},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[community_id],
+        mid_regions={community_id: [a.id, b.id, c.id]},
+    )
+    resp = recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="c",
+        session_id="s",
+    )
+    # The top-cosine seed is A; its 2-hop neighbourhood is {B, C}. Trace must contain A.
+    assert a.id in resp.activation_trace
+
+
+def test_pick_seeds_ranks_by_blended_score(tmp_path) -> None:
+    """Stage 3 blend: 0.6*cos + 0.4*centrality picks the high-blend record first.
+
+    redesign: `_pick_seeds` now operates over a
+    precomputed shared cosine array; positions, not UUIDs, flow through.
+    Reproduces the pre-redesign assertion: r2 (cos=0.707, cen=1.0,
+    blend=0.82) beats r1 (cos=1.0, cen=0.0, blend=0.6) at n=1.
+    """
+    import numpy as np
+
+    # Pool layout: position 0 = r1, position 1 = r2.
+    # cue = axis 0 -> shared_cos = [1.0, 0.707].
+    shared_cos = np.array([1.0, 0.7071068], dtype=np.float32)
+    centrality_arr = np.array([0.0, 1.0], dtype=np.float32)
+    candidate_indices = np.array([0, 1], dtype=np.int64)
+    seed_indices = _pick_seeds(
+        candidate_indices, shared_cos, centrality_arr, n=1,
+    )
+    # r2 (position 1): blend = 0.6 * 0.707 + 0.4 * 1.0 = 0.824 > r1's 0.6.
+    assert list(seed_indices) == [1]
+
+
+def test_pipeline_core_dispatch_integration(tmp_path, monkeypatch) -> None:
+    """core.dispatch("memory_recall", ...) routes to pipeline for non-empty store."""
+    import iai_mcp.pipeline as pipeline_mod
+    from iai_mcp.core import dispatch
+
+    store = MemoryStore(path=tmp_path)
+    r = _make([1.0] + [0.0] * (EMBED_DIM - 1), text="integration")
+    store.insert(r)
+
+    # Stub out Embedder inside core to avoid HF download.
+    class _StubEmbedder:
+        DIM = 384
+
+        def embed(self, text: str) -> list[float]:
+            return [1.0] + [0.0] * (EMBED_DIM - 1)
+
+    # core.py imports Embedder lazily inside dispatch -> patch at module level.
+    import iai_mcp.embed as embed_mod
+    monkeypatch.setattr(embed_mod, "Embedder", _StubEmbedder)
+
+    resp = dispatch(
+        store,
+        "memory_recall",
+        {"cue": "integration", "session_id": "s-int"},
+    )
+    assert "hits" in resp
+    assert isinstance(resp["hits"], list)
+    # activation_trace field always present, list of string UUIDs.
+    assert isinstance(resp["activation_trace"], list)
+    assert "budget_used" in resp
+
+
+def test_pipeline_empty_gate_falls_back_to_all_nodes(tmp_path) -> None:
+    """If community gate returns no candidates, pipeline falls back to all nodes."""
+    store = MemoryStore(path=tmp_path)
+    r = _make([1.0] + [0.0] * (EMBED_DIM - 1), text="lonely")
+    store.insert(r)
+    graph = MemoryGraph()
+    graph.add_node(r.id, community_id=None, embedding=r.embedding)
+    # Assignment whose mid_regions is empty (degenerate) -> pipeline must fall back.
+    assignment = CommunityAssignment(
+        node_to_community={},
+        community_centroids={},
+        modularity=0.0,
+        backend="flat",
+        top_communities=[],
+        mid_regions={},
+    )
+    resp = recall_for_response(
+        store=store,
+        graph=graph,
+        assignment=assignment,
+        rich_club=[],
+        embedder=_FakeEmbedder(),
+        cue="c",
+        session_id="s",
+    )
+    # The lone record is still reachable via the fallback.
+    assert len(resp.hits) == 1