@pentatonic-ai/ai-agent-sdk 0.10.4 → 0.10.6

package/packages/memory-engine-v2/tests/test_entity_resolution_v2.py

@@ -0,0 +1,507 @@
+"""Unit tests for scripts/entity_resolution_v2.py (BET 1b: blocking +
+embedding similarity + LLM adjudication — dry-run tooling).
+
+Pure-python: no DB, no network, no embeddings endpoint, no LLM. The
+HTTP embedding backend and the Anthropic adjudicator are NEVER called
+here — tests use fakes. Importable without psycopg (both scripts guard
+the driver import).
+
+Run: pytest packages/memory-engine-v2/tests/test_entity_resolution_v2.py
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+_SCRIPTS = Path(__file__).resolve().parent.parent / "scripts"
+
+
+def _load(name: str):
+    if name in sys.modules:
+        return sys.modules[name]
+    spec = importlib.util.spec_from_file_location(name, _SCRIPTS / f"{name}.py")
+    assert spec and spec.loader
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[name] = mod  # register before exec (py3.13+ dataclasses)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+er2 = _load("entity_resolution_v2")
+v1 = _load("backfill_entity_reconciliation")
+
+
+def make_entity(eid: str, name: str, aliases: list[str] | None = None,
+                facts: int = 0, rels: int = 0) -> "v1.Entity":
+    return v1.Entity(
+        id=eid, canonical_name=name, aliases=aliases or [],
+        provenance_event_ids=[], fact_count=facts, rel_count=rels,
+        norm_forms={v1._normalize_surface(name),
+                    *(v1._normalize_surface(a) for a in (aliases or []))},
+    )
+
+
+class FakeAdjudicator(er2.Adjudicator):
+    """Deterministic verdicts keyed by frozenset of entity ids."""
+
+    def __init__(self, verdicts: dict[frozenset, str]) -> None:
+        self.verdicts = verdicts
+        self.calls: list[frozenset] = []
+
+    def adjudicate(self, a, a_facts, b, b_facts):
+        key = frozenset({a.id, b.id})
+        self.calls.append(key)
+        return er2.Adjudication(self.verdicts.get(key, "unsure"),
+                                f"fake verdict for {sorted(key)}")
+
+
+# ----------------------------------------------------------------------
+# 1. Candidate-generation normalization — the exact Johann variants.
+#    NEVER identity: also locked in below.
+# ----------------------------------------------------------------------
+
+def test_block_normalize_underscore():
+    assert er2.block_normalize("Johann_Boedecker") == "johann boedecker"
+
+
+def test_block_normalize_diacritic_and_comma():
+    # diacritic fold (NFD + strip combining) + punctuation strip
+    assert er2.block_normalize("Bödecker, Johann") == "bodecker johann"
+
+
+def test_token_sort_handles_comma_inversion():
+    assert er2.token_sort("Bödecker, Johann") == er2.token_sort("Bodecker Johann")
+    assert er2.token_sort("Bödecker, Johann") == "bodecker johann"
+
+
+def test_char_trigrams_bridge_oe_vs_diacritic_fold():
+    # "Boedecker" (oe transliteration) vs "Bödecker" (folds to bodecker)
+    # differ as tokens but must share trigram buckets.
+    shared = er2.char_trigrams("Boedecker") & er2.char_trigrams("Bödecker")
+    assert {"dec", "eck", "cke", "ker"} <= shared
+
+
+def test_blocking_normalization_is_not_identity_normalization():
+    # The identity scheme (v1 / entity_id.py) preserves diacritics and
+    # punctuation; the blocking form must NEVER replace it.
+    s = "Bödecker, Johann"
+    assert v1._normalize_surface(s) == "bödecker, johann"   # identity: untouched ö + comma
+    assert er2.block_normalize(s) == "bodecker johann"       # blocking only
+    assert v1._normalize_surface(s) != er2.block_normalize(s)
+
+
+def test_blocking_groups_johann_variants():
+    a = make_entity("e_a", "Johann_Boedecker")
+    b = make_entity("e_b", "Bödecker, Johann")
+    c = make_entity("e_c", "Johann Boedecker")
+    bare = make_entity("e_bare", "Johann")
+    near = make_entity("e_near", "Johanna Phil")
+    unrelated = make_entity("e_zoe", "Zoe Quist")
+
+    pairs = er2.generate_candidate_pairs([a, b, c, bare, near, unrelated])
+    keys = {p.key for p in pairs}
+
+    # All Boedecker spellings co-block pairwise.
+    assert frozenset({"e_a", "e_b"}) in keys
+    assert frozenset({"e_a", "e_c"}) in keys
+    assert frozenset({"e_b", "e_c"}) in keys
+    # Bare "Johann" blocks with the full names (first/last token key).
+    assert frozenset({"e_bare", "e_c"}) in keys
+    # Near-miss "Johanna Phil" co-blocks via trigrams — it must reach
+    # the similarity/adjudication stage (and be rejected there), not
+    # be silently invisible.
+    assert frozenset({"e_near", "e_c"}) in keys
+    # Unrelated person doesn't pair with the Boedecker cluster.
+    assert not any("e_zoe" in k and k != frozenset({"e_zoe"}) for k in keys)
+
+
+def test_blocking_email_local_part_key():
+    a = make_entity("e_a", "johann.boedecker@pentatonic.com",
+                    aliases=["johann.boedecker@pentatonic.com"])
+    b = make_entity("e_b", "Johann Boedecker")
+    pairs = er2.generate_candidate_pairs([a, b])
+    assert frozenset({"e_a", "e_b"}) in {p.key for p in pairs}
+
+
+def test_blocking_skips_pairs_already_grouped_by_v1():
+    a = make_entity("e_a", "Johann_Boedecker")
+    b = make_entity("e_b", "Johann Boedecker")
+    pairs = er2.generate_candidate_pairs(
+        [a, b], already_grouped={frozenset({"e_a", "e_b"})})
+    assert pairs == []
+
+
+def test_oversized_blocks_are_skipped():
+    clones = [make_entity(f"e_{i}", f"Johann Clone{i}") for i in range(10)]
+    pairs = er2.generate_candidate_pairs(clones, max_block=3)
+    # every key these share ("first:johann"/"last:..."/trigrams of
+    # "johann...") is oversized or unique per clone → no pair explosion
+    # from the shared-first-name block.
+    shared_first = [p for p in pairs if "first:clone0" in p.shared_keys]
+    assert shared_first == []
+
+
+# ----------------------------------------------------------------------
+# 2. Threshold band routing
+# ----------------------------------------------------------------------
+
+def test_threshold_constants():
+    assert er2.HIGH_THRESHOLD == 0.92
+    assert er2.LOW_THRESHOLD == 0.75
+
+
+def test_route_band_boundaries():
+    assert er2.route_band(0.95) == "high"
+    assert er2.route_band(0.92) == "high"        # inclusive
+    assert er2.route_band(0.9199) == "ambiguous"
+    assert er2.route_band(0.75) == "ambiguous"   # inclusive
+    assert er2.route_band(0.7499) == "drop"
+    assert er2.route_band(0.10) == "drop"
+
+
+def _scored_pair(a, b, sim):
+    p = er2.CandidatePair(a=a, b=b)
+    p.similarity = sim
+    return p
+
+
+def test_route_pairs_band_routing():
+    a1, a2 = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Bödecker, Johann")
+    b1, b2 = make_entity("e_3", "Carla Voss"), make_entity("e_4", "Carla Vosse")
+    c1, c2 = make_entity("e_5", "Johanna Phil"), make_entity("e_6", "Johann Boedeker")
+    pairs = [
+        _scored_pair(a1, a2, 0.96),   # high → auto merge, no LLM
+        _scored_pair(b1, b2, 0.85),   # ambiguous → adjudicated yes
+        _scored_pair(c1, c2, 0.60),   # below low → dropped, no LLM
+    ]
+    adj = FakeAdjudicator({frozenset({"e_3", "e_4"}): "yes"})
+    routed = er2.route_pairs(pairs, adj, {})
+
+    assert {p.key for p in routed.merge} == {frozenset({"e_1", "e_2"}),
+                                             frozenset({"e_3", "e_4"})}
+    assert {p.key for p in routed.dropped} == {frozenset({"e_5", "e_6"})}
+    assert routed.human_review == []
+    # LLM consulted ONLY for the ambiguous pair.
+    assert adj.calls == [frozenset({"e_3", "e_4"})]
+    # high-confidence merge carries its reasoning
+    auto = next(p for p in routed.merge if p.key == frozenset({"e_1", "e_2"}))
+    assert auto.verdict == "auto" and "0.960" in auto.reason
+
+
+def test_route_pairs_unsure_never_merges():
+    a, b = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Johanna Phil")
+    routed = er2.route_pairs(
+        [_scored_pair(a, b, 0.85)],
+        FakeAdjudicator({frozenset({"e_1", "e_2"}): "unsure"}), {})
+    assert routed.merge == []
+    assert [p.key for p in routed.human_review] == [frozenset({"e_1", "e_2"})]
+
+
+def test_route_pairs_no_verdict_means_no_merge():
+    a, b = make_entity("e_1", "Johann Boedecker"), make_entity("e_2", "Johanna Phil")
+    routed = er2.route_pairs(
+        [_scored_pair(a, b, 0.85)],
+        FakeAdjudicator({frozenset({"e_1", "e_2"}): "no"}), {})
+    assert routed.merge == [] and routed.human_review == []
+    assert [p.key for p in routed.dropped] == [frozenset({"e_1", "e_2"})]
+
+
+def test_no_llm_routes_whole_ambiguous_band_to_human_review():
+    a, b = make_entity("e_1", "Carla Voss"), make_entity("e_2", "Carla Vosse")
+    c, d = make_entity("e_3", "Mark Diaz"), make_entity("e_4", "Marc Diaz")
+    routed = er2.route_pairs(
+        [_scored_pair(a, b, 0.85), _scored_pair(c, d, 0.80)],
+        er2.NoLLMAdjudicator(), {})
+    assert routed.merge == []
+    assert len(routed.human_review) == 2
+    assert all(p.verdict == "unsure" for p in routed.human_review)
+
+
+# ----------------------------------------------------------------------
+# 3. Bare-first-name policy
+# ----------------------------------------------------------------------
+
+def test_is_bare_name():
+    assert er2.is_bare_name(make_entity("e", "Johann"))
+    assert not er2.is_bare_name(make_entity("e", "Johann Boedecker"))
+    # an email-only entity is not a bare *name*
+    assert not er2.is_bare_name(make_entity("e", "johann@x.com"))
+    # a single-token canonical with a multi-token alias is not bare
+    assert not er2.is_bare_name(
+        make_entity("e", "Johann", aliases=["Johann Boedecker"]))
+
+
+def test_bare_name_single_candidate_and_yes_merges():
+    bare = make_entity("e_bare", "Johann")
+    full = make_entity("e_full", "Johann Boedecker", facts=10)
+    routed = er2.route_pairs(
+        [_scored_pair(bare, full, 0.95)],
+        FakeAdjudicator({frozenset({"e_bare", "e_full"}): "yes"}), {})
+    assert [p.key for p in routed.merge] == [frozenset({"e_bare", "e_full"})]
+
+
+def test_bare_name_high_similarity_still_requires_adjudication():
+    bare = make_entity("e_bare", "Johann")
+    full = make_entity("e_full", "Johann Boedecker")
+    adj = FakeAdjudicator({frozenset({"e_bare", "e_full"}): "no"})
+    routed = er2.route_pairs([_scored_pair(bare, full, 0.99)], adj, {})
+    # 0.99 >= high, but bare names never auto-merge: LLM said no → drop
+    assert routed.merge == []
+    assert adj.calls == [frozenset({"e_bare", "e_full"})]
+
+
+def test_bare_name_two_candidates_never_merges():
+    bare = make_entity("e_bare", "Johann")
+    full1 = make_entity("e_f1", "Johann Boedecker")
+    full2 = make_entity("e_f2", "Johann Mueller")
+    adj = FakeAdjudicator({
+        frozenset({"e_bare", "e_f1"}): "yes",
+        frozenset({"e_bare", "e_f2"}): "yes",
+    })
+    routed = er2.route_pairs(
+        [_scored_pair(bare, full1, 0.93), _scored_pair(bare, full2, 0.90)],
+        adj, {})
+    assert routed.merge == []
+    assert len(routed.human_review) == 2
+    assert adj.calls == []   # ambiguity short-circuits before the LLM
+
+
+def test_bare_name_with_no_llm_goes_to_review():
+    bare = make_entity("e_bare", "Johann")
+    full = make_entity("e_full", "Johann Boedecker")
+    routed = er2.route_pairs([_scored_pair(bare, full, 0.95)],
+                             er2.NoLLMAdjudicator(), {})
+    assert routed.merge == []
+    assert len(routed.human_review) == 1
+
+
+# ----------------------------------------------------------------------
+# 4. Proposal assembly (richest-row-wins, same ordering as v1)
+# ----------------------------------------------------------------------
+
+def test_pairs_to_proposals_richest_wins_and_transitive():
+    a = make_entity("e_a", "Johann_Boedecker", facts=2)
+    b = make_entity("e_b", "Bödecker, Johann", facts=50, rels=3)
+    c = make_entity("e_c", "Johann Boedecker", facts=7)
+    p1, p2 = _scored_pair(a, b, 0.95), _scored_pair(b, c, 0.94)
+    proposals = er2.pairs_to_proposals([p1, p2])
+    assert len(proposals) == 1
+    assert proposals[0].canonical.id == "e_b"      # richest
+    assert {d.id for d in proposals[0].deprecated} == {"e_a", "e_c"}
+    assert proposals[0].signal == "embedding_llm"
+
+
+# ----------------------------------------------------------------------
+# 5. Adjudication JSON parsing — strict; malformed → unsure
+# ----------------------------------------------------------------------
+
+def test_parse_adjudication_valid():
+    out = er2.parse_adjudication(
+        '{"same_person": "yes", "reason": "same email domain and role"}')
+    assert out.same_person == "yes"
+    assert "email domain" in out.reason
+
+
+def test_parse_adjudication_json_embedded_in_prose():
+    out = er2.parse_adjudication(
+        'Sure! Here is my answer: {"same_person": "no", "reason": "different people"}')
+    assert out.same_person == "no"
+
+
+def test_parse_adjudication_malformed_is_unsure():
+    assert er2.parse_adjudication("I think they are the same").same_person == "unsure"
+    assert er2.parse_adjudication('{"same_person": "maybe"}').same_person == "unsure"
+    assert er2.parse_adjudication('{"broken json').same_person == "unsure"
+
+
+# ----------------------------------------------------------------------
+# 6. Embedding plumbing (no network — backends are constructed only)
+# ----------------------------------------------------------------------
+
+def test_cosine():
+    assert abs(er2.cosine([1.0, 0.0], [1.0, 0.0]) - 1.0) < 1e-9
+    assert abs(er2.cosine([1.0, 0.0], [0.0, 1.0])) < 1e-9
+    assert er2.cosine([0.0, 0.0], [1.0, 0.0]) == 0.0
+
+
+def test_embedding_bundle_contains_surfaces_and_facts():
+    e = make_entity("e", "Johann Boedecker", aliases=["johann@x.com"])
+    bundle = er2.embedding_bundle(e, ["Johann leads the Berlin office"])
+    assert "Johann Boedecker" in bundle
+    assert "johann@x.com" in bundle
+    assert "Berlin office" in bundle
+
+
+def test_http_backend_requires_url_and_is_never_defaulted():
+    try:
+        er2.HttpEmbeddingBackend("")
+        raise AssertionError("expected ValueError for empty --embed-url")
+    except ValueError:
+        pass
+
+
+def test_local_backend_is_an_explicit_stub():
+    try:
+        er2.LocalEmbeddingBackend()
+        raise AssertionError("expected NotImplementedError")
+    except NotImplementedError as e:
+        assert "--embed-backend http" in str(e)
+
+
+# ----------------------------------------------------------------------
+# 7. Safety gates: --apply refused without --i-have-a-snapshot
+# ----------------------------------------------------------------------
+
+def test_validate_args_refuses_apply_without_snapshot():
+    args = er2.parse_args(["--arena", "test-arena", "--pg-dsn", "x",
+                           "--embed-url", "http://e", "--apply"])
+    err = er2.validate_args(args)
+    assert err is not None and "--i-have-a-snapshot" in err
+
+
+def test_validate_args_accepts_apply_with_snapshot():
+    args = er2.parse_args(["--arena", "test-arena", "--pg-dsn", "x",
+                           "--embed-url", "http://e",
+                           "--apply", "--i-have-a-snapshot"])
+    assert er2.validate_args(args) is None
+
+
+def test_cli_subprocess_refuses_apply_without_snapshot():
+    proc = subprocess.run(
+        [sys.executable, str(_SCRIPTS / "entity_resolution_v2.py"),
+         "--arena", "test-arena", "--pg-dsn", "postgresql://x", "--apply"],
+        capture_output=True, text=True)
+    assert proc.returncode == 2
+    assert "--i-have-a-snapshot" in proc.stderr
+
+
+def test_cli_requires_arena():
+    proc = subprocess.run(
+        [sys.executable, str(_SCRIPTS / "entity_resolution_v2.py"),
+         "--pg-dsn", "postgresql://x"],
+        capture_output=True, text=True)
+    assert proc.returncode == 2
+    assert "--arena" in proc.stderr
+
+
+# ----------------------------------------------------------------------
+# 8. Arena scoping — every SQL statement carries the arena predicate
+# ----------------------------------------------------------------------
+
+def test_v2_sql_registry_is_arena_scoped():
+    assert er2.ARENA_SCOPED_SQL, "registry must not be empty"
+    for name, sql in er2.ARENA_SCOPED_SQL.items():
+        assert "arena = %s" in sql, f"v2 SQL '{name}' lost its arena predicate"
+    er2.assert_arena_scoped()   # must not raise
+
+
+def test_v1_load_and_repoint_sql_is_arena_scoped():
+    src = (_SCRIPTS / "backfill_entity_reconciliation.py").read_text()
+    # Every fact/relationship repoint in apply_proposals is arena-scoped.
+    import re
+    repoints = re.findall(
+        r"UPDATE (?:facts|relationships) SET \w+ = %s\s+WHERE ([^\"]+?)\n",
+        src)
+    assert len(repoints) == 4, "expected the 4 v1 repoint statements"
+    for predicate in repoints:
+        assert "arena = %s" in predicate
+    # The entity load + co-occurrence scan are arena-scoped too.
+    assert "WHERE arena = %s AND entity_type = %s" in src
+    assert '"SELECT attributes FROM events WHERE arena = %s"' in src
+
+
+# ----------------------------------------------------------------------
+# 9. v1 CLI byte-compatibility — v2 must not have changed it
+# ----------------------------------------------------------------------
+
+def test_v1_cli_surface_unchanged():
+    old_argv = sys.argv
+    try:
+        sys.argv = ["backfill_entity_reconciliation.py",
+                    "--arena", "a", "--pg-dsn", "dsn"]
+        ns = v1.parse_args()
+    finally:
+        sys.argv = old_argv
+    assert vars(ns) == {
+        "arena": "a", "pg_dsn": "dsn", "entity_type": "person",
+        "apply": False, "heuristic_merge": False, "out": None,
+        "merged_by": None,
+    }, "v1's CLI namespace changed — it must stay byte-compatible"
+
+
+def test_v1_cli_rejects_v2_only_flags():
+    old_argv = sys.argv
+    try:
+        sys.argv = ["backfill_entity_reconciliation.py",
+                    "--arena", "a", "--pg-dsn", "dsn", "--i-have-a-snapshot"]
+        try:
+            v1.parse_args()
+            raise AssertionError("v1 accepted a v2-only flag")
+        except SystemExit as e:
+            assert e.code == 2
+    finally:
+        sys.argv = old_argv
+
+
+def test_v1_machinery_is_imported_not_copied():
+    # v2 must reuse v1's load/apply machinery, not fork it.
+    assert er2.v1 is v1
+    assert er2.Entity is v1.Entity
+    assert er2.MergeProposal is v1.MergeProposal
+    src = (_SCRIPTS / "entity_resolution_v2.py").read_text()
+    for fn in ("def load_entities", "def apply_proposals",
+               "def collect_cooccurrence_pairs", "def build_proposals"):
+        assert fn not in src, f"v2 copied v1's {fn} instead of importing it"
+
+
+# ----------------------------------------------------------------------
+# 10. Tiered report
+# ----------------------------------------------------------------------
+
+def test_report_tiers_and_counts(tmp_path=None):
+    rich = make_entity("e_rich", "Johann Boedecker", facts=50)
+    dup = make_entity("e_dup", "Johann_Boedecker", facts=1)
+    v1_prop = v1.MergeProposal(canonical=rich, deprecated=[dup],
+                               signal="co_occurrence")
+    m1, m2 = make_entity("e_m1", "Carla Voss", facts=9), make_entity("e_m2", "Carla Vosse")
+    pair = _scored_pair(m1, m2, 0.85)
+    pair.band, pair.verdict, pair.reason = "ambiguous", "yes", "same person per facts"
+    routed = er2.RoutedPairs(merge=[pair])
+    v2_props = er2.pairs_to_proposals([pair])
+
+    records = er2.build_report_records(
+        "test-arena", [v1_prop], v2_props, routed, before_count=10)
+    header = records[0]
+    assert header["arena"] == "test-arena"
+    assert header["entity_count_before"] == 10
+    assert header["entity_count_after_if_applied"] == 8   # 2 deprecations
+    assert "untouched" in header["other_arenas"]
+    assert header["tiers"] == ["co_occurrence", "alias_overlap",
+                               "embedding_llm", "heuristic"]
+
+    proposals = [r for r in records if r["type"] == "merge_proposal"]
+    assert [p["tier"] for p in proposals] == ["co_occurrence", "embedding_llm"]
+    emb = proposals[1]
+    assert emb["evidence"][0]["similarity"] == 0.85
+    assert emb["evidence"][0]["verdict"] == "yes"
+    assert "same person per facts" in emb["evidence"][0]["reason"]
+    # All records JSONL-serializable.
+    for r in records:
+        json.dumps(r)
+
+
+def test_markdown_summary_mentions_frozen_arena(tmp_path=None):
+    import tempfile, os
+    records = er2.build_report_records(
+        "test-arena", [], [], er2.RoutedPairs(), before_count=3)
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "summary.md")
+        er2.write_markdown_summary(records, path)
+        text = Path(path).read_text()
+    assert "Other arenas" in text
+    assert "pip-agents" in text and "FROZEN" in text
+    assert "`test-arena`" in text