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. otherwise create, and park ambiguous cases in the review queue as merge candidates7"""8from __future__ import annotations910import logging11from datetime import UTC, datetime12from typing import Any1314from sqlalchemy.ext.asyncio import AsyncConnection1516from aiatlas.db import execute, fetch_all, fetch_one, jsonb17from aiatlas.ids import ENTITY_TYPES, new_id, normalize_alias, slugify18from aiatlas.sdk.facts import EntityRef1920log = logging.getLogger(__name__)2122# Types whose slug should be prefixed by the organization slug to stay unique and readable (models: `qwen-qwen3-8b` is ugly →23# we keep the model name and only prefix on collision).24GENERIC_NAMES = {"model", "models", "api", "pricing", "docs", "blog", "news", "research", "overview"}252627class Resolver:28 def __init__(self, conn: AsyncConnection, *, snapshot_id: str | None = None, source_tier: int = 2):29 self.conn = conn30 self.snapshot_id = snapshot_id31 self.tier = source_tier32 self._cache: dict[str, str] = {}33 self.created: list[str] = []34 self.updated: set[str] = set()3536 async def resolve(self, ref: EntityRef, *, create: bool = True) -> str | None:37 if ref.id:38 return ref.id39 if ref.entity_type not in ENTITY_TYPES:40 raise ValueError(f"unknown entity type {ref.entity_type!r}")41 key = ref.key()42 if key in self._cache:43 ref.id = self._cache[key]44 await self._refresh_links(ref)45 return ref.id4647 org_id = await self.resolve(ref.organization) if ref.organization else None4849 found = await self._by_identifiers(ref)50 if found is None:51 found = await self._by_alias(ref, org_id)52 if found is None:53 found = await self._by_slug(ref)54 if found is None and not create:55 return None56 if found is None:57 found = await self._create(ref, org_id)58 else:59 await self._touch(found, ref, org_id)60 ref.id = found61 self._cache[key] = found62 await self._refresh_links(ref)63 return found6465 # ---------------------------------------------------------------------------------------------- lookups66 async def _by_identifiers(self, ref: EntityRef) -> str | None:67 for scheme, value in ref.identifiers.items():68 row = await fetch_one(self.conn, """select ei.entity_id, e.entity_type, e.merged_into from entity_identifiers ei69 join entities e on e.id = ei.entity_id where ei.scheme = :s and ei.value = :v""",70 s=scheme, v=str(value))71 if row:72 if row["entity_type"] != ref.entity_type:73 log.warning("identifier type mismatch", extra={"scheme": scheme, "value": value, "have": row["entity_type"], "want": ref.entity_type})74 continue75 return row["merged_into"] or row["entity_id"]76 return None7778 async def _conflicting_identifier(self, eid: str, ref: EntityRef) -> bool:79 """True when `eid` already carries an identifier of a scheme the ref also has, with a different value80 (e.g. two Codestral versions sharing a display name but different API ids) — never fuse those."""81 if not ref.identifiers:82 return False83 rows = await fetch_all(self.conn, "select scheme, value from entity_identifiers where entity_id = :e and scheme = any(cast(:s as text[]))",84 e=eid, s=list(ref.identifiers))85 return any(str(ref.identifiers[r["scheme"]]) != r["value"] for r in rows)8687 async def _by_alias(self, ref: EntityRef, org_id: str | None) -> str | None:88 names = [ref.name, *ref.aliases]89 norms = {normalize_alias(n) for n in names if n and normalize_alias(n)}90 if not norms:91 return None92 rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into from entity_aliases a93 join entities e on e.id = a.entity_id94 where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = :t""", norms=list(norms), t=ref.entity_type)95 if not rows:96 return None97 rows = [{**r, "id": r["merged_into"] or r["id"]} for r in rows]98 kept = []99 for r in rows:100 if await self._conflicting_identifier(r["id"], ref):101 await self._review("merge_candidate", sorted({r["id"]}), f"'{ref.name}' shares a name with {r['canonical_name']} but has different identifiers",102 {"name": ref.name, "identifiers": ref.identifiers})103 else:104 kept.append(r)105 rows = kept106 if not rows:107 return None108 ids = {r["id"] for r in rows}109 if len(ids) == 1:110 return rows[0]["id"]111 if org_id:112 same_org = [r for r in rows if r["organization_id"] == org_id]113 if len({r["id"] for r in same_org}) == 1:114 return same_org[0]["id"]115 # ambiguous: do not guess — create and flag116 await self._review("merge_candidate", sorted(ids), f"alias '{ref.name}' matches {len(ids)} {ref.entity_type} entities",117 {"name": ref.name, "identifiers": ref.identifiers})118 return None119120 async def _by_slug(self, ref: EntityRef) -> str | None:121 slug = ref.slug_hint or slugify(ref.name)122 row = await fetch_one(self.conn, "select id, entity_type, organization_id, merged_into from entities where slug = :s", s=slug)123 if row and row["entity_type"] == ref.entity_type and not await self._conflicting_identifier(row["merged_into"] or row["id"], ref):124 return row["merged_into"] or row["id"]125 return None126127 # ---------------------------------------------------------------------------------------------- writes128 async def _create(self, ref: EntityRef, org_id: str | None) -> str:129 eid = new_id(ref.entity_type)130 slug = await self._unique_slug(ref, org_id)131 first_seen = datetime.now(UTC)132 await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at)133 values (:id, :t, :n, :slug, :d, :status, :org, :fs, now())""",134 id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active",135 org=org_id, fs=first_seen)136 self.created.append(eid)137 return eid138139 async def _touch(self, eid: str, ref: EntityRef, org_id: str | None) -> None:140 sets = ["last_seen_at = now()"]141 params: dict[str, Any] = {"id": eid}142 if org_id:143 sets.append("organization_id = coalesce(organization_id, :org)")144 params["org"] = org_id145 if ref.description:146 sets.append("description = case when description is null or length(description) < 40 then :d else description end")147 params["d"] = ref.description148 await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params)149 self.updated.add(eid)150151 async def _refresh_links(self, ref: EntityRef) -> None:152 assert ref.id153 for alias in {ref.name, *ref.aliases}:154 norm = normalize_alias(alias)155 if not norm or len(norm) < 2:156 continue157 await execute(self.conn, """insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) values (:e, :a, :n, 'alias', :s)158 on conflict (entity_id, alias_norm) do nothing""", e=ref.id, a=alias.strip()[:300], n=norm, s=self.snapshot_id)159 for scheme, value in ref.identifiers.items():160 await execute(self.conn, """insert into entity_identifiers (entity_id, scheme, value, snapshot_id) values (:e, :s, :v, :snap)161 on conflict (scheme, value) do nothing""", e=ref.id, s=scheme, v=str(value)[:500], snap=self.snapshot_id)162163 async def _unique_slug(self, ref: EntityRef, org_id: str | None) -> str:164 base = ref.slug_hint or slugify(ref.name)165 if base in GENERIC_NAMES or len(base) < 2:166 base = f"{ref.entity_type}-{base}"167 candidates = [base]168 if org_id:169 org_slug = await fetch_one(self.conn, "select slug from entities where id = :id", id=org_id)170 if org_slug and not base.startswith(org_slug["slug"]):171 candidates.append(f"{org_slug['slug']}-{base}")172 for c in candidates:173 if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):174 return c175 n = 2176 while True:177 c = f"{base}-{n}"178 if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):179 return c180 n += 1181182 async def _review(self, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any]) -> None:183 dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}"184 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)185 on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe)186187188__all__ = ["Resolver"]189