SPB Git

spb/ai-atlas

Public
8.5 KB · 170 lines python
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 _by_alias(self, ref: EntityRef, org_id: str | None) -> str | None:79        names = [ref.name, *ref.aliases]80        norms = {normalize_alias(n) for n in names if n and normalize_alias(n)}81        if not norms:82            return None83        rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into from entity_aliases a84                                            join entities e on e.id = a.entity_id85                                            where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = :t""", norms=list(norms), t=ref.entity_type)86        if not rows:87            return None88        rows = [{**r, "id": r["merged_into"] or r["id"]} for r in rows]89        ids = {r["id"] for r in rows}90        if len(ids) == 1:91            return rows[0]["id"]92        if org_id:93            same_org = [r for r in rows if r["organization_id"] == org_id]94            if len({r["id"] for r in same_org}) == 1:95                return same_org[0]["id"]96        # ambiguous: do not guess — create and flag97        await self._review("merge_candidate", sorted(ids), f"alias '{ref.name}' matches {len(ids)} {ref.entity_type} entities",98                           {"name": ref.name, "identifiers": ref.identifiers})99        return None100101    async def _by_slug(self, ref: EntityRef) -> str | None:102        slug = ref.slug_hint or slugify(ref.name)103        row = await fetch_one(self.conn, "select id, entity_type, organization_id, merged_into from entities where slug = :s", s=slug)104        if row and row["entity_type"] == ref.entity_type:105            return row["merged_into"] or row["id"]106        return None107108    # ---------------------------------------------------------------------------------------------- writes109    async def _create(self, ref: EntityRef, org_id: str | None) -> str:110        eid = new_id(ref.entity_type)111        slug = await self._unique_slug(ref, org_id)112        first_seen = datetime.now(UTC)113        await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at)114                                    values (:id, :t, :n, :slug, :d, :status, :org, :fs, now())""",115                      id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active",116                      org=org_id, fs=first_seen)117        self.created.append(eid)118        return eid119120    async def _touch(self, eid: str, ref: EntityRef, org_id: str | None) -> None:121        sets = ["last_seen_at = now()"]122        params: dict[str, Any] = {"id": eid}123        if org_id:124            sets.append("organization_id = coalesce(organization_id, :org)")125            params["org"] = org_id126        if ref.description:127            sets.append("description = case when description is null or length(description) < 40 then :d else description end")128            params["d"] = ref.description129        await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params)130        self.updated.add(eid)131132    async def _refresh_links(self, ref: EntityRef) -> None:133        assert ref.id134        for alias in {ref.name, *ref.aliases}:135            norm = normalize_alias(alias)136            if not norm or len(norm) < 2:137                continue138            await execute(self.conn, """insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) values (:e, :a, :n, 'alias', :s)139                                        on conflict (entity_id, alias_norm) do nothing""", e=ref.id, a=alias.strip()[:300], n=norm, s=self.snapshot_id)140        for scheme, value in ref.identifiers.items():141            await execute(self.conn, """insert into entity_identifiers (entity_id, scheme, value, snapshot_id) values (:e, :s, :v, :snap)142                                        on conflict (scheme, value) do nothing""", e=ref.id, s=scheme, v=str(value)[:500], snap=self.snapshot_id)143144    async def _unique_slug(self, ref: EntityRef, org_id: str | None) -> str:145        base = ref.slug_hint or slugify(ref.name)146        if base in GENERIC_NAMES or len(base) < 2:147            base = f"{ref.entity_type}-{base}"148        candidates = [base]149        if org_id:150            org_slug = await fetch_one(self.conn, "select slug from entities where id = :id", id=org_id)151            if org_slug and not base.startswith(org_slug["slug"]):152                candidates.append(f"{org_slug['slug']}-{base}")153        for c in candidates:154            if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):155                return c156        n = 2157        while True:158            c = f"{base}-{n}"159            if not await fetch_one(self.conn, "select 1 from entities where slug = :s", s=c):160                return c161            n += 1162163    async def _review(self, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any]) -> None:164        dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}"165        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)166                                    on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe)167168169__all__ = ["Resolver"]170
Files · e6b6b4aeb02f00c4c6493363aecc51e3a804d968
Menu

Theme