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"})404142def 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={"name": ref.name, "candidate": r["canonical_name"]})138 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 base_norms = {normalize_alias(base), normalize_alias(a.base_key)} - {""}202 rows = await fetch_all(self.conn, """select distinct e.id, e.canonical_name, e.organization_id, e.merged_into from entities e203 left join entity_identifiers ei on ei.entity_id = e.id204 left join entity_aliases al on al.entity_id = e.id205 where e.entity_type = 'model' and e.merged_into is null206 and ((ei.scheme = 'artificial_analysis' and ei.value = any(cast(:bases as text[])))207 or al.alias_norm = any(cast(:norms as text[])) or e.slug = any(cast(:bases as text[])))""",208 bases=sorted({base.lower(), a.base_key}), norms=sorted(base_norms))209 candidates = {r["id"]: r for r in rows if r["id"] != ref.id}210 if self.variant_index:211 vk = variant_key(base)212 cid = self.variant_index.get(vk)213 if cid and cid != ref.id and cid not in candidates:214 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)215 if row:216 candidates[cid] = row217 # never fold onto another effort variant218 candidates = {k: v for k, v in candidates.items() if not analyze_model_name(v["canonical_name"]).is_effort_variant}219 if not candidates:220 return None221 if len(candidates) > 1 and org_id:222 same = {k: v for k, v in candidates.items() if v["organization_id"] == org_id}223 if same:224 candidates = same225 if len(candidates) != 1:226 await self._review("merge_candidate", sorted(candidates), f"effort variant '{ref.name}' matches {len(candidates)} base models",227 {"name": ref.name, "base": base})228 return None229 cid = next(iter(candidates))230 return cid, dict(a.effort)231232 # ---------------------------------------------------------------------------------------------- writes233 async def _create(self, ref: EntityRef, org_id: str | None) -> str:234 eid = new_id(ref.entity_type)235 slug = await self._unique_slug(ref, org_id)236 now = datetime.now(UTC)237 first_seen = now238 if ref.first_seen_hint is not None:239 hint = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC)240 first_seen = min(now, hint)241 await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at,242 identity_confidence, artifact_kind)243 values (:id, :t, :n, :slug, :d, :status, :org, :fs, now(), :ic, :ak)""",244 id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active",245 org=org_id, fs=first_seen, ic=ref.identity_confidence or "high", ak=ref.artifact_kind if ref.entity_type == "artifact" else None)246 self.created.append(eid)247 return eid248249 async def _touch(self, eid: str, ref: EntityRef, org_id: str | None) -> None:250 sets = ["last_seen_at = now()"]251 params: dict[str, Any] = {"id": eid}252 if org_id:253 sets.append("organization_id = coalesce(organization_id, :org)")254 params["org"] = org_id255 if ref.description:256 sets.append("description = case when description is null or length(description) < 40 then :d else description end")257 params["d"] = ref.description258 if ref.first_seen_hint is not None:259 sets.append("first_seen_at = least(first_seen_at, :fs)")260 params["fs"] = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC)261 await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params)262 self.updated.add(eid)263264 async def _refresh_links(self, ref: EntityRef) -> None:265 assert ref.id266 for alias in {ref.name, *ref.aliases}:267 norm = normalize_alias(alias)268 if not norm or len(norm) < 2:269 continue270 await execute(self.conn, """insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) values (:e, :a, :n, 'alias', :s)271 on conflict (entity_id, alias_norm) do nothing""", e=ref.id, a=alias.strip()[:300], n=norm, s=self.snapshot_id)272 for scheme, value in ref.identifiers.items():273 await execute(self.conn, """insert into entity_identifiers (entity_id, scheme, value, snapshot_id) values (:e, :s, :v, :snap)274 on conflict (scheme, value) do nothing""", e=ref.id, s=scheme, v=str(value)[:500], snap=self.snapshot_id)275276 async def _unique_slug(self, ref: EntityRef, org_id: str | None) -> str:277 base = ref.slug_hint or slugify(ref.name)278 if base in GENERIC_NAMES or len(base) < 2:279 base = f"{ref.entity_type}-{base}"280 candidates = [base]281 if org_id:282 org_slug = await fetch_one(self.conn, "select slug from entities where id = :id", id=org_id)283 if org_slug and not base.startswith(org_slug["slug"]):284 candidates.append(f"{org_slug['slug']}-{base}")285 for c in candidates:286 if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):287 return c288 n = 2289 while True:290 c = f"{base}-{n}"291 if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):292 return c293 n += 1294295 async def _review(self, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any]) -> None:296 dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}"297 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)298 on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe)299300301__all__ = ["EVALUATOR_SCHEMES", "Resolver", "compatible_types"]302