Livrare LOT 1 - Didi

ai_platform/modules/didi_brain/lint/pairs.py

@@ -0,0 +1,93 @@
+"""Candidate pair generation for the Lint pass.
+
+For each source claim atom, we pull its nearest neighbors via the Atomic
+/api/atoms/{id}/similar endpoint (which uses pgvector kNN under the hood).
+Only Type/Claim neighbors are kept; Type/Document neighbors are dropped so
+we only compare claim-to-claim.
+
+Pairs are canonicalized as (lower_id, higher_id) so (A, B) and (B, A)
+produce the same CandidatePair and don't get evaluated twice.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from shared.atomic_api import AtomicClient
+from shared.logging import get_logger
+
+log = get_logger(__name__)
+
+# Atoms with similarity below this are considered too unrelated to be worth
+# an NLI call. Between 0.55 and 0.95 is the interesting band — below is
+# "probably different topic", above is "almost certainly the same text".
+MIN_PAIR_SIMILARITY = 0.55
+NEIGHBORS_PER_CLAIM = 15
+
+
+@dataclass(slots=True, frozen=True)
+class CandidatePair:
+    """A pair of claim atom ids to evaluate, canonicalized lo<hi."""
+
+    atom_a_id: str
+    atom_b_id: str
+    similarity: float
+
+
+def _is_claim_hit(hit_tags: list[dict]) -> bool:
+    return any(t.get("name") == "Claim" for t in hit_tags)
+
+
+async def generate_candidates(
+    atomic: AtomicClient,
+    *,
+    source_atom_ids: list[str],
+    min_similarity: float = MIN_PAIR_SIMILARITY,
+    neighbors_per: int = NEIGHBORS_PER_CLAIM,
+) -> list[CandidatePair]:
+    """Walk each source atom and collect canonical candidate pairs.
+
+    Limitation: Atomic's /similar endpoint returns Document and Claim atoms
+    mixed together. We filter client-side for Type/Claim. It might be worth
+    upstream adding a tag filter later, but for this corpus size the current
+    approach is fine.
+    """
+    seen: set[tuple[str, str]] = set()
+    out: list[CandidatePair] = []
+
+    for idx, src_id in enumerate(source_atom_ids):
+        try:
+            hits = await atomic.find_similar(
+                src_id, threshold=min_similarity, limit=neighbors_per
+            )
+        except Exception as e:  # noqa: BLE001
+            log.warning(
+                "find_similar_failed", atom_id=src_id, error=f"{type(e).__name__}:{e}"
+            )
+            continue
+
+        for h in hits:
+            if h.atom_id == src_id:
+                continue
+            if not _is_claim_hit(h.tags):
+                continue
+            lo, hi = sorted((src_id, h.atom_id))
+            key = (lo, hi)
+            if key in seen:
+                continue
+            seen.add(key)
+            out.append(
+                CandidatePair(
+                    atom_a_id=lo, atom_b_id=hi, similarity=round(h.similarity, 4)
+                )
+            )
+
+        if (idx + 1) % 50 == 0:
+            log.info(
+                "candidates_progress",
+                processed=idx + 1,
+                total=len(source_atom_ids),
+                pairs=len(out),
+            )
+
+    return out