"""Entity merging (curation): fold a duplicate `source` into `target`. Nothing is deleted — the source row stays with `status='merged'` and `merged_into=target` so old slugs and ids keep resolving; every dependent row is re-pointed. Modes merge (default) full merge; records a `merge` decision alias same as merge; records an `alias` decision (the source was another name of the target) variant the source is an *artifact* (quantisation / conversion / packaging) of the target: it keeps its own row, slug and facts, gets `entity_type='artifact'`, `canonical_id=target` and an `artifact_of` relation family_member the target is a `model_family`: the source gets `family_id=target` and a `member_of_family` relation Every call writes a `resolution_decisions` row (applied=true) and an `admin_audit_log` row; `keep_separate` decisions block merges. """ from __future__ import annotations import json from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection from aiatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb from aiatlas.ids import new_id, normalize_alias from aiatlas.ontology import benchmarks as bench_ontology ORG_TYPES = {"company", "organization", "lab", "university"} MODES = ("merge", "alias", "variant", "family_member") DERIVED_SOURCE_KEY = "ai-atlas.registry" async def registry_source_id(conn: AsyncConnection) -> str | None: row = await fetch_one(conn, "select id from sources where key = :k", k=DERIVED_SOURCE_KEY) return row["id"] if row else None async def record_decision(conn: AsyncConnection, a_id: str, b_id: str, decision: str, *, actor: str = "curation", note: str | None = None, payload: dict[str, Any] | None = None, applied: bool = True) -> None: await execute(conn, """insert into resolution_decisions (id, a_id, b_id, decision, actor, note, payload, applied) values (:id, :a, :b, :d, :actor, :note, cast(:p as jsonb), :applied) on conflict (a_id, b_id, decision) do update set applied = resolution_decisions.applied or excluded.applied, note = coalesce(excluded.note, resolution_decisions.note), payload = resolution_decisions.payload || excluded.payload""", id=f"rd_{new_id('review').split('_', 1)[1]}", a=a_id, b=b_id, d=decision, actor=actor, note=note, p=jsonb(payload or {}), applied=applied) async def audit(conn: AsyncConnection, action: str, target: str | None, payload: dict[str, Any] | None = None, *, actor: str = "curation") -> None: await execute(conn, "insert into admin_audit_log (actor, action, target, payload) values (:actor, :action, :target, cast(:p as jsonb))", actor=actor, action=action, target=target, p=jsonb(payload or {})) async def kept_separate(conn: AsyncConnection, a: str, b: str) -> bool: row = await fetch_one(conn, """select 1 from resolution_decisions where decision = 'keep_separate' and ((a_id = :a and b_id = :b) or (a_id = :b and b_id = :a)) limit 1""", a=a, b=b) return row is not None async def upsert_relation(conn: AsyncConnection, subject_id: str, predicate: str, object_id: str, attributes: dict[str, Any] | None = None, *, source_id: str | None = None, tier: int = 2, confidence: str = "high") -> bool: """Insert a live relation unless an identical live edge exists. Returns True when a row was inserted.""" if subject_id == object_id: return False existing = await fetch_one(conn, "select id from relations where subject_id = :s and predicate = :p and object_id = :o and valid_to is null", s=subject_id, p=predicate, o=object_id) if existing: if attributes: await execute(conn, "update relations set attributes = attributes || cast(:a as jsonb) where id = :id", a=jsonb(attributes), id=existing["id"]) return False await execute(conn, """insert into relations (id, subject_id, predicate, object_id, attributes, source_id, tier, confidence, observed_at, valid_from) values (:id, :s, :p, :o, cast(:a as jsonb), :src, :tier, :conf, now(), now())""", id=new_id("relation"), s=subject_id, p=predicate, o=object_id, a=jsonb(attributes or {}), src=source_id, tier=tier, conf=confidence) return True async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str, *, mode: str = "merge", actor: str = "curation", note: str | None = None, payload: dict[str, Any] | None = None) -> dict[str, Any]: if mode not in MODES: raise ValueError(f"unknown merge mode {mode!r}") if source_id == target_id: raise ValueError("source and target are the same entity") src = await fetch_one(conn, "select id, entity_type, canonical_name, slug, merged_into, attributes, provenance from entities where id = :id", id=source_id) dst = await fetch_one(conn, "select id, entity_type, canonical_name, merged_into from entities where id = :id", id=target_id) if not src or not dst: raise LookupError("source or target entity not found") if dst["merged_into"]: raise ValueError("target is itself merged; merge into its survivor instead") if src["merged_into"]: raise ValueError("source is already merged") if mode in ("merge", "alias") and await kept_separate(conn, source_id, target_id): raise ValueError("a keep_separate decision exists for this pair; refusing to merge") if mode == "variant": return await _mark_artifact(conn, src, dst, actor=actor, note=note, payload=payload) if mode == "family_member": return await _mark_family_member(conn, src, dst, actor=actor, note=note, payload=payload) same_type = src["entity_type"] == dst["entity_type"] org_pair = src["entity_type"] in ORG_TYPES and dst["entity_type"] in ORG_TYPES model_pair = {src["entity_type"], dst["entity_type"]} <= {"model", "artifact"} if not (same_type or org_pair or model_pair): raise ValueError(f"cannot merge a {src['entity_type']} into a {dst['entity_type']}") moved: dict[str, int] = {} async def count_update(label: str, sql: str) -> None: n = await fetch_val(conn, f"with u as ({sql} returning 1) select count(*) from u", s=source_id, t=target_id) moved[label] = int(n or 0) await count_update("aliases", "insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) " "select :t, alias, alias_norm, kind, snapshot_id from entity_aliases where entity_id = :s on conflict (entity_id, alias_norm) do nothing") await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:t, :a, :n, 'former_name') " "on conflict (entity_id, alias_norm) do update set kind = 'former_name'", t=target_id, a=src["canonical_name"], n=normalize_alias(src["canonical_name"])) await execute(conn, "delete from entity_aliases where entity_id = :s", s=source_id) await count_update("identifiers", "update entity_identifiers x set entity_id = :t where entity_id = :s " "and not exists (select 1 from entity_identifiers y where y.scheme = x.scheme and y.value = x.value and y.entity_id = :t)") await execute(conn, "delete from entity_identifiers where entity_id = :s", s=source_id) await count_update("claims", "update claims set entity_id = :t where entity_id = :s") # relations: drop those that would duplicate a live edge on the target, then re-point await execute(conn, """update relations r set valid_to = now() where r.valid_to is null and (r.subject_id = :s or r.object_id = :s) and exists ( select 1 from relations x where x.valid_to is null and x.predicate = r.predicate and x.subject_id = case when r.subject_id = :s then :t else r.subject_id end and x.object_id = case when r.object_id = :s then :t else r.object_id end)""", s=source_id, t=target_id) await count_update("relations", "update relations set subject_id = case when subject_id = :s then :t else subject_id end, " "object_id = case when object_id = :s then :t else object_id end where subject_id = :s or object_id = :s") await execute(conn, "update relations set valid_to = now() where subject_id = object_id and valid_to is null and subject_id = :t", t=target_id) await count_update("events", "update change_events set entity_id = :t where entity_id = :s") await execute(conn, "update prices set valid_to = now() where valid_to is null and model_id = :s and exists (select 1 from prices q where q.valid_to is null " "and q.model_id = :t and q.provider_id = prices.provider_id and coalesce(q.provider_model_id,'') = coalesce(prices.provider_model_id,''))", s=source_id, t=target_id) await count_update("prices", "update prices set model_id = case when model_id = :s then :t else model_id end, " "provider_id = case when provider_id = :s then :t else provider_id end where model_id = :s or provider_id = :s") await count_update("results", "update benchmark_results set model_id = case when model_id = :s then :t else model_id end, " "benchmark_id = case when benchmark_id = :s then :t else benchmark_id end where model_id = :s or benchmark_id = :s") moved["result_dedupe_collisions"] = await recompute_result_keys(conn, target_id, source_id=source_id) moved["results_closed"] = await enforce_current_results(conn, model_id=target_id) await count_update("documents", "update documents set entity_id = :t where entity_id = :s") await count_update("children", "update entities set organization_id = :t where organization_id = :s") await count_update("family_members", "update entities set family_id = :t where family_id = :s") await count_update("artifacts", "update entities set canonical_id = :t where canonical_id = :s") await count_update("llm_jobs", "update llm_jobs set entity_id = :t where entity_id = :s") await count_update("sources", "update sources set organization_id = :t where organization_id = :s") await count_update("domains", "update domains set organization_id = :t where organization_id = :s") await count_update("review_items", "update review_queue set entity_ids = array_replace(entity_ids, :s, :t) where :s = any(entity_ids)") # embeddings (table only exists where pgvector is installed): the target keeps its own vector; a source-only vector moves over if await fetch_val(conn, "select to_regclass('entity_embeddings') is not null"): has_target_vec = await fetch_one(conn, "select 1 from entity_embeddings where entity_id = :t", t=target_id) if has_target_vec: await count_update("embeddings", "delete from entity_embeddings where entity_id = :s") else: await count_update("embeddings", "update entity_embeddings set entity_id = :t where entity_id = :s") # attributes the target lacks are inherited (with their provenance); target values always win await execute(conn, """update entities t set attributes = coalesce(s.attributes, '{}'::jsonb) || t.attributes, provenance = coalesce(s.provenance, '{}'::jsonb) || t.provenance, last_seen_at = greatest(t.last_seen_at, s.last_seen_at), first_seen_at = least(t.first_seen_at, s.first_seen_at), updated_at = now() from entities s where t.id = :t and s.id = :s""", s=source_id, t=target_id) await execute(conn, "update entities set merged_into = :t, status = 'merged', canonical_id = null, family_id = null, updated_at = now() where id = :s", s=source_id, t=target_id) await execute(conn, "update entities set merged_into = :t where merged_into = :s", s=source_id, t=target_id) await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, connector_name, dedupe_key, meta, is_backfill) values (:id, :t, 'ENTITY_MERGED', 'source', :sum, 0, 'curation', :dk, cast(:m as jsonb), true) on conflict (dedupe_key) do nothing""", id=new_id("change_event"), t=target_id, sum=f"Merged duplicate '{src['canonical_name']}' ({src['slug']})", dk=f"merge:{source_id}:{target_id}", m=json.dumps({"source_id": source_id, "source_slug": src["slug"], "mode": mode})) decision = "alias" if mode == "alias" else "merge" await record_decision(conn, source_id, target_id, decision, actor=actor, note=note, payload={"moved": moved, **(payload or {})}) await audit(conn, f"entity.{decision}", target_id, {"source_id": source_id, "source_slug": src["slug"], "target_id": target_id, "moved": moved, "note": note}, actor=actor) return {"source_id": source_id, "target_id": target_id, "mode": mode, "moved": moved} async def _mark_artifact(conn: AsyncConnection, src: dict[str, Any], dst: dict[str, Any], *, actor: str, note: str | None, payload: dict[str, Any] | None) -> dict[str, Any]: if dst["entity_type"] not in ("model",): raise ValueError("the canonical entity of an artifact must be a model") kind = (payload or {}).get("artifact_kind") or (src["attributes"] or {}).get("artifact_kind") or "conversion" await execute(conn, """update entities set entity_type = 'artifact', canonical_id = :t, artifact_kind = coalesce(artifact_kind, :k), identity_confidence = coalesce(cast(:ic as text), identity_confidence), updated_at = now() where id = :s""", t=dst["id"], k=kind, ic=(payload or {}).get("identity_confidence"), s=src["id"]) inserted = await upsert_relation(conn, src["id"], "artifact_of", dst["id"], {"artifact_kind": kind}, source_id=await registry_source_id(conn)) await record_decision(conn, src["id"], dst["id"], "variant_of", actor=actor, note=note, payload={"artifact_kind": kind, **(payload or {})}) await audit(conn, "entity.artifact_of", src["id"], {"canonical_id": dst["id"], "artifact_kind": kind, "note": note}, actor=actor) return {"source_id": src["id"], "target_id": dst["id"], "mode": "variant", "moved": {"relations": int(inserted)}} async def _mark_family_member(conn: AsyncConnection, src: dict[str, Any], dst: dict[str, Any], *, actor: str, note: str | None, payload: dict[str, Any] | None) -> dict[str, Any]: if dst["entity_type"] != "model_family": raise ValueError("family_member requires a model_family target") await execute(conn, "update entities set family_id = :t, updated_at = now() where id = :s and family_id is distinct from :t", t=dst["id"], s=src["id"]) inserted = await upsert_relation(conn, src["id"], "member_of_family", dst["id"], source_id=await registry_source_id(conn)) await record_decision(conn, src["id"], dst["id"], "family_member", actor=actor, note=note, payload=payload) await audit(conn, "entity.family_member", src["id"], {"family_id": dst["id"], "note": note}, actor=actor) return {"source_id": src["id"], "target_id": dst["id"], "mode": "family_member", "moved": {"relations": int(inserted)}} # ---------------------------------------------------------------------------------------------- benchmark result bookkeeping def _cfg_hash(config: dict[str, Any] | None) -> str: import hashlib return hashlib.sha1(json.dumps(config or {}, sort_keys=True, default=str).encode()).hexdigest()[:12] async def recompute_result_keys(conn: AsyncConnection, model_id: str, *, source_id: str | None = None) -> int: """Recompute `dedupe_key` (= model:benchmark:cfg_hash:metric) and `config_key` for the rows now attached to `model_id`. Collisions (the target already had the same result from the same source) close the older row. Returns the number of collisions.""" rows = await fetch_all(conn, """select id, benchmark_id, metric, config, dedupe_key, config_key, run_group, variant, observed_at, valid_to, is_current from benchmark_results where model_id = :m order by observed_at asc, id asc""", m=model_id) collisions = 0 seen: dict[str, dict[str, Any]] = {} for r in rows: base = f"{model_id}:{r['benchmark_id']}:{_cfg_hash(r['config'])}:{r['metric'] or ''}" ck = bench_ontology.config_key(r["config"], r["metric"]) rg = r["run_group"] or bench_ontology.run_group_from_config(r["config"]) variant = r["variant"] or bench_ontology.variant_from_config(r["config"]) if (rg, variant) != (r["run_group"], r["variant"]): await execute(conn, "update benchmark_results set run_group = :rg, variant = :v where id = :id", rg=rg, v=variant, id=r["id"]) if r["valid_to"] is not None: # historical rows keep a unique suffixed key key = base if r["dedupe_key"] == base and base not in seen else f"{base}:{r['id'][-10:]}" await execute(conn, "update benchmark_results set dedupe_key = :d, config_key = :ck where id = :id and (dedupe_key <> :d or config_key is distinct from :ck)", d=key, ck=ck, id=r["id"]) continue prev = seen.get(base) if prev is not None: # two live rows for the same key: keep the newer (rows are sorted by observed_at asc), close the older older, newer = prev, r await execute(conn, "update benchmark_results set valid_to = :o, is_current = false, dedupe_key = :d, config_key = :ck where id = :id", o=newer["observed_at"], d=f"{base}:{older['id'][-10:]}", ck=ck, id=older["id"]) collisions += 1 seen[base] = r await execute(conn, "update benchmark_results set dedupe_key = :d, config_key = :ck where id = :id and (dedupe_key <> :d or config_key is distinct from :ck)", d=base, ck=ck, id=r["id"]) return collisions async def enforce_current_results(conn: AsyncConnection, *, model_id: str | None = None, dry_run: bool = False) -> int: """One current row per (model, benchmark, metric, config_key): rows from an older run group than the latest observed one are closed (`is_current=false`, `valid_to` = the newer observation). Returns the number of rows (that would be) closed.""" scope = "and r.model_id = :m" if model_id else "" sql = f"""with latest as ( select distinct on (model_id, benchmark_id, coalesce(metric, ''), config_key) id, model_id, benchmark_id, coalesce(metric, '') as metric, config_key, coalesce(run_group, '') as run_group, observed_at from benchmark_results r where valid_to is null and is_current and config_key is not null {scope} order by model_id, benchmark_id, coalesce(metric, ''), config_key, observed_at desc, coalesce(run_group, '') desc, id desc) select r.id, l.observed_at as close_at from benchmark_results r join latest l on l.model_id = r.model_id and l.benchmark_id = r.benchmark_id and l.metric = coalesce(r.metric, '') and l.config_key = r.config_key where r.valid_to is null and r.is_current and r.id <> l.id and coalesce(r.run_group, '') <> l.run_group {scope}""" rows = await fetch_all(conn, sql, m=model_id) if model_id else await fetch_all(conn, sql) if dry_run or not rows: return len(rows) for r in rows: await execute(conn, "update benchmark_results set is_current = false, valid_to = :o where id = :id", o=r["close_at"], id=r["id"]) return len(rows) __all__ = ["MODES", "ORG_TYPES", "audit", "enforce_current_results", "kept_separate", "merge_entities", "recompute_result_keys", "record_decision", "registry_source_id", "upsert_relation"]