HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""Entity merging (curation): fold a duplicate `source` into `target`. Nothing is deleted — the source row stays with2`status='merged'` and `merged_into=target` so old slugs and ids keep resolving; every dependent row is re-pointed.34Modes5 merge (default) full merge; records a `merge` decision6 alias same as merge; records an `alias` decision (the source was another name of the target)7 variant the source is an *artifact* (quantisation / conversion / packaging) of the target: it keeps its own row, slug and8 facts, gets `entity_type='artifact'`, `canonical_id=target` and an `artifact_of` relation9 family_member the target is a `model_family`: the source gets `family_id=target` and a `member_of_family` relation1011Every call writes a `resolution_decisions` row (applied=true) and an `admin_audit_log` row; `keep_separate` decisions block merges.12"""13from __future__ import annotations1415import json16from typing import Any1718from sqlalchemy.ext.asyncio import AsyncConnection1920from aiatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb21from aiatlas.ids import new_id, normalize_alias22from aiatlas.ontology import benchmarks as bench_ontology2324ORG_TYPES = {"company", "organization", "lab", "university"}25MODES = ("merge", "alias", "variant", "family_member")26DERIVED_SOURCE_KEY = "ai-atlas.registry"272829async def registry_source_id(conn: AsyncConnection) -> str | None:30 row = await fetch_one(conn, "select id from sources where key = :k", k=DERIVED_SOURCE_KEY)31 return row["id"] if row else None323334async def record_decision(conn: AsyncConnection, a_id: str, b_id: str, decision: str, *, actor: str = "curation", note: str | None = None,35 payload: dict[str, Any] | None = None, applied: bool = True) -> None:36 await execute(conn, """insert into resolution_decisions (id, a_id, b_id, decision, actor, note, payload, applied)37 values (:id, :a, :b, :d, :actor, :note, cast(:p as jsonb), :applied)38 on conflict (a_id, b_id, decision) do update set applied = resolution_decisions.applied or excluded.applied,39 note = coalesce(excluded.note, resolution_decisions.note), payload = resolution_decisions.payload || excluded.payload""",40 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)414243async def audit(conn: AsyncConnection, action: str, target: str | None, payload: dict[str, Any] | None = None, *, actor: str = "curation") -> None:44 await execute(conn, "insert into admin_audit_log (actor, action, target, payload) values (:actor, :action, :target, cast(:p as jsonb))",45 actor=actor, action=action, target=target, p=jsonb(payload or {}))464748async def kept_separate(conn: AsyncConnection, a: str, b: str) -> bool:49 row = await fetch_one(conn, """select 1 from resolution_decisions where decision = 'keep_separate'50 and ((a_id = :a and b_id = :b) or (a_id = :b and b_id = :a)) limit 1""", a=a, b=b)51 return row is not None525354async def upsert_relation(conn: AsyncConnection, subject_id: str, predicate: str, object_id: str, attributes: dict[str, Any] | None = None, *,55 source_id: str | None = None, tier: int = 2, confidence: str = "high") -> bool:56 """Insert a live relation unless an identical live edge exists. Returns True when a row was inserted."""57 if subject_id == object_id:58 return False59 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",60 s=subject_id, p=predicate, o=object_id)61 if existing:62 if attributes:63 await execute(conn, "update relations set attributes = attributes || cast(:a as jsonb) where id = :id", a=jsonb(attributes), id=existing["id"])64 return False65 await execute(conn, """insert into relations (id, subject_id, predicate, object_id, attributes, source_id, tier, confidence, observed_at, valid_from)66 values (:id, :s, :p, :o, cast(:a as jsonb), :src, :tier, :conf, now(), now())""",67 id=new_id("relation"), s=subject_id, p=predicate, o=object_id, a=jsonb(attributes or {}), src=source_id, tier=tier, conf=confidence)68 return True697071async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str, *, mode: str = "merge", actor: str = "curation",72 note: str | None = None, payload: dict[str, Any] | None = None) -> dict[str, Any]:73 if mode not in MODES:74 raise ValueError(f"unknown merge mode {mode!r}")75 if source_id == target_id:76 raise ValueError("source and target are the same entity")77 src = await fetch_one(conn, "select id, entity_type, canonical_name, slug, merged_into, attributes, provenance from entities where id = :id", id=source_id)78 dst = await fetch_one(conn, "select id, entity_type, canonical_name, merged_into from entities where id = :id", id=target_id)79 if not src or not dst:80 raise LookupError("source or target entity not found")81 if dst["merged_into"]:82 raise ValueError("target is itself merged; merge into its survivor instead")83 if src["merged_into"]:84 raise ValueError("source is already merged")85 if mode in ("merge", "alias") and await kept_separate(conn, source_id, target_id):86 raise ValueError("a keep_separate decision exists for this pair; refusing to merge")87 if mode == "variant":88 return await _mark_artifact(conn, src, dst, actor=actor, note=note, payload=payload)89 if mode == "family_member":90 return await _mark_family_member(conn, src, dst, actor=actor, note=note, payload=payload)91 same_type = src["entity_type"] == dst["entity_type"]92 org_pair = src["entity_type"] in ORG_TYPES and dst["entity_type"] in ORG_TYPES93 model_pair = {src["entity_type"], dst["entity_type"]} <= {"model", "artifact"}94 if not (same_type or org_pair or model_pair):95 raise ValueError(f"cannot merge a {src['entity_type']} into a {dst['entity_type']}")9697 moved: dict[str, int] = {}9899 async def count_update(label: str, sql: str) -> None:100 n = await fetch_val(conn, f"with u as ({sql} returning 1) select count(*) from u", s=source_id, t=target_id)101 moved[label] = int(n or 0)102103 await count_update("aliases", "insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) "104 "select :t, alias, alias_norm, kind, snapshot_id from entity_aliases where entity_id = :s on conflict (entity_id, alias_norm) do nothing")105 await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:t, :a, :n, 'former_name') "106 "on conflict (entity_id, alias_norm) do update set kind = 'former_name'",107 t=target_id, a=src["canonical_name"], n=normalize_alias(src["canonical_name"]))108 await execute(conn, "delete from entity_aliases where entity_id = :s", s=source_id)109 await count_update("identifiers", "update entity_identifiers x set entity_id = :t where entity_id = :s "110 "and not exists (select 1 from entity_identifiers y where y.scheme = x.scheme and y.value = x.value and y.entity_id = :t)")111 await execute(conn, "delete from entity_identifiers where entity_id = :s", s=source_id)112 await count_update("claims", "update claims set entity_id = :t where entity_id = :s")113 # relations: drop those that would duplicate a live edge on the target, then re-point114 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 (115 select 1 from relations x where x.valid_to is null and x.predicate = r.predicate116 and x.subject_id = case when r.subject_id = :s then :t else r.subject_id end117 and x.object_id = case when r.object_id = :s then :t else r.object_id end)""", s=source_id, t=target_id)118 await count_update("relations", "update relations set subject_id = case when subject_id = :s then :t else subject_id end, "119 "object_id = case when object_id = :s then :t else object_id end where subject_id = :s or object_id = :s")120 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)121 await count_update("events", "update change_events set entity_id = :t where entity_id = :s")122 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 "123 "and q.model_id = :t and q.provider_id = prices.provider_id and coalesce(q.provider_model_id,'') = coalesce(prices.provider_model_id,''))",124 s=source_id, t=target_id)125 await count_update("prices", "update prices set model_id = case when model_id = :s then :t else model_id end, "126 "provider_id = case when provider_id = :s then :t else provider_id end where model_id = :s or provider_id = :s")127 await count_update("results", "update benchmark_results set model_id = case when model_id = :s then :t else model_id end, "128 "benchmark_id = case when benchmark_id = :s then :t else benchmark_id end where model_id = :s or benchmark_id = :s")129 moved["result_dedupe_collisions"] = await recompute_result_keys(conn, target_id, source_id=source_id)130 moved["results_closed"] = await enforce_current_results(conn, model_id=target_id)131 await count_update("documents", "update documents set entity_id = :t where entity_id = :s")132 await count_update("children", "update entities set organization_id = :t where organization_id = :s")133 await count_update("family_members", "update entities set family_id = :t where family_id = :s")134 await count_update("artifacts", "update entities set canonical_id = :t where canonical_id = :s")135 await count_update("llm_jobs", "update llm_jobs set entity_id = :t where entity_id = :s")136 await count_update("sources", "update sources set organization_id = :t where organization_id = :s")137 await count_update("domains", "update domains set organization_id = :t where organization_id = :s")138 await count_update("review_items", "update review_queue set entity_ids = array_replace(entity_ids, :s, :t) where :s = any(entity_ids)")139 # embeddings (table only exists where pgvector is installed): the target keeps its own vector; a source-only vector moves over140 if await fetch_val(conn, "select to_regclass('entity_embeddings') is not null"):141 has_target_vec = await fetch_one(conn, "select 1 from entity_embeddings where entity_id = :t", t=target_id)142 if has_target_vec:143 await count_update("embeddings", "delete from entity_embeddings where entity_id = :s")144 else:145 await count_update("embeddings", "update entity_embeddings set entity_id = :t where entity_id = :s")146 # attributes the target lacks are inherited (with their provenance); target values always win147 await execute(conn, """update entities t set attributes = coalesce(s.attributes, '{}'::jsonb) || t.attributes,148 provenance = coalesce(s.provenance, '{}'::jsonb) || t.provenance, last_seen_at = greatest(t.last_seen_at, s.last_seen_at),149 first_seen_at = least(t.first_seen_at, s.first_seen_at), updated_at = now()150 from entities s where t.id = :t and s.id = :s""", s=source_id, t=target_id)151 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)152 await execute(conn, "update entities set merged_into = :t where merged_into = :s", s=source_id, t=target_id)153 await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, connector_name, dedupe_key, meta, is_backfill)154 values (:id, :t, 'ENTITY_MERGED', 'source', :sum, 0, 'curation', :dk, cast(:m as jsonb), true)155 on conflict (dedupe_key) do nothing""",156 id=new_id("change_event"), t=target_id, sum=f"Merged duplicate '{src['canonical_name']}' ({src['slug']})", dk=f"merge:{source_id}:{target_id}",157 m=json.dumps({"source_id": source_id, "source_slug": src["slug"], "mode": mode}))158 decision = "alias" if mode == "alias" else "merge"159 await record_decision(conn, source_id, target_id, decision, actor=actor, note=note, payload={"moved": moved, **(payload or {})})160 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)161 return {"source_id": source_id, "target_id": target_id, "mode": mode, "moved": moved}162163164async 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]:165 if dst["entity_type"] not in ("model",):166 raise ValueError("the canonical entity of an artifact must be a model")167 kind = (payload or {}).get("artifact_kind") or (src["attributes"] or {}).get("artifact_kind") or "conversion"168 await execute(conn, """update entities set entity_type = 'artifact', canonical_id = :t, artifact_kind = coalesce(artifact_kind, :k),169 identity_confidence = coalesce(cast(:ic as text), identity_confidence), updated_at = now() where id = :s""",170 t=dst["id"], k=kind, ic=(payload or {}).get("identity_confidence"), s=src["id"])171 inserted = await upsert_relation(conn, src["id"], "artifact_of", dst["id"], {"artifact_kind": kind}, source_id=await registry_source_id(conn))172 await record_decision(conn, src["id"], dst["id"], "variant_of", actor=actor, note=note, payload={"artifact_kind": kind, **(payload or {})})173 await audit(conn, "entity.artifact_of", src["id"], {"canonical_id": dst["id"], "artifact_kind": kind, "note": note}, actor=actor)174 return {"source_id": src["id"], "target_id": dst["id"], "mode": "variant", "moved": {"relations": int(inserted)}}175176177async 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]:178 if dst["entity_type"] != "model_family":179 raise ValueError("family_member requires a model_family target")180 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"])181 inserted = await upsert_relation(conn, src["id"], "member_of_family", dst["id"], source_id=await registry_source_id(conn))182 await record_decision(conn, src["id"], dst["id"], "family_member", actor=actor, note=note, payload=payload)183 await audit(conn, "entity.family_member", src["id"], {"family_id": dst["id"], "note": note}, actor=actor)184 return {"source_id": src["id"], "target_id": dst["id"], "mode": "family_member", "moved": {"relations": int(inserted)}}185186187# ---------------------------------------------------------------------------------------------- benchmark result bookkeeping188def _cfg_hash(config: dict[str, Any] | None) -> str:189 import hashlib190191 return hashlib.sha1(json.dumps(config or {}, sort_keys=True, default=str).encode()).hexdigest()[:12]192193194async def recompute_result_keys(conn: AsyncConnection, model_id: str, *, source_id: str | None = None) -> int:195 """Recompute `dedupe_key` (= model:benchmark:cfg_hash:metric) and `config_key` for the rows now attached to `model_id`.196 Collisions (the target already had the same result from the same source) close the older row. Returns the number of collisions."""197 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_results198 where model_id = :m order by observed_at asc, id asc""", m=model_id)199 collisions = 0200 seen: dict[str, dict[str, Any]] = {}201 for r in rows:202 base = f"{model_id}:{r['benchmark_id']}:{_cfg_hash(r['config'])}:{r['metric'] or ''}"203 ck = bench_ontology.config_key(r["config"], r["metric"])204 rg = r["run_group"] or bench_ontology.run_group_from_config(r["config"])205 variant = r["variant"] or bench_ontology.variant_from_config(r["config"])206 if (rg, variant) != (r["run_group"], r["variant"]):207 await execute(conn, "update benchmark_results set run_group = :rg, variant = :v where id = :id", rg=rg, v=variant, id=r["id"])208 if r["valid_to"] is not None:209 # historical rows keep a unique suffixed key210 key = base if r["dedupe_key"] == base and base not in seen else f"{base}:{r['id'][-10:]}"211 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)",212 d=key, ck=ck, id=r["id"])213 continue214 prev = seen.get(base)215 if prev is not None:216 # two live rows for the same key: keep the newer (rows are sorted by observed_at asc), close the older217 older, newer = prev, r218 await execute(conn, "update benchmark_results set valid_to = :o, is_current = false, dedupe_key = :d, config_key = :ck where id = :id",219 o=newer["observed_at"], d=f"{base}:{older['id'][-10:]}", ck=ck, id=older["id"])220 collisions += 1221 seen[base] = r222 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)",223 d=base, ck=ck, id=r["id"])224 return collisions225226227async def enforce_current_results(conn: AsyncConnection, *, model_id: str | None = None, dry_run: bool = False) -> int:228 """One current row per (model, benchmark, metric, config_key): rows from an older run group than the latest observed one are closed229 (`is_current=false`, `valid_to` = the newer observation). Returns the number of rows (that would be) closed."""230 scope = "and r.model_id = :m" if model_id else ""231 sql = f"""with latest as (232 select distinct on (model_id, benchmark_id, coalesce(metric, ''), config_key) id, model_id, benchmark_id, coalesce(metric, '') as metric, config_key,233 coalesce(run_group, '') as run_group, observed_at234 from benchmark_results r where valid_to is null and is_current and config_key is not null {scope}235 order by model_id, benchmark_id, coalesce(metric, ''), config_key, observed_at desc, coalesce(run_group, '') desc, id desc)236 select r.id, l.observed_at as close_at from benchmark_results r join latest l237 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_key238 where r.valid_to is null and r.is_current and r.id <> l.id and coalesce(r.run_group, '') <> l.run_group {scope}"""239 rows = await fetch_all(conn, sql, m=model_id) if model_id else await fetch_all(conn, sql)240 if dry_run or not rows:241 return len(rows)242 for r in rows:243 await execute(conn, "update benchmark_results set is_current = false, valid_to = :o where id = :id", o=r["close_at"], id=r["id"])244 return len(rows)245246247__all__ = ["MODES", "ORG_TYPES", "audit", "enforce_current_results", "kept_separate", "merge_entities", "recompute_result_keys", "record_decision",248 "registry_source_id", "upsert_relation"]249