"""Entity resolution — deterministic first, never merge blindly. 1. identifiers (scheme, value) → exact entity 2. normalized alias within the same entity type (disambiguated by organization when several match) 3. slug collision within the same type 4. otherwise create, and park ambiguous cases in the review queue as merge candidates """ from __future__ import annotations import logging from datetime import UTC, datetime from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection from aiatlas.db import execute, fetch_all, fetch_one, jsonb from aiatlas.ids import ENTITY_TYPES, new_id, normalize_alias, slugify from aiatlas.sdk.facts import EntityRef log = logging.getLogger(__name__) # Types whose slug should be prefixed by the organization slug to stay unique and readable (models: `qwen-qwen3-8b` is ugly → # we keep the model name and only prefix on collision). GENERIC_NAMES = {"model", "models", "api", "pricing", "docs", "blog", "news", "research", "overview"} class Resolver: def __init__(self, conn: AsyncConnection, *, snapshot_id: str | None = None, source_tier: int = 2): self.conn = conn self.snapshot_id = snapshot_id self.tier = source_tier self._cache: dict[str, str] = {} self.created: list[str] = [] self.updated: set[str] = set() async def resolve(self, ref: EntityRef, *, create: bool = True) -> str | None: if ref.id: return ref.id if ref.entity_type not in ENTITY_TYPES: raise ValueError(f"unknown entity type {ref.entity_type!r}") key = ref.key() if key in self._cache: ref.id = self._cache[key] await self._refresh_links(ref) return ref.id org_id = await self.resolve(ref.organization) if ref.organization else None found = await self._by_identifiers(ref) if found is None: found = await self._by_alias(ref, org_id) if found is None: found = await self._by_slug(ref) if found is None and not create: return None if found is None: found = await self._create(ref, org_id) else: await self._touch(found, ref, org_id) ref.id = found self._cache[key] = found await self._refresh_links(ref) return found # ---------------------------------------------------------------------------------------------- lookups async def _by_identifiers(self, ref: EntityRef) -> str | None: for scheme, value in ref.identifiers.items(): row = await fetch_one(self.conn, """select ei.entity_id, e.entity_type, e.merged_into from entity_identifiers ei join entities e on e.id = ei.entity_id where ei.scheme = :s and ei.value = :v""", s=scheme, v=str(value)) if row: if row["entity_type"] != ref.entity_type: log.warning("identifier type mismatch", extra={"scheme": scheme, "value": value, "have": row["entity_type"], "want": ref.entity_type}) continue return row["merged_into"] or row["entity_id"] return None async def _by_alias(self, ref: EntityRef, org_id: str | None) -> str | None: names = [ref.name, *ref.aliases] norms = {normalize_alias(n) for n in names if n and normalize_alias(n)} if not norms: return None rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into from entity_aliases a join entities e on e.id = a.entity_id where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = :t""", norms=list(norms), t=ref.entity_type) if not rows: return None rows = [{**r, "id": r["merged_into"] or r["id"]} for r in rows] ids = {r["id"] for r in rows} if len(ids) == 1: return rows[0]["id"] if org_id: same_org = [r for r in rows if r["organization_id"] == org_id] if len({r["id"] for r in same_org}) == 1: return same_org[0]["id"] # ambiguous: do not guess — create and flag await self._review("merge_candidate", sorted(ids), f"alias '{ref.name}' matches {len(ids)} {ref.entity_type} entities", {"name": ref.name, "identifiers": ref.identifiers}) return None async def _by_slug(self, ref: EntityRef) -> str | None: slug = ref.slug_hint or slugify(ref.name) row = await fetch_one(self.conn, "select id, entity_type, organization_id, merged_into from entities where slug = :s", s=slug) if row and row["entity_type"] == ref.entity_type: return row["merged_into"] or row["id"] return None # ---------------------------------------------------------------------------------------------- writes async def _create(self, ref: EntityRef, org_id: str | None) -> str: eid = new_id(ref.entity_type) slug = await self._unique_slug(ref, org_id) first_seen = datetime.now(UTC) await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at) values (:id, :t, :n, :slug, :d, :status, :org, :fs, now())""", id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active", org=org_id, fs=first_seen) self.created.append(eid) return eid async def _touch(self, eid: str, ref: EntityRef, org_id: str | None) -> None: sets = ["last_seen_at = now()"] params: dict[str, Any] = {"id": eid} if org_id: sets.append("organization_id = coalesce(organization_id, :org)") params["org"] = org_id if ref.description: sets.append("description = case when description is null or length(description) < 40 then :d else description end") params["d"] = ref.description await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params) self.updated.add(eid) async def _refresh_links(self, ref: EntityRef) -> None: assert ref.id for alias in {ref.name, *ref.aliases}: norm = normalize_alias(alias) if not norm or len(norm) < 2: continue await execute(self.conn, """insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) values (:e, :a, :n, 'alias', :s) on conflict (entity_id, alias_norm) do nothing""", e=ref.id, a=alias.strip()[:300], n=norm, s=self.snapshot_id) for scheme, value in ref.identifiers.items(): await execute(self.conn, """insert into entity_identifiers (entity_id, scheme, value, snapshot_id) values (:e, :s, :v, :snap) on conflict (scheme, value) do nothing""", e=ref.id, s=scheme, v=str(value)[:500], snap=self.snapshot_id) async def _unique_slug(self, ref: EntityRef, org_id: str | None) -> str: base = ref.slug_hint or slugify(ref.name) if base in GENERIC_NAMES or len(base) < 2: base = f"{ref.entity_type}-{base}" candidates = [base] if org_id: org_slug = await fetch_one(self.conn, "select slug from entities where id = :id", id=org_id) if org_slug and not base.startswith(org_slug["slug"]): candidates.append(f"{org_slug['slug']}-{base}") for c in candidates: if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c): return c n = 2 while True: c = f"{base}-{n}" if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c): return c n += 1 async def _review(self, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any]) -> None: dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}" 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) on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe) __all__ = ["Resolver"]