SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
19.0 KB · 318 lines python
Raw Blame History
1"""Entity resolution — deterministic first, never merge blindly.23    1. identifiers (scheme, value) → exact entity4    2. normalized alias within the same entity type (disambiguated by organization when several match)5    3. slug collision within the same type6    4. evaluation-effort variants of an existing model ("gpt-5-4-mini-medium") fold onto the canonical model — they are a result7       configuration, never an entity of their own8    5. otherwise create, and park ambiguous cases in the review queue as merge candidates910Guards: persisted `resolution_decisions` (keep_separate) are honoured — two entities an operator kept apart are never re-fused by an11alias match; aliases whose normalisation collapses digits and dots ("Qwen3-8B" ≡ "Qwen 38B") additionally require the same12`variant_key`; `model` and `artifact` are compatible types for lookups (an artifact used to be typed model — the same hub repo must13keep resolving to the same row).14"""15from __future__ import annotations1617import logging18import re19from datetime import UTC, datetime20from typing import Any2122from sqlalchemy.ext.asyncio import AsyncConnection2324from aiatlas.db import execute, fetch_all, fetch_one, jsonb25from aiatlas.ids import ENTITY_TYPES, new_id, normalize_alias, slugify26from aiatlas.ontology.models import analyze_model_name, base_name, variant_key27from aiatlas.sdk.facts import EntityRef2829log = logging.getLogger(__name__)3031# Types whose slug should be prefixed by the organization slug to stay unique and readable (models: `qwen-qwen3-8b` is ugly →32# we keep the model name and only prefix on collision).33GENERIC_NAMES = {"model", "models", "api", "pricing", "docs", "blog", "news", "research", "overview"}34# lookup-compatible type groups: a ref of one type may resolve to a stored entity of another type in the same group35_COMPATIBLE: dict[str, tuple[str, ...]] = {"model": ("model", "artifact"), "artifact": ("artifact", "model")}36# names whose normalised alias is ambiguous: a digit, a separator, a digit ("Qwen3-8B" / "Qwen 38B" → "qwen38b")37_DIGIT_SEP_DIGIT = re.compile(r"\d[.\-\s_]\d")38# identifier schemes issued by evaluators: an entity known ONLY through these may be an evaluation configuration rather than a model39EVALUATOR_SCHEMES = frozenset({"artificial_analysis", "livebench_model_id", "aider_model", "lmarena", "openrouter"})   # OpenRouter lists `o3-mini-high` as an endpoint of o3-mini404142def compatible_types(entity_type: str) -> tuple[str, ...]:43    return _COMPATIBLE.get(entity_type, (entity_type,))444546class Resolver:47    def __init__(self, conn: AsyncConnection, *, snapshot_id: str | None = None, source_tier: int = 2, variant_index: dict[str, str] | None = None):48        self.conn = conn49        self.snapshot_id = snapshot_id50        self.tier = source_tier51        self.variant_index = variant_index          # optional precomputed variant_key → canonical model id (canonicalization)52        self._cache: dict[str, str] = {}53        self.created: list[str] = []54        self.updated: set[str] = set()55        self.folded: dict[str, tuple[str, dict[str, str]]] = {}    # ref key → (canonical id, effort config) for folded variants5657    async def resolve(self, ref: EntityRef, *, create: bool = True) -> str | None:58        if ref.id:59            return ref.id60        if ref.entity_type not in ENTITY_TYPES:61            raise ValueError(f"unknown entity type {ref.entity_type!r}")62        key = ref.key()63        if key in self._cache:64            ref.id = self._cache[key]65            await self._refresh_links(ref)66            return ref.id6768        org_id = await self.resolve(ref.organization) if ref.organization else None6970        found = await self._by_identifiers(ref)71        if found is None:72            found = await self._by_alias(ref, org_id)73        if found is None:74            found = await self._by_slug(ref)75        if found is None and ref.entity_type == "model" and self.tier >= 2 and set(ref.identifiers) <= EVALUATOR_SCHEMES and not ref.attributes.get("hf_repo"):76            # only evaluator-only references fold: a ref carrying official/hub identifiers (or a tier-1 source) names a real model77            variant = await self.resolve_variant(ref, org_id=org_id)78            if variant:79                found = variant[0]80                self.folded[key] = variant81        if found is None and not create:82            return None83        if found is None:84            found = await self._create(ref, org_id)85        else:86            await self._touch(found, ref, org_id)87        ref.id = found88        self._cache[key] = found89        await self._refresh_links(ref)90        return found9192    # ---------------------------------------------------------------------------------------------- lookups93    async def _by_identifiers(self, ref: EntityRef) -> str | None:94        types = compatible_types(ref.entity_type)95        for scheme, value in ref.identifiers.items():96            row = await fetch_one(self.conn, """select ei.entity_id, e.entity_type, e.merged_into from entity_identifiers ei97                                                 join entities e on e.id = ei.entity_id where ei.scheme = :s and ei.value = :v""",98                                  s=scheme, v=str(value))99            if row:100                if row["entity_type"] not in types:101                    log.warning("identifier type mismatch", extra={"scheme": scheme, "value": value, "have": row["entity_type"], "want": ref.entity_type})102                    continue103                return row["merged_into"] or row["entity_id"]104        return None105106    async def _conflicting_identifier(self, eid: str, ref: EntityRef) -> bool:107        """True when `eid` already carries an identifier of a scheme the ref also has, with a different value108        (e.g. two Codestral versions sharing a display name but different API ids) — never fuse those."""109        if not ref.identifiers:110            return False111        rows = await fetch_all(self.conn, "select scheme, value from entity_identifiers where entity_id = :e and scheme = any(cast(:s as text[]))",112                               e=eid, s=list(ref.identifiers))113        return any(str(ref.identifiers[r["scheme"]]) != r["value"] for r in rows)114115    async def _by_alias(self, ref: EntityRef, org_id: str | None) -> str | None:116        names = [ref.name, *ref.aliases]117        norms = {normalize_alias(n) for n in names if n and normalize_alias(n)}118        if not norms:119            return None120        rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into, a.alias from entity_aliases a121                                            join entities e on e.id = a.entity_id122                                            where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = any(cast(:types as text[]))""",123                               norms=list(norms), types=list(compatible_types(ref.entity_type)))124        if not rows:125            return None126        rows = [{**r, "id": r["merged_into"] or r["id"]} for r in rows]127        # collision safety: when the normalised alias erases digit separators, the candidate must share the ref's variant key128        # (or match one of the ref's names textually) — "Qwen 38B" must not land on "Qwen3-8B"129        if ref.entity_type in ("model", "artifact") and any(_DIGIT_SEP_DIGIT.search(n) for n in names if n):130            lower_names = {n.strip().lower() for n in names if n}131            ref_vk = variant_key(ref.name)132            safe = []133            for r in rows:134                if (r["alias"] or "").strip().lower() in lower_names or variant_key(r["canonical_name"]) == ref_vk:135                    safe.append(r)136                else:137                    log.info("alias collision refused", extra={"ref_name": ref.name, "candidate": r["canonical_name"]})   # `name` is reserved by LogRecord138            rows = safe139            if not rows:140                return None141        kept = []142        for r in rows:143            if await self._conflicting_identifier(r["id"], ref):144                await self._review("merge_candidate", sorted({r["id"]}), f"'{ref.name}' shares a name with {r['canonical_name']} but has different identifiers",145                                   {"name": ref.name, "identifiers": ref.identifiers})146            else:147                kept.append(r)148        rows = kept149        if not rows:150            return None151        ids = {r["id"] for r in rows}152        # persisted decisions: a candidate an operator kept separate from another entity is only accepted with a matching organisation153        separate = await self._kept_separate(sorted(ids))154        if separate:155            rows = [r for r in rows if r["id"] not in separate or (org_id and r["organization_id"] == org_id)]156            ids = {r["id"] for r in rows}157            if not rows:158                return None159        if len(ids) == 1:160            return rows[0]["id"]161        if org_id:162            same_org = [r for r in rows if r["organization_id"] == org_id]163            if len({r["id"] for r in same_org}) == 1:164                return same_org[0]["id"]165        # ambiguous: do not guess — create and flag166        await self._review("merge_candidate", sorted(ids), f"alias '{ref.name}' matches {len(ids)} {ref.entity_type} entities",167                           {"name": ref.name, "identifiers": ref.identifiers})168        return None169170    async def _kept_separate(self, ids: list[str]) -> set[str]:171        """Ids among `ids` that carry a `keep_separate` decision with any entity."""172        if not ids:173            return set()174        rows = await fetch_all(self.conn, """select a_id, b_id from resolution_decisions where decision = 'keep_separate'175                                            and (a_id = any(cast(:ids as text[])) or b_id = any(cast(:ids as text[])))""", ids=ids)176        out: set[str] = set()177        for r in rows:178            out.update(x for x in (r["a_id"], r["b_id"]) if x in ids)179        return out180181    async def kept_separate(self, a: str, b: str) -> bool:182        row = await fetch_one(self.conn, """select 1 from resolution_decisions where decision = 'keep_separate'183                                            and ((a_id = :a and b_id = :b) or (a_id = :b and b_id = :a)) limit 1""", a=a, b=b)184        return row is not None185186    async def _by_slug(self, ref: EntityRef) -> str | None:187        slug = ref.slug_hint or slugify(ref.name)188        row = await fetch_one(self.conn, "select id, entity_type, organization_id, merged_into from entities where slug = :s", s=slug)189        if row and row["entity_type"] in compatible_types(ref.entity_type) and not await self._conflicting_identifier(row["merged_into"] or row["id"], ref):190            return row["merged_into"] or row["id"]191        return None192193    async def resolve_variant(self, ref: EntityRef, *, org_id: str | None = None) -> tuple[str, dict[str, str]] | None:194        """Evaluation-effort variant ("claude-opus-5-medium", "qwen3-6-27b-non-reasoning") → (canonical model id, effort config) when a195        canonical model exists: same `artificial_analysis` identifier or alias as the base name, or the same `variant_key`196        (through the optional precomputed index). Returns None when the base cannot be resolved or is ambiguous."""197        a = analyze_model_name(ref.name)198        if not a.is_effort_variant:199            return None200        base = base_name(ref.name)201        bases = {base, base.lower(), a.base_key, slugify(base)}202        # a dated snapshot in the base ("gpt-5.2-2025-12-11-high") also names the undated model203        undated = re.sub(r"[-_ ]?(20\d{2}[-_.]?\d{2}[-_.]?\d{2}|\d{4})$", "", base).strip("-_ ")204        if undated and undated != base and re.search(r"[a-z]", undated, re.I):205            bases |= {undated, undated.lower(), slugify(undated)}206        bases = sorted(bases)207        base_norms = sorted({normalize_alias(b) for b in bases} - {""})208        # the base may be known by any alias, any identifier value (AA slug, api model id, provider id…), its slug or its api_model_id209        rows = await fetch_all(self.conn, """select distinct coalesce(e.merged_into, e.id) as id, e.canonical_name, e.organization_id from entities e210                                            left join entity_identifiers ei on ei.entity_id = e.id211                                            left join entity_aliases al on al.entity_id = e.id212                                            where e.entity_type = 'model'213                                              and (ei.value = any(cast(:bases as text[])) or al.alias_norm = any(cast(:norms as text[]))214                                                   or e.slug = any(cast(:bases as text[])) or lower(e.attributes->>'api_model_id') = any(cast(:bases as text[])))""",215                               bases=bases, norms=base_norms)216        candidates = {r["id"]: r for r in rows if r["id"] != ref.id}217        if self.variant_index:218            vk = variant_key(base)219            cid = self.variant_index.get(vk)220            if cid and cid != ref.id and cid not in candidates:221                row = await fetch_one(self.conn, "select id, canonical_name, organization_id, merged_into from entities where id = :id and merged_into is null", id=cid)222                if row:223                    candidates[cid] = row224        # candidates must be live models (a merged row resolved to its survivor above) and never another effort variant225        live = {r["id"] for r in await fetch_all(self.conn, "select id from entities where id = any(cast(:ids as text[])) and merged_into is null and entity_type = 'model'",226                                                 ids=sorted(candidates))} if candidates else set()227        candidates = {k: v for k, v in candidates.items() if k in live and not analyze_model_name(v["canonical_name"]).is_effort_variant}228        if not candidates:229            return None230        if len(candidates) > 1:231            # prefer the candidate an evaluator/vendor identifies by the exact base name, then the variant's organisation232            exact = await fetch_all(self.conn, "select distinct entity_id from entity_identifiers where value = any(cast(:b as text[])) and entity_id = any(cast(:ids as text[]))",233                                    b=bases, ids=sorted(candidates))234            exact_ids = {r["entity_id"] for r in exact}235            if len(exact_ids) == 1:236                candidates = {k: v for k, v in candidates.items() if k in exact_ids}237        if len(candidates) > 1 and org_id:238            same = {k: v for k, v in candidates.items() if v["organization_id"] == org_id}239            if same:240                candidates = same241        if len(candidates) != 1:242            await self._review("merge_candidate", sorted(candidates), f"effort variant '{ref.name}' matches {len(candidates)} base models",243                               {"name": ref.name, "base": base})244            return None245        cid = next(iter(candidates))246        return cid, dict(a.effort)247248    # ---------------------------------------------------------------------------------------------- writes249    async def _create(self, ref: EntityRef, org_id: str | None) -> str:250        eid = new_id(ref.entity_type)251        slug = await self._unique_slug(ref, org_id)252        now = datetime.now(UTC)253        first_seen = now254        if ref.first_seen_hint is not None:255            hint = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC)256            first_seen = min(now, hint)257        await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at,258                                    identity_confidence, artifact_kind)259                                    values (:id, :t, :n, :slug, :d, :status, :org, :fs, now(), :ic, :ak)""",260                      id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active",261                      org=org_id, fs=first_seen, ic=ref.identity_confidence or "high", ak=ref.artifact_kind if ref.entity_type == "artifact" else None)262        self.created.append(eid)263        return eid264265    async def _touch(self, eid: str, ref: EntityRef, org_id: str | None) -> None:266        sets = ["last_seen_at = now()"]267        params: dict[str, Any] = {"id": eid}268        if org_id:269            sets.append("organization_id = coalesce(organization_id, :org)")270            params["org"] = org_id271        if ref.description:272            sets.append("description = case when description is null or length(description) < 40 then :d else description end")273            params["d"] = ref.description274        if ref.first_seen_hint is not None:275            sets.append("first_seen_at = least(first_seen_at, :fs)")276            params["fs"] = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC)277        await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params)278        self.updated.add(eid)279280    async def _refresh_links(self, ref: EntityRef) -> None:281        assert ref.id282        for alias in {ref.name, *ref.aliases}:283            norm = normalize_alias(alias)284            if not norm or len(norm) < 2:285                continue286            await execute(self.conn, """insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) values (:e, :a, :n, 'alias', :s)287                                        on conflict (entity_id, alias_norm) do nothing""", e=ref.id, a=alias.strip()[:300], n=norm, s=self.snapshot_id)288        for scheme, value in ref.identifiers.items():289            await execute(self.conn, """insert into entity_identifiers (entity_id, scheme, value, snapshot_id) values (:e, :s, :v, :snap)290                                        on conflict (scheme, value) do nothing""", e=ref.id, s=scheme, v=str(value)[:500], snap=self.snapshot_id)291292    async def _unique_slug(self, ref: EntityRef, org_id: str | None) -> str:293        base = ref.slug_hint or slugify(ref.name)294        if base in GENERIC_NAMES or len(base) < 2:295            base = f"{ref.entity_type}-{base}"296        candidates = [base]297        if org_id:298            org_slug = await fetch_one(self.conn, "select slug from entities where id = :id", id=org_id)299            if org_slug and not base.startswith(org_slug["slug"]):300                candidates.append(f"{org_slug['slug']}-{base}")301        for c in candidates:302            if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):303                return c304        n = 2305        while True:306            c = f"{base}-{n}"307            if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):308                return c309            n += 1310311    async def _review(self, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any]) -> None:312        dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}"313        await execute(self.conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key) values (:id, :k, :ids, cast(:p as jsonb), :r, :d)314                                    on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe)315316317__all__ = ["EVALUATOR_SCHEMES", "Resolver", "compatible_types"]318