| |
1 | +"""Canonicalization engine — `aia canonicalize [--apply] [--step …]`. |
| |
2 | + |
| |
3 | +Turns the flat "everything is a model" corpus into the canonical hierarchy (model_family → model → artifact, effort variants folded |
| |
4 | +into result configurations), normalises taxonomies, links licences, classifies events (backfill vs live), enforces benchmark result |
| |
5 | +comparability and flags anomalies. Rules are documented in docs/CANONICALIZATION.md. |
| |
6 | + |
| |
7 | +Invariants |
| |
8 | + * dry-run by default: every step computes a plan from reads only and reports counts + examples; `--apply` executes it |
| |
9 | + * nothing is ever deleted; raw snapshots are never touched; every write is idempotent (a second `--apply` is a no-op) |
| |
10 | + * merges go through `services.merge.merge_entities` (persisted decision + audit log); relations through `upsert_relation` |
| |
11 | + * derived claims are written by a tier-2 `derived` FactWriter (source `ai-atlas.registry`) and never supersede tier-1 statements |
| |
12 | +""" |
| |
13 | +from __future__ import annotations |
| |
14 | + |
| |
15 | +import json |
| |
16 | +import logging |
| |
17 | +from collections import defaultdict |
| |
18 | +from dataclasses import dataclass, field |
| |
19 | +from datetime import UTC, datetime, timedelta |
| |
20 | +from typing import Any |
| |
21 | + |
| |
22 | +from sqlalchemy.ext.asyncio import AsyncConnection |
| |
23 | + |
| |
24 | +from aiatlas.db import execute, fetch_all, fetch_one, jsonb, transaction |
| |
25 | +from aiatlas.ids import new_id, normalize_alias, slugify |
| |
26 | +from aiatlas.ontology import benchmarks as bench_ontology |
| |
27 | +from aiatlas.ontology.anomalies import Anomaly, check_hardware, check_model, check_price, check_result |
| |
28 | +from aiatlas.ontology.licenses import LICENSES, normalize_license |
| |
29 | +from aiatlas.ontology.models import ( |
| |
30 | + CONVERTER_ORGS, |
| |
31 | + analyze_model_name, |
| |
32 | + base_name, |
| |
33 | + effort_config, |
| |
34 | + family_hint, |
| |
35 | + family_release_hint, |
| |
36 | + official_orgs, |
| |
37 | + variant_key, |
| |
38 | +) |
| |
39 | +from aiatlas.ontology.openness import derive_openness, normalize_openness, openness_dimensions |
| |
40 | +from aiatlas.ontology.taxonomy import ORG_TYPE_DEFAULT_KIND, TAXONOMY_PROPERTIES, normalize_property |
| |
41 | +from aiatlas.sdk.facts import EntityRef, Facts, facts_from_json |
| |
42 | +from aiatlas.sdk.resolution import EVALUATOR_SCHEMES, Resolver |
| |
43 | +from aiatlas.sdk.writer import FactWriter, _same |
| |
44 | +from aiatlas.services.anomalies import record |
| |
45 | +from aiatlas.services.events import BACKFILL_LAG_DAYS, group_key_for |
| |
46 | +from aiatlas.services.merge import ( |
| |
47 | + ORG_TYPES, |
| |
48 | + audit, |
| |
49 | + enforce_current_results, |
| |
50 | + kept_separate, |
| |
51 | + merge_entities, |
| |
52 | + record_decision, |
| |
53 | + registry_source_id, |
| |
54 | + upsert_relation, |
| |
55 | +) |
| |
56 | + |
| |
57 | +log = logging.getLogger(__name__) |
| |
58 | + |
| |
59 | +CANON_VERSION = "2026.09" |
| |
60 | +STRUCTURAL_CONNECTORS = {"curation", "canonicalize"} |
| |
61 | +STEPS = ("duplicates", "variants", "artifacts", "families", "licenses", "taxonomy", "results", "events", "anomalies") |
| |
62 | +DEDUPE_TYPES = ("model", "provider", "benchmark", "dataset", "framework", "library", "hardware") |
| |
63 | +LICENSED_TYPES = ("model", "artifact", "dataset", "framework", "library", "repository") |
| |
64 | +QUANT_ATTR_FORMATS = {"gguf", "awq", "gptq", "exl2", "exl3", "int4", "int8", "fp8", "nvfp4", "mxfp4", "fp4", "bnb", "quantized"} |
| |
65 | + |
| |
66 | + |
| |
67 | +@dataclass |
| |
68 | +class StepReport: |
| |
69 | + name: str |
| |
70 | + counts: dict[str, int] = field(default_factory=dict) |
| |
71 | + examples: list[str] = field(default_factory=list) |
| |
72 | + notes: list[str] = field(default_factory=list) |
| |
73 | + |
| |
74 | + def bump(self, key: str, n: int = 1) -> None: |
| |
75 | + self.counts[key] = self.counts.get(key, 0) + n |
| |
76 | + |
| |
77 | + def example(self, text: str, *, limit: int = 12) -> None: |
| |
78 | + if len(self.examples) < limit: |
| |
79 | + self.examples.append(text) |
| |
80 | + |
| |
81 | + @property |
| |
82 | + def changes(self) -> int: |
| |
83 | + return sum(v for k, v in self.counts.items() if not k.startswith("_")) |
| |
84 | + |
| |
85 | + |
| |
86 | +@dataclass |
| |
87 | +class Report: |
| |
88 | + apply: bool |
| |
89 | + steps: list[StepReport] = field(default_factory=list) |
| |
90 | + started_at: datetime = field(default_factory=lambda: datetime.now(UTC)) |
| |
91 | + |
| |
92 | + @property |
| |
93 | + def changes(self) -> int: |
| |
94 | + return sum(s.changes for s in self.steps) |
| |
95 | + |
| |
96 | + def render(self) -> str: |
| |
97 | + mode = "APPLY" if self.apply else "DRY-RUN" |
| |
98 | + lines = [f"aia canonicalize — {mode} — {self.started_at:%Y-%m-%d %H:%M:%S} UTC — {self.changes} change(s)"] |
| |
99 | + for s in self.steps: |
| |
100 | + lines.append(f"\n[{s.name}] {s.changes} change(s)") |
| |
101 | + for k, v in sorted(s.counts.items()): |
| |
102 | + lines.append(f" {k:<40} {v}") |
| |
103 | + for n in s.notes: |
| |
104 | + lines.append(f" · {n}") |
| |
105 | + for e in s.examples: |
| |
106 | + lines.append(f" - {e}") |
| |
107 | + return "\n".join(lines) |
| |
108 | + |
| |
109 | + def as_dict(self) -> dict[str, Any]: |
| |
110 | + return {"apply": self.apply, "started_at": self.started_at.isoformat(), "changes": self.changes, |
| |
111 | + "steps": [{"name": s.name, "counts": s.counts, "examples": s.examples, "notes": s.notes} for s in self.steps]} |
| |
112 | + |
| |
113 | + |
| |
114 | +async def canonicalize(*, apply: bool = False, steps: list[str] | None = None, scope: set[str] | None = None) -> Report: |
| |
115 | + """Run the steps in canonical order. `scope` (entity ids) restricts the rows a step *acts on* — lookups (bases, families, licences) |
| |
116 | + still see the whole corpus; used by tests and targeted re-runs.""" |
| |
117 | + wanted = [s for s in STEPS if not steps or s in steps] |
| |
118 | + unknown = set(steps or []) - set(STEPS) |
| |
119 | + if unknown: |
| |
120 | + raise ValueError(f"unknown step(s): {sorted(unknown)}; known: {STEPS}") |
| |
121 | + report = Report(apply=apply) |
| |
122 | + for name in wanted: |
| |
123 | + rep = StepReport(name=name) |
| |
124 | + fn = _STEP_FUNCTIONS[name] |
| |
125 | + async with transaction() as conn: |
| |
126 | + await fn(conn, rep, apply, scope=scope) |
| |
127 | + if apply and rep.changes: |
| |
128 | + await audit(conn, f"canonicalize.{name}", None, {"counts": rep.counts, "version": CANON_VERSION}, actor="canonicalize") |
| |
129 | + report.steps.append(rep) |
| |
130 | + log.info("canonicalize step done", extra={"step": name, "apply": apply, **{k: v for k, v in rep.counts.items()}}) |
| |
131 | + return report |
| |
132 | + |
| |
133 | + |
| |
134 | +# ---------------------------------------------------------------------------------------------- shared helpers |
| |
135 | +def _derived_writer(conn: AsyncConnection, source_id: str | None) -> FactWriter: |
| |
136 | + return FactWriter(conn, source_id=source_id, snapshot_id=None, source_url=None, tier=2, connector_name="canonicalize", extractor="derived", |
| |
137 | + extractor_version=CANON_VERSION, run_id=f"canon_{datetime.now(UTC):%Y%m%d}") |
| |
138 | + |
| |
139 | + |
| |
140 | +async def _review(conn: AsyncConnection, kind: str, entity_ids: list[str], reason: str, payload: dict[str, Any], *, apply: bool = True) -> bool: |
| |
141 | + """Queue a review item once (dedupe key). Returns True when the item is new (or would be, in dry-run) so reports stay idempotent.""" |
| |
142 | + dedupe = f"{kind}:{':'.join(entity_ids)}:{normalize_alias(reason)[:80]}" |
| |
143 | + if await fetch_one(conn, "select 1 from review_queue where dedupe_key = :d", d=dedupe): |
| |
144 | + return False |
| |
145 | + if not apply: |
| |
146 | + return True |
| |
147 | + row = await fetch_one(conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key) values (:id, :k, :ids, cast(:p as jsonb), :r, :d) |
| |
148 | + on conflict (dedupe_key) do nothing returning id""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe) |
| |
149 | + return row is not None |
| |
150 | + |
| |
151 | + |
| |
152 | +async def _set_attributes(conn: AsyncConnection, entity_id: str, attrs: dict[str, Any], *, source_id: str | None) -> None: |
| |
153 | + """Direct attribute write for structural/derived facts (family label, license_key, org_kind default) with `derived` provenance.""" |
| |
154 | + prov = {k: {"source_id": source_id, "tier": 2, "confidence": "high", "extractor": "derived", "observed_at": datetime.now(UTC).isoformat(timespec="seconds")} for k in attrs} |
| |
155 | + await execute(conn, "update entities set attributes = attributes || cast(:a as jsonb), provenance = provenance || cast(:p as jsonb), updated_at = now() where id = :id", |
| |
156 | + a=jsonb(attrs), p=jsonb(prov), id=entity_id) |
| |
157 | + |
| |
158 | + |
| |
159 | +async def _claim_counts(conn: AsyncConnection, ids: list[str]) -> dict[str, int]: |
| |
160 | + if not ids: |
| |
161 | + return {} |
| |
162 | + rows = await fetch_all(conn, "select entity_id, count(*) as n from claims where status = 'current' and entity_id = any(cast(:ids as text[])) group by 1", ids=ids) |
| |
163 | + return {r["entity_id"]: int(r["n"]) for r in rows} |
| |
164 | + |
| |
165 | + |
| |
166 | +async def _identifiers(conn: AsyncConnection, ids: list[str]) -> dict[str, dict[str, set[str]]]: |
| |
167 | + out: dict[str, dict[str, set[str]]] = defaultdict(lambda: defaultdict(set)) |
| |
168 | + if not ids: |
| |
169 | + return out |
| |
170 | + rows = await fetch_all(conn, "select entity_id, scheme, value from entity_identifiers where entity_id = any(cast(:ids as text[]))", ids=ids) |
| |
171 | + for r in rows: |
| |
172 | + out[r["entity_id"]][r["scheme"]].add(r["value"]) |
| |
173 | + return out |
| |
174 | + |
| |
175 | + |
| |
176 | +def _identifiers_conflict(a: dict[str, set[str]], b: dict[str, set[str]]) -> bool: |
| |
177 | + return any(a[s] and b[s] and a[s] != b[s] for s in set(a) & set(b)) |
| |
178 | + |
| |
179 | + |
| |
180 | +def _org_lookup_keys(row: dict[str, Any]) -> set[str]: |
| |
181 | + attrs = row.get("attributes") or {} |
| |
182 | + keys = {row["slug"], normalize_alias(row["canonical_name"])} |
| |
183 | + for k in ("hf_org", "github_org"): |
| |
184 | + v = attrs.get(k) |
| |
185 | + if isinstance(v, str) and v.strip(): |
| |
186 | + keys.add(v.strip().lower()) |
| |
187 | + return {k for k in keys if k} |
| |
188 | + |
| |
189 | + |
| |
190 | +# ---------------------------------------------------------------------------------------------- step: duplicates |
| |
191 | +def _in_scope(scope: set[str] | None, entity_id: str) -> bool: |
| |
192 | + return scope is None or entity_id in scope |
| |
193 | + |
| |
194 | + |
| |
195 | +async def step_duplicates(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
196 | + source_id = await registry_source_id(conn) |
| |
197 | + rows = await fetch_all(conn, """select id, entity_type, canonical_name, slug, organization_id, attributes, first_seen_at from entities |
| |
198 | + where merged_into is null and entity_type = any(cast(:types as text[]))""", types=list(DEDUPE_TYPES) + sorted(ORG_TYPES)) |
| |
199 | + if scope is not None: |
| |
200 | + rows = [r for r in rows if r["id"] in scope] |
| |
201 | + # group: exact normalised name within a type; organisations across the org group (+ shared hf_org/github_org) |
| |
202 | + groups: dict[str, list[dict[str, Any]]] = defaultdict(list) |
| |
203 | + for r in rows: |
| |
204 | + norm = normalize_alias(r["canonical_name"]) |
| |
205 | + if r["entity_type"] in ORG_TYPES: |
| |
206 | + groups[f"org:{norm}"].append(r) |
| |
207 | + for k in ("hf_org", "github_org"): |
| |
208 | + v = (r["attributes"] or {}).get(k) |
| |
209 | + if isinstance(v, str) and v.strip() and normalize_alias(v) != norm: |
| |
210 | + groups[f"org:{normalize_alias(v)}"].append(r) |
| |
211 | + else: |
| |
212 | + groups[f"{r['entity_type']}:{norm}"].append(r) |
| |
213 | + # union overlapping org groups |
| |
214 | + parent: dict[str, str] = {} |
| |
215 | + |
| |
216 | + def find(x: str) -> str: |
| |
217 | + while parent.setdefault(x, x) != x: |
| |
218 | + x = parent[x] |
| |
219 | + return x |
| |
220 | + |
| |
221 | + for members in groups.values(): |
| |
222 | + ids = [m["id"] for m in members] |
| |
223 | + for other in ids[1:]: |
| |
224 | + parent[find(other)] = find(ids[0]) |
| |
225 | + clusters: dict[str, list[dict[str, Any]]] = defaultdict(list) |
| |
226 | + seen: set[str] = set() |
| |
227 | + for members in groups.values(): |
| |
228 | + for m in members: |
| |
229 | + if m["id"] not in seen: |
| |
230 | + seen.add(m["id"]) |
| |
231 | + clusters[find(m["id"])].append(m) |
| |
232 | + dup_clusters = [c for c in clusters.values() if len(c) > 1] |
| |
233 | + all_ids = [m["id"] for c in dup_clusters for m in c] |
| |
234 | + claims = await _claim_counts(conn, all_ids) |
| |
235 | + idents = await _identifiers(conn, all_ids) |
| |
236 | + type_rank = {"company": 0, "lab": 0, "university": 0, "organization": 1} # curated org types beat the generic hub "organization" |
| |
237 | + for cluster in dup_clusters: |
| |
238 | + cluster.sort(key=lambda m: (type_rank.get(m["entity_type"], 0), -claims.get(m["id"], 0), m["first_seen_at"])) |
| |
239 | + survivor = cluster[0] |
| |
240 | + for other in cluster[1:]: |
| |
241 | + pair = sorted([survivor["id"], other["id"]]) |
| |
242 | + if await kept_separate(conn, survivor["id"], other["id"]): |
| |
243 | + rep.bump("_kept_separate") |
| |
244 | + continue |
| |
245 | + if _identifiers_conflict(idents[survivor["id"]], idents[other["id"]]): |
| |
246 | + if await _review(conn, "merge_candidate", pair, f"'{other['canonical_name']}' duplicates '{survivor['canonical_name']}' but identifiers conflict", |
| |
247 | + {"slugs": [other["slug"], survivor["slug"]], "step": "duplicates"}, apply=apply): |
| |
248 | + rep.bump("review_conflicting_identifiers") |
| |
249 | + rep.example(f"review: {other['canonical_name']} ({other['slug']}) vs {survivor['canonical_name']} ({survivor['slug']}) — identifiers differ") |
| |
250 | + continue |
| |
251 | + if survivor["organization_id"] and other["organization_id"] and survivor["organization_id"] != other["organization_id"] and other["entity_type"] not in ORG_TYPES: |
| |
252 | + if await _review(conn, "merge_candidate", pair, f"'{other['canonical_name']}' and '{survivor['canonical_name']}' share a name but have different organisations", |
| |
253 | + {"slugs": [other["slug"], survivor["slug"]], "step": "duplicates"}, apply=apply): |
| |
254 | + rep.bump("review_different_organizations") |
| |
255 | + rep.example(f"review: {other['slug']} and {survivor['slug']} share a name but belong to different organisations") |
| |
256 | + continue |
| |
257 | + rep.bump("merged") |
| |
258 | + rep.example(f"merge {other['entity_type']} {other['slug']} → {survivor['slug']}") |
| |
259 | + if apply: |
| |
260 | + await merge_entities(conn, other["id"], survivor["id"], mode="merge", actor="canonicalize", note="exact normalised-name duplicate", |
| |
261 | + payload={"step": "duplicates"}) |
| |
262 | + # junk organisations (single-letter names) → review only |
| |
263 | + for r in rows: |
| |
264 | + if r["entity_type"] in ORG_TYPES and len(r["canonical_name"].strip()) <= 1: |
| |
265 | + if await _review(conn, "junk_entity", [r["id"]], f"organisation '{r['canonical_name']}' ({r['slug']}) looks like extraction noise", {"slug": r["slug"]}, apply=apply): |
| |
266 | + rep.bump("review_junk_organization") |
| |
267 | + rep.example(f"review junk organisation {r['slug']!r}") |
| |
268 | + # org_kind defaults / normalisation |
| |
269 | + for r in rows: |
| |
270 | + if r["entity_type"] not in ORG_TYPES: |
| |
271 | + continue |
| |
272 | + attrs = r["attributes"] or {} |
| |
273 | + current = attrs.get("org_kind") |
| |
274 | + default = ORG_TYPE_DEFAULT_KIND.get(r["entity_type"]) |
| |
275 | + canon, _raw, _m = normalize_property(r["entity_type"], "org_kind", current) if current else (None, None, []) |
| |
276 | + target = canon if isinstance(canon, str) and canon in ("company", "lab", "university", "nonprofit", "government", "community", "consortium", "individual") else default |
| |
277 | + if target and current != target: |
| |
278 | + rep.bump("org_kind_set") |
| |
279 | + if apply: |
| |
280 | + await _set_attributes(conn, r["id"], {"org_kind": target}, source_id=source_id) |
| |
281 | + |
| |
282 | + |
| |
283 | +# ---------------------------------------------------------------------------------------------- step: variants |
| |
284 | +async def _models(conn: AsyncConnection, *, types: tuple[str, ...] = ("model",)) -> list[dict[str, Any]]: |
| |
285 | + return await fetch_all(conn, """select e.id, e.entity_type, e.canonical_name, e.slug, e.organization_id, e.attributes, e.first_seen_at, e.family_id, e.canonical_id, |
| |
286 | + e.identity_confidence, o.slug as org_slug, o.canonical_name as org_name, o.attributes->>'hf_org' as org_hf |
| |
287 | + from entities e left join entities o on o.id = e.organization_id |
| |
288 | + where e.merged_into is null and e.entity_type = any(cast(:t as text[])) order by e.first_seen_at, e.id""", t=list(types)) |
| |
289 | + |
| |
290 | + |
| |
291 | +def _variant_index(models: list[dict[str, Any]]) -> tuple[dict[str, str], dict[str, list[dict[str, Any]]]]: |
| |
292 | + """variant_key → canonical model id (unique keys only) and the full multi-map.""" |
| |
293 | + multi: dict[str, list[dict[str, Any]]] = defaultdict(list) |
| |
294 | + for m in models: |
| |
295 | + a = analyze_model_name(m["canonical_name"]) |
| |
296 | + if a.is_effort_variant or a.is_artifact: |
| |
297 | + continue |
| |
298 | + multi[variant_key(m["canonical_name"])].append(m) |
| |
299 | + return {k: v[0]["id"] for k, v in multi.items() if len(v) == 1}, multi |
| |
300 | + |
| |
301 | + |
| |
302 | +async def step_variants(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
303 | + models = await _models(conn) |
| |
304 | + index, _ = _variant_index(models) |
| |
305 | + resolver = Resolver(conn, source_tier=2, variant_index=index) |
| |
306 | + candidates = [m for m in models if _in_scope(scope, m["id"]) and analyze_model_name(m["canonical_name"]).is_effort_variant] |
| |
307 | + idents = await _identifiers(conn, [m["id"] for m in candidates]) |
| |
308 | + official = {r["entity_id"] for r in await fetch_all(conn, """select distinct entity_id from claims where tier = 1 and status = 'current' |
| |
309 | + and entity_id = any(cast(:ids as text[]))""", ids=[m["id"] for m in candidates])} if candidates else set() |
| |
310 | + for m in candidates: |
| |
311 | + a = analyze_model_name(m["canonical_name"]) |
| |
312 | + rep.bump("_effort_variants_seen") |
| |
313 | + schemes = set(idents.get(m["id"], {})) |
| |
314 | + if not schemes <= EVALUATOR_SCHEMES or (m["attributes"] or {}).get("hf_repo") or m["id"] in official: |
| |
315 | + # known to an official source, a hub or a provider → a real model that happens to end in "thinking"/"high"; never folded |
| |
316 | + rep.bump("_kept_real_model") |
| |
317 | + continue |
| |
318 | + ref = EntityRef(entity_type="model", name=m["canonical_name"], id=m["id"]) |
| |
319 | + folded = await resolver.resolve_variant(ref, org_id=m["organization_id"]) |
| |
320 | + if folded is None: |
| |
321 | + hint = base_name(m["canonical_name"]) |
| |
322 | + attrs = m["attributes"] or {} |
| |
323 | + if attrs.get("evaluation_variant_of_hint") == hint and m["identity_confidence"] == "medium": |
| |
324 | + continue |
| |
325 | + rep.bump("unresolved_flagged") |
| |
326 | + rep.example(f"unresolved variant {m['slug']} (base '{hint}' not found) → identity medium + review") |
| |
327 | + if apply: |
| |
328 | + await _set_attributes(conn, m["id"], {"evaluation_variant_of_hint": hint}, source_id=await registry_source_id(conn)) |
| |
329 | + await execute(conn, "update entities set identity_confidence = 'medium' where id = :id", id=m["id"]) |
| |
330 | + await _review(conn, "variant_candidate", [m["id"]], f"'{m['canonical_name']}' looks like an evaluation-effort variant of '{hint}' but no such model exists", |
| |
331 | + {"slug": m["slug"], "base": hint, "effort": a.effort}, apply=apply) |
| |
332 | + continue |
| |
333 | + cid, effort = folded |
| |
334 | + rep.bump("folded") |
| |
335 | + target = next((x for x in models if x["id"] == cid), None) |
| |
336 | + rep.example(f"fold {m['slug']} → {target['slug'] if target else cid} {effort}") |
| |
337 | + if apply: |
| |
338 | + rows = await fetch_all(conn, "select id, config from benchmark_results where model_id = :m", m=m["id"]) |
| |
339 | + for r in rows: |
| |
340 | + cfg = effort_config(m["canonical_name"], r["config"]) |
| |
341 | + if not _same(cfg, r["config"]): |
| |
342 | + await execute(conn, "update benchmark_results set config = cast(:c as jsonb) where id = :id", c=jsonb(cfg), id=r["id"]) |
| |
343 | + rep.bump("results_reconfigured", len(rows)) |
| |
344 | + await merge_entities(conn, m["id"], cid, mode="merge", actor="canonicalize", note="evaluation-effort variant folded into its canonical model", |
| |
345 | + payload={"step": "variants", "effort": effort, "variant_slug": m["slug"]}) |
| |
346 | + |
| |
347 | + |
| |
348 | +# ---------------------------------------------------------------------------------------------- step: artifacts |
| |
349 | +async def step_artifacts(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
350 | + source_id = await registry_source_id(conn) |
| |
351 | + models = await _models(conn) |
| |
352 | + _, multi = _variant_index(models) |
| |
353 | + by_id = {m["id"]: m for m in models} |
| |
354 | + # only `quantized_from` names the packaged model; `derived_from`/`fine_tuned_from` point at a *base* model (a fine-tune is a new model) |
| |
355 | + rel_rows = await fetch_all(conn, """select r.subject_id, r.object_id from relations r join entities o on o.id = r.object_id |
| |
356 | + where r.predicate = 'quantized_from' and r.valid_to is null and o.merged_into is null and o.entity_type = 'model'""") |
| |
357 | + base_of: dict[str, str] = {} |
| |
358 | + for r in rel_rows: |
| |
359 | + base_of.setdefault(r["subject_id"], r["object_id"]) |
| |
360 | + official_idents = await _identifiers(conn, [m["id"] for m in models if (m["attributes"] or {}).get("hf_repo")]) |
| |
361 | + for m in models: |
| |
362 | + if not _in_scope(scope, m["id"]): |
| |
363 | + continue |
| |
364 | + attrs = m["attributes"] or {} |
| |
365 | + hf_repo = attrs.get("hf_repo") if isinstance(attrs.get("hf_repo"), str) else None |
| |
366 | + probe = hf_repo or m["canonical_name"] |
| |
367 | + a = analyze_model_name(probe) |
| |
368 | + a_name = analyze_model_name(m["canonical_name"]) if hf_repo else a |
| |
369 | + repo_org = (a.repo_org or "").lower() |
| |
370 | + # repo attributes (is_quantized / quant_format) describe a hub repository: without hf_repo they come from a provider endpoint |
| |
371 | + # (OpenRouter lists the serving precision) and say nothing about the model's identity |
| |
372 | + quant_attr = str(attrs.get("quant_format") or "").lower() if hf_repo else "" |
| |
373 | + attr_quantized = attrs.get("is_quantized") is True and bool(hf_repo) |
| |
374 | + has_token = bool(a.quant_formats or a.precision or a_name.quant_formats or a_name.precision) |
| |
375 | + flagged = has_token or a.is_artifact or a_name.is_artifact or attr_quantized or bool(quant_attr) or repo_org in CONVERTER_ORGS |
| |
376 | + if not flagged: |
| |
377 | + continue |
| |
378 | + # "official" = the repo belongs to the model's developer: the family's known publisher, or the entity's own organisation when that |
| |
379 | + # organisation is not a redistributor (hub-derived entities are attached to the redistributor org, e.g. `bartowski`) |
| |
380 | + own_org_keys = {k for k in ((m["org_slug"] or "").lower(), (m["org_hf"] or "").lower(), normalize_alias(m["org_name"] or "")) if k} |
| |
381 | + official = bool(repo_org) and repo_org not in CONVERTER_ORGS and (repo_org in set(official_orgs(probe)) or repo_org in own_org_keys) |
| |
382 | + vk_candidates = [c for c in multi.get(variant_key(probe), []) if c["id"] != m["id"]] |
| |
383 | + cid: str | None = None |
| |
384 | + if official: |
| |
385 | + # the developer's own repo is the model (NVIDIA-Nemotron-3-Ultra-…-BF16, DeepSeek-R1 in native FP8) — it only becomes an |
| |
386 | + # artifact when the same organisation also has the plain model entity (tencent/Hy-MT2-7B-GGUF next to Hy-MT2-7B) AND nobody |
| |
387 | + # else (evaluators, providers) knows it under its own identity; otherwise the pair is a merge candidate for review |
| |
388 | + same_org = [c for c in vk_candidates if c["organization_id"] == m["organization_id"] and not analyze_model_name((c["attributes"] or {}).get("hf_repo") or c["canonical_name"]).is_artifact] |
| |
389 | + if len(same_org) != 1 or not has_token: |
| |
390 | + rep.bump("_official_checkpoint_kept_as_model") |
| |
391 | + continue |
| |
392 | + schemes = set(official_idents.get(m["id"], {})) - {"hf_repo"} |
| |
393 | + if schemes: |
| |
394 | + if await _review(conn, "merge_candidate", sorted([m["id"], same_org[0]["id"]]), |
| |
395 | + f"'{m['canonical_name']}' and '{same_org[0]['canonical_name']}' look like one model published as two official checkpoints", |
| |
396 | + {"slugs": [m["slug"], same_org[0]["slug"]], "step": "artifacts"}, apply=apply): |
| |
397 | + rep.bump("review_official_checkpoint_pair") |
| |
398 | + rep.example(f"review: {m['slug']} (official repo with dtype tag, known to {sorted(schemes)}) vs {same_org[0]['slug']}") |
| |
399 | + continue |
| |
400 | + cid = same_org[0]["id"] |
| |
401 | + if attr_quantized or a.is_quantized or a_name.is_quantized or quant_attr in QUANT_ATTR_FORMATS: |
| |
402 | + kind = "quantization" |
| |
403 | + elif a.is_conversion or a_name.is_conversion or a.precision or quant_attr: |
| |
404 | + kind = "conversion" |
| |
405 | + else: |
| |
406 | + kind = "packaging" |
| |
407 | + if cid is None: |
| |
408 | + cid = base_of.get(m["id"]) |
| |
409 | + if cid and (cid not in by_id or cid == m["id"]): |
| |
410 | + cid = None |
| |
411 | + if cid is None: |
| |
412 | + candidates = vk_candidates |
| |
413 | + if len(candidates) > 1: |
| |
414 | + same_org = [c for c in candidates if c["organization_id"] == m["organization_id"]] |
| |
415 | + official_c = [c for c in candidates if (c["org_slug"] or "").lower() in set(official_orgs(probe)) or (c["org_hf"] or "").lower() in set(official_orgs(probe))] |
| |
416 | + candidates = same_org or official_c or candidates |
| |
417 | + if len(candidates) == 1: |
| |
418 | + cid = candidates[0]["id"] |
| |
419 | + if m["entity_type"] == "artifact" and m["canonical_id"] == cid: |
| |
420 | + continue |
| |
421 | + rep.bump("artifacts_marked" if cid else "artifacts_unresolved") |
| |
422 | + rep.example(f"artifact[{kind}] {m['slug']} → {by_id[cid]['slug'] if cid else 'UNRESOLVED'}") |
| |
423 | + if apply: |
| |
424 | + await execute(conn, """update entities set entity_type = 'artifact', artifact_kind = coalesce(artifact_kind, :k), canonical_id = :c, |
| |
425 | + identity_confidence = :ic, updated_at = now() where id = :id""", |
| |
426 | + k=kind, c=cid, ic="high" if cid else "low", id=m["id"]) |
| |
427 | + if cid: |
| |
428 | + await upsert_relation(conn, m["id"], "artifact_of", cid, {"artifact_kind": kind}, source_id=source_id) |
| |
429 | + await record_decision(conn, m["id"], cid, "variant_of", actor="canonicalize", payload={"artifact_kind": kind, "step": "artifacts"}) |
| |
430 | + if not cid: |
| |
431 | + await _review(conn, "unresolved_artifact", [m["id"]], f"'{m['canonical_name']}' is a {kind} artifact but its base model is unknown", |
| |
432 | + {"slug": m["slug"], "hf_repo": hf_repo, "variant_key": variant_key(probe)}, apply=apply) |
| |
433 | + |
| |
434 | + |
| |
435 | +# ---------------------------------------------------------------------------------------------- step: families |
| |
436 | +async def step_families(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
437 | + source_id = await registry_source_id(conn) |
| |
438 | + models = [m for m in await _models(conn) if _in_scope(scope, m["id"])] |
| |
439 | + orgs = await fetch_all(conn, "select id, slug, canonical_name, attributes from entities where merged_into is null and entity_type = any(cast(:t as text[]))", t=sorted(ORG_TYPES)) |
| |
440 | + org_by_key: dict[str, str] = {} |
| |
441 | + for o in orgs: |
| |
442 | + for k in _org_lookup_keys(o): |
| |
443 | + org_by_key.setdefault(k, o["id"]) |
| |
444 | + fam_rows = await fetch_all(conn, "select id, slug, canonical_name, organization_id from entities where entity_type = 'model_family' and merged_into is null") |
| |
445 | + families: dict[str, dict[str, Any]] = {f["canonical_name"].lower(): f for f in fam_rows} # one family per label |
| |
446 | + slugs_taken = {r["slug"] for r in await fetch_all(conn, "select slug from entities")} |
| |
447 | + for m in models: |
| |
448 | + label = family_release_hint(m["canonical_name"]) |
| |
449 | + if not label: |
| |
450 | + rep.bump("_no_family_hint") |
| |
451 | + continue |
| |
452 | + root = family_hint(m["canonical_name"]) or label |
| |
453 | + official = official_orgs(m["canonical_name"]) |
| |
454 | + org_id = next((org_by_key[k] for k in official if k in org_by_key), None) or m["organization_id"] |
| |
455 | + fam = families.get(label.lower()) |
| |
456 | + if fam is None: |
| |
457 | + # slug = slugify(label); on collision with any other entity (the model "gpt-5.5" itself) prefix with the organisation slug, then "family-" |
| |
458 | + base_slug = slugify(label) |
| |
459 | + org_slug = next((o["slug"] for o in orgs if o["id"] == org_id), None) |
| |
460 | + candidates = [base_slug] + ([f"{org_slug}-{base_slug}"] if org_slug else []) + [f"family-{base_slug}"] + [f"family-{base_slug}-{n}" for n in range(2, 20)] |
| |
461 | + slug = next(c for c in candidates if c not in slugs_taken) |
| |
462 | + rep.bump("families_created") |
| |
463 | + rep.example(f"family '{label}' ({slug}) root={root}") |
| |
464 | + fam = {"id": new_id("model_family"), "slug": slug, "canonical_name": label, "organization_id": org_id, "_new": True} |
| |
465 | + families[label.lower()] = fam |
| |
466 | + slugs_taken.add(slug) |
| |
467 | + if apply: |
| |
468 | + await execute(conn, """insert into entities (id, entity_type, canonical_name, slug, status, organization_id, attributes, provenance, first_seen_at, last_seen_at, identity_confidence) |
| |
469 | + values (:id, 'model_family', :n, :slug, 'active', :org, cast(:a as jsonb), cast(:p as jsonb), :fs, now(), 'high')""", |
| |
470 | + id=fam["id"], n=label, slug=slug, org=org_id, a=jsonb({"family_root": root, "label": label}), |
| |
471 | + p=jsonb({"family_root": {"source_id": source_id, "tier": 2, "extractor": "derived"}}), fs=m["first_seen_at"]) |
| |
472 | + await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:e, :a, :n, 'alias') on conflict do nothing", |
| |
473 | + e=fam["id"], a=label, n=normalize_alias(label)) |
| |
474 | + if m["family_id"] != fam["id"]: |
| |
475 | + rep.bump("models_linked") |
| |
476 | + if apply: |
| |
477 | + await execute(conn, "update entities set family_id = :f, first_seen_at = first_seen_at where id = :id", f=fam["id"], id=m["id"]) |
| |
478 | + await execute(conn, "update entities set first_seen_at = least(first_seen_at, :fs) where id = :f", fs=m["first_seen_at"], f=fam["id"]) |
| |
479 | + await upsert_relation(conn, m["id"], "member_of_family", fam["id"], source_id=source_id) |
| |
480 | + if not (m["attributes"] or {}).get("family"): |
| |
481 | + rep.bump("family_attribute_set") |
| |
482 | + if apply: |
| |
483 | + await _set_attributes(conn, m["id"], {"family": label}, source_id=source_id) |
| |
484 | + |
| |
485 | + |
| |
486 | +# ---------------------------------------------------------------------------------------------- step: licenses |
| |
487 | +async def step_licenses(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
488 | + source_id = await registry_source_id(conn) |
| |
489 | + rows = await fetch_all(conn, """select id, entity_type, canonical_name, slug, attributes from entities where merged_into is null |
| |
490 | + and entity_type = any(cast(:t as text[])) and (attributes ? 'license' or attributes ? 'openness' or attributes ? 'hf_repo') |
| |
491 | + order by first_seen_at, id""", t=list(LICENSED_TYPES)) |
| |
492 | + if scope is not None: |
| |
493 | + rows = [r for r in rows if r["id"] in scope] |
| |
494 | + lic_rows = await fetch_all(conn, "select id, slug from entities where entity_type = 'license'") |
| |
495 | + license_entities: dict[str, str] = {r["slug"]: r["id"] for r in lic_rows} |
| |
496 | + existing_rel = {(r["subject_id"], r["object_id"]) for r in await fetch_all(conn, "select subject_id, object_id from relations where predicate = 'uses_license' and valid_to is null")} |
| |
497 | + # derived claims that a higher-tier source already contradicts: stored once as `conflicting`, never re-proposed |
| |
498 | + contested: dict[tuple[str, str], Any] = {(r["entity_id"], r["property"]): r["value"] for r in await fetch_all( |
| |
499 | + conn, "select entity_id, property, value from claims where extractor = 'derived' and status = 'conflicting'")} |
| |
500 | + writer = _derived_writer(conn, source_id) if apply else None |
| |
501 | + for r in rows: |
| |
502 | + attrs = r["attributes"] or {} |
| |
503 | + raw = attrs.get("license") |
| |
504 | + key = normalize_license(raw) if isinstance(raw, str) else None |
| |
505 | + if key is None and isinstance(attrs.get("license_key"), str) and attrs["license_key"] in LICENSES: |
| |
506 | + key = attrs["license_key"] |
| |
507 | + if isinstance(raw, str) and key is None: |
| |
508 | + rep.bump("_license_unclassified") |
| |
509 | + rep.example(f"unclassified licence {raw!r} on {r['slug']}") if len(rep.examples) < 4 else None |
| |
510 | + if key: |
| |
511 | + lslug = key.lower() |
| |
512 | + lid = license_entities.get(lslug) |
| |
513 | + if lid is None: |
| |
514 | + info = LICENSES[key] |
| |
515 | + lid = new_id("license") |
| |
516 | + license_entities[lslug] = lid |
| |
517 | + rep.bump("license_entities_created") |
| |
518 | + if apply: |
| |
519 | + await execute(conn, """insert into entities (id, entity_type, canonical_name, slug, status, description, attributes, provenance, identity_confidence) |
| |
520 | + values (:id, 'license', :n, :slug, 'active', :d, cast(:a as jsonb), cast(:p as jsonb), 'high')""", |
| |
521 | + id=lid, n=info.label, slug=lslug, d=f"{info.label} — {info.category} licence" + (f" (SPDX {info.spdx})" if info.spdx else ""), |
| |
522 | + a=jsonb({**info.as_dict(), "license_key": key}), p=jsonb({"license_key": {"source_id": source_id, "tier": 2, "extractor": "derived"}})) |
| |
523 | + await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:e, :a, :n, 'alias') on conflict do nothing", |
| |
524 | + e=lid, a=info.label, n=normalize_alias(info.label)) |
| |
525 | + for alias in (key, info.spdx or key): |
| |
526 | + await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:e, :a, :n, 'alias') on conflict do nothing", |
| |
527 | + e=lid, a=alias, n=normalize_alias(alias)) |
| |
528 | + if (r["id"], lid) not in existing_rel: |
| |
529 | + rep.bump("uses_license_relations") |
| |
530 | + existing_rel.add((r["id"], lid)) |
| |
531 | + if apply: |
| |
532 | + await upsert_relation(conn, r["id"], "uses_license", lid, source_id=source_id) |
| |
533 | + if attrs.get("license_key") != key: |
| |
534 | + rep.bump("license_key_attribute_set") |
| |
535 | + if apply: |
| |
536 | + await _set_attributes(conn, r["id"], {"license_key": key}, source_id=source_id) |
| |
537 | + if r["entity_type"] not in ("model", "artifact"): |
| |
538 | + continue |
| |
539 | + # openness dimensions (derived claims, tier 2, never supersede a tier-1 statement) |
| |
540 | + openness_raw = attrs.get("openness") |
| |
541 | + openness_now = normalize_openness(openness_raw) if isinstance(openness_raw, str) else None |
| |
542 | + weights: bool | None = None |
| |
543 | + if attrs.get("hf_repo") or attrs.get("model_card_url") or attrs.get("weights_url") or (openness_now or "").startswith(("open", "restricted")): |
| |
544 | + weights = True |
| |
545 | + elif openness_now == "proprietary": |
| |
546 | + weights = False |
| |
547 | + if weights is None: |
| |
548 | + rep.bump("_openness_unknown_skipped") |
| |
549 | + continue |
| |
550 | + dims = openness_dimensions(weights_available=weights, license_key=key) |
| |
551 | + derived = derive_openness(dims, license_key=key) |
| |
552 | + wanted: dict[str, Any] = {"weights_available": dims["weights_available"]} |
| |
553 | + for k in ("commercial_use_allowed", "redistribution_allowed", "derivatives_allowed"): |
| |
554 | + if dims[k] is not None: |
| |
555 | + wanted[k] = dims[k] |
| |
556 | + if derived != "unknown": |
| |
557 | + wanted["openness"] = derived |
| |
558 | + changed = {k: v for k, v in wanted.items() if not _same(attrs.get(k), v) and not _same(contested.get((r["id"], k)), v)} |
| |
559 | + if not changed: |
| |
560 | + continue |
| |
561 | + rep.bump("openness_claims_written", len(changed)) |
| |
562 | + if "openness" in changed: |
| |
563 | + rep.bump("openness_category_changed") |
| |
564 | + rep.example(f"{r['slug']}: openness {openness_raw!r} → {derived} (licence {key})") |
| |
565 | + if writer is not None: |
| |
566 | + facts = Facts() |
| |
567 | + ref = EntityRef(entity_type=r["entity_type"], name=r["canonical_name"], id=r["id"]) |
| |
568 | + for k, v in changed.items(): |
| |
569 | + facts.claim(ref, k, v, confidence="high") |
| |
570 | + await writer.write(facts) |
| |
571 | + if writer is not None: |
| |
572 | + rep.bump("_conflicting_derived_claims", writer.stats.conflicts) |
| |
573 | + |
| |
574 | + |
| |
575 | +# ---------------------------------------------------------------------------------------------- step: taxonomy |
| |
576 | +async def step_taxonomy(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
577 | + props = sorted(TAXONOMY_PROPERTIES) |
| |
578 | + rows = await fetch_all(conn, """select c.id, c.entity_id, c.property, c.value, c.value_raw, e.entity_type from claims c join entities e on e.id = c.entity_id |
| |
579 | + where c.status = 'current' and c.property = any(cast(:p as text[]))""", p=props) |
| |
580 | + if scope is not None: |
| |
581 | + rows = [r for r in rows if r["entity_id"] in scope] |
| |
582 | + mappings: dict[tuple[str, str], str | None] = {} |
| |
583 | + |
| |
584 | + def note(maps: list[tuple[str, str, str | None]]) -> None: |
| |
585 | + for d, rw, cn in maps: |
| |
586 | + if cn is not None and rw == cn: |
| |
587 | + continue # identity mapping: nothing to learn |
| |
588 | + mappings[(d, rw)] = cn if cn is not None else mappings.get((d, rw)) |
| |
589 | + |
| |
590 | + for r in rows: |
| |
591 | + canon, raw, maps = normalize_property(r["entity_type"], r["property"], r["value"]) |
| |
592 | + note(maps) |
| |
593 | + if _same(canon, r["value"]): |
| |
594 | + continue |
| |
595 | + rep.bump(f"claims_normalized:{r['property']}") |
| |
596 | + rep.example(f"{r['entity_type']} {r['property']}: {json.dumps(r['value'], ensure_ascii=False)[:40]} → {json.dumps(canon, ensure_ascii=False)[:40]}") |
| |
597 | + if apply: |
| |
598 | + keep_raw = r["value_raw"] or (r["value"] if isinstance(r["value"], str) else json.dumps(r["value"], ensure_ascii=False)) |
| |
599 | + keep_raw = keep_raw if keep_raw != canon else None |
| |
600 | + await execute(conn, "update claims set value = cast(:v as jsonb), value_text = :vt, value_raw = :raw where id = :id", |
| |
601 | + v=jsonb(canon), vt=canon[:2000] if isinstance(canon, str) else None, raw=keep_raw, id=r["id"]) |
| |
602 | + attrs = {r["property"]: canon} |
| |
603 | + if keep_raw: |
| |
604 | + attrs[f"{r['property']}_raw"] = keep_raw |
| |
605 | + await execute(conn, "update entities set attributes = attributes || cast(:a as jsonb), updated_at = now() where id = :id", a=jsonb(attrs), id=r["entity_id"]) |
| |
606 | + if r["property"] == "status" and isinstance(canon, str): |
| |
607 | + await execute(conn, "update entities set status = :s where id = :id", s=canon[:40], id=r["entity_id"]) |
| |
608 | + # attributes without a current claim (seeded/imported values) |
| |
609 | + for prop in props: |
| |
610 | + ents = await fetch_all(conn, """select e.id, e.entity_type, e.attributes->:p as value from entities e where e.merged_into is null and e.attributes ? :p |
| |
611 | + and not exists (select 1 from claims c where c.entity_id = e.id and c.property = :p and c.status = 'current')""", p=prop) |
| |
612 | + for e in ents: |
| |
613 | + if not _in_scope(scope, e["id"]): |
| |
614 | + continue |
| |
615 | + canon, raw, maps = normalize_property(e["entity_type"], prop, e["value"]) |
| |
616 | + note(maps) |
| |
617 | + if _same(canon, e["value"]): |
| |
618 | + continue |
| |
619 | + rep.bump(f"attributes_normalized:{prop}") |
| |
620 | + if apply: |
| |
621 | + keep_raw = e["value"] if isinstance(e["value"], str) else json.dumps(e["value"], ensure_ascii=False) |
| |
622 | + attrs = {prop: canon, f"{prop}_raw": keep_raw} |
| |
623 | + await execute(conn, "update entities set attributes = attributes || cast(:a as jsonb), updated_at = now() where id = :id", a=jsonb(attrs), id=e["id"]) |
| |
624 | + if prop == "status" and isinstance(canon, str): |
| |
625 | + await execute(conn, "update entities set status = :s where id = :id", s=canon[:40], id=e["id"]) |
| |
626 | + existing = {(m["domain"], m["raw"]): m["canonical"] for m in await fetch_all(conn, "select domain, raw, canonical from taxonomy_mappings")} |
| |
627 | + for (domain, raw_full), canon in mappings.items(): |
| |
628 | + raw = raw_full[:300] |
| |
629 | + if (domain, raw) in existing and (existing[(domain, raw)] == canon or canon is None): |
| |
630 | + continue |
| |
631 | + rep.bump("taxonomy_mappings_upserted") |
| |
632 | + if apply: |
| |
633 | + await execute(conn, """insert into taxonomy_mappings (domain, raw, canonical) values (:d, :r, :c) |
| |
634 | + on conflict (domain, raw) do update set canonical = coalesce(excluded.canonical, taxonomy_mappings.canonical), last_seen_at = now()""", |
| |
635 | + d=domain, r=raw, c=canon) |
| |
636 | + unknown = sorted({f"{d}:{r}" for (d, r), c in mappings.items() if c is None}) |
| |
637 | + if unknown: |
| |
638 | + rep.notes.append(f"{len(unknown)} raw value(s) without canonical mapping kept as-is: {', '.join(unknown[:15])}{'…' if len(unknown) > 15 else ''}") |
| |
639 | + |
| |
640 | + |
| |
641 | +# ---------------------------------------------------------------------------------------------- step: results |
| |
642 | +async def step_results(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
643 | + rows = await fetch_all(conn, """select r.id, r.model_id, r.config, r.metric, r.unit, r.score, r.config_key, r.trust_level, r.variant, r.run_group, r.extractor, r.valid_to, |
| |
644 | + r.is_current, s.key as source_key from benchmark_results r left join sources s on s.id = r.source_id""") |
| |
645 | + if scope is not None: |
| |
646 | + rows = [r for r in rows if r["model_id"] in scope] |
| |
647 | + for r in rows: |
| |
648 | + cfg = r["config"] or {} |
| |
649 | + ck = bench_ontology.config_key(cfg, r["metric"]) |
| |
650 | + trust = r["trust_level"] or bench_ontology.trust_level(r["source_key"], cfg, extractor=r["extractor"] or "deterministic") |
| |
651 | + variant = r["variant"] or bench_ontology.variant_from_config(cfg) |
| |
652 | + rg = r["run_group"] or bench_ontology.run_group_from_config(cfg) |
| |
653 | + current = r["is_current"] and r["valid_to"] is None |
| |
654 | + if (ck, trust, variant, rg, current) == (r["config_key"], r["trust_level"], r["variant"], r["run_group"], r["is_current"]): |
| |
655 | + continue |
| |
656 | + rep.bump("results_backfilled") |
| |
657 | + if apply: |
| |
658 | + await execute(conn, "update benchmark_results set config_key = :ck, trust_level = :t, variant = :v, run_group = :rg, is_current = :cur where id = :id", |
| |
659 | + ck=ck, t=trust, v=variant, rg=rg, cur=current, id=r["id"]) |
| |
660 | + if scope is None: |
| |
661 | + closed = await enforce_current_results(conn, dry_run=not apply) |
| |
662 | + else: |
| |
663 | + closed = sum([await enforce_current_results(conn, model_id=mid, dry_run=not apply) for mid in sorted(scope)]) |
| |
664 | + if closed: |
| |
665 | + rep.bump("older_run_rows_closed", closed) |
| |
666 | + live = await fetch_one(conn, "select count(*) as n from benchmark_results where valid_to is null and is_current") |
| |
667 | + rep.notes.append(f"current benchmark results after step: {int(live['n']) if live else 0}") |
| |
668 | + |
| |
669 | + |
| |
670 | +# ---------------------------------------------------------------------------------------------- step: events |
| |
671 | +async def step_events(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
672 | + lag = timedelta(days=BACKFILL_LAG_DAYS) |
| |
673 | + runs = await fetch_all(conn, """select connector_name, started_at, finished_at from connector_runs where status in ('success', 'unchanged', 'suspect', 'released') |
| |
674 | + order by connector_name, started_at""") |
| |
675 | + second_start: dict[str, datetime | None] = {} |
| |
676 | + per: dict[str, list[dict[str, Any]]] = defaultdict(list) |
| |
677 | + for r in runs: |
| |
678 | + per[r["connector_name"]].append(r) |
| |
679 | + for name, lst in per.items(): |
| |
680 | + second_start[name] = lst[1]["started_at"] if len(lst) > 1 else None |
| |
681 | + before = await fetch_one(conn, """select count(*) as n from change_events where is_backfill = false and event_type <> 'DOCUMENT_CHANGED' |
| |
682 | + and observed_at > now() - interval '24 hours'""") |
| |
683 | + rows = await fetch_all(conn, """select ev.id, ev.event_type, ev.connector_name, ev.observed_at, ev.effective_at, ev.is_backfill, ev.group_key, ev.entity_id, |
| |
684 | + e.attributes->>'release_date' as release_date, e.attributes->>'published_at' as published_at |
| |
685 | + from change_events ev left join entities e on e.id = ev.entity_id""") |
| |
686 | + from aiatlas.sdk.extract.dates import parse_datetime |
| |
687 | + |
| |
688 | + if scope is not None: |
| |
689 | + rows = [r for r in rows if r["entity_id"] in scope] |
| |
690 | + for ev in rows: |
| |
691 | + obs = ev["observed_at"] |
| |
692 | + bf = False |
| |
693 | + if ev["connector_name"] in STRUCTURAL_CONNECTORS: |
| |
694 | + bf = True # merges, folds, derived corrections: bookkeeping about our own data, never news |
| |
695 | + elif ev["connector_name"] and ev["connector_name"] in per: |
| |
696 | + s2 = second_start.get(ev["connector_name"]) |
| |
697 | + if s2 is None or obs < s2: |
| |
698 | + bf = True |
| |
699 | + elif ev["connector_name"] and ev["connector_name"] not in per: |
| |
700 | + bf = True # connector without any successful run yet → initial load |
| |
701 | + if not bf and ev["effective_at"] is not None and ev["effective_at"] < obs - lag: |
| |
702 | + bf = True |
| |
703 | + if not bf and ev["event_type"].startswith("NEW_"): |
| |
704 | + hint = parse_datetime(ev["release_date"]) if ev["release_date"] else (parse_datetime(ev["published_at"]) if ev["published_at"] else None) |
| |
705 | + if hint is not None: |
| |
706 | + hint = hint if hint.tzinfo else hint.replace(tzinfo=UTC) |
| |
707 | + if hint < obs - lag: |
| |
708 | + bf = True |
| |
709 | + gk = group_key_for(ev["event_type"], ev["entity_id"], ev["effective_at"], obs) |
| |
710 | + if bf == ev["is_backfill"] and gk == ev["group_key"]: |
| |
711 | + continue |
| |
712 | + if bf != ev["is_backfill"]: |
| |
713 | + rep.bump("backfill_flag_set" if bf else "backfill_flag_cleared") |
| |
714 | + if gk != ev["group_key"]: |
| |
715 | + rep.bump("group_key_set") |
| |
716 | + if apply: |
| |
717 | + await execute(conn, "update change_events set is_backfill = :bf, group_key = :gk where id = :id", bf=bf, gk=gk, id=ev["id"]) |
| |
718 | + after = await fetch_one(conn, """select count(*) as n from change_events where is_backfill = false and event_type <> 'DOCUMENT_CHANGED' |
| |
719 | + and observed_at > now() - interval '24 hours'""") |
| |
720 | + rep.notes.append(f"live (non-backfill) events in the last 24 h: before {int(before['n']) if before else 0} → after {int(after['n']) if after else 0}" |
| |
721 | + + ("" if apply else " (dry-run: after = before)")) |
| |
722 | + |
| |
723 | + |
| |
724 | +# ---------------------------------------------------------------------------------------------- step: anomalies |
| |
725 | +async def _collect_anomalies(conn: AsyncConnection) -> list[Anomaly]: |
| |
726 | + found: list[Anomaly] = [] |
| |
727 | + for r in await fetch_all(conn, "select id, canonical_name, attributes from entities where entity_type in ('model','artifact') and merged_into is null"): |
| |
728 | + found.extend(check_model(r["id"], r["canonical_name"], r["attributes"] or {})) |
| |
729 | + for r in await fetch_all(conn, "select id, canonical_name, attributes from entities where entity_type = 'hardware' and merged_into is null"): |
| |
730 | + found.extend(check_hardware(r["id"], r["canonical_name"], r["attributes"] or {})) |
| |
731 | + for r in await fetch_all(conn, """select p.*, m.canonical_name as model_name, v.canonical_name as provider_name from prices p |
| |
732 | + join entities m on m.id = p.model_id join entities v on v.id = p.provider_id where p.valid_to is null"""): |
| |
733 | + found.extend(check_price(r)) |
| |
734 | + for r in await fetch_all(conn, """select r.id, r.model_id, r.benchmark_id, r.score, r.metric, r.unit, r.evaluated_at, m.canonical_name as model_name, |
| |
735 | + b.canonical_name as benchmark_name, m.attributes->>'release_date' as model_release_date |
| |
736 | + from benchmark_results r join entities m on m.id = r.model_id join entities b on b.id = r.benchmark_id |
| |
737 | + where r.valid_to is null and r.is_current"""): |
| |
738 | + found.extend(check_result(r)) |
| |
739 | + return found |
| |
740 | + |
| |
741 | + |
| |
742 | +async def step_anomalies(conn: AsyncConnection, rep: StepReport, apply: bool, *, scope: set[str] | None = None) -> None: |
| |
743 | + found = await _collect_anomalies(conn) |
| |
744 | + if scope is not None: |
| |
745 | + found = [a for a in found if a.entity_id in scope] |
| |
746 | + keys: dict[str, Anomaly] = {} |
| |
747 | + for a in found: |
| |
748 | + keys.setdefault(a.dedupe_key[:400], a) |
| |
749 | + existing = {r["dedupe_key"]: r for r in await fetch_all(conn, "select dedupe_key, status, severity from anomalies")} |
| |
750 | + for key, a in keys.items(): |
| |
751 | + rep.bump(f"_by_severity:{a.severity}") |
| |
752 | + prev = existing.get(key) |
| |
753 | + if prev is None: |
| |
754 | + rep.bump("anomalies_opened") |
| |
755 | + rep.example(f"[{a.severity}] {a.message}") |
| |
756 | + elif prev["status"] == "resolved": |
| |
757 | + rep.bump("anomalies_reopened") |
| |
758 | + if apply: |
| |
759 | + await record(conn, a) |
| |
760 | + stale = [k for k, r in existing.items() if r["status"] == "open" and k not in keys] if scope is None else [] |
| |
761 | + if stale: |
| |
762 | + rep.bump("anomalies_resolved", len(stale)) |
| |
763 | + if apply: |
| |
764 | + await execute(conn, """update anomalies set status = 'resolved', resolved_at = now(), resolution = 'check no longer fires' |
| |
765 | + where status = 'open' and dedupe_key = any(cast(:k as text[]))""", k=stale) |
| |
766 | + |
| |
767 | + |
| |
768 | +_STEP_FUNCTIONS = { |
| |
769 | + "duplicates": step_duplicates, "variants": step_variants, "artifacts": step_artifacts, "families": step_families, "licenses": step_licenses, |
| |
770 | + "taxonomy": step_taxonomy, "results": step_results, "events": step_events, "anomalies": step_anomalies, |
| |
771 | +} |
| |
772 | + |
| |
773 | + |
| |
774 | +# ---------------------------------------------------------------------------------------------- quarantine release / discard |
| |
775 | +async def release_quarantine(conn: AsyncConnection, quarantine_id: str, *, actor: str = "admin") -> dict[str, Any]: |
| |
776 | + """Write the held facts of a quarantined run exactly as the connector would have (same source, snapshot, tier, run id).""" |
| |
777 | + from aiatlas.connectors import get as get_connector |
| |
778 | + |
| |
779 | + q = await fetch_one(conn, "select * from quarantined_runs where id = :id", id=quarantine_id) |
| |
780 | + if not q: |
| |
781 | + raise LookupError(f"quarantined run {quarantine_id} not found") |
| |
782 | + if q["status"] != "pending": |
| |
783 | + raise ValueError(f"quarantined run {quarantine_id} is already {q['status']}") |
| |
784 | + connector = get_connector(q["connector_name"]) |
| |
785 | + state = await fetch_one(conn, """select c.source_id, s.key as source_key from connectors c left join sources s on s.id = c.source_id where c.name = :n""", n=q["connector_name"]) |
| |
786 | + totals: dict[str, int] = defaultdict(int) |
| |
787 | + for item in q["facts"] or []: |
| |
788 | + facts = facts_from_json(item["facts"]) |
| |
789 | + fetched_at = datetime.fromisoformat(item["fetched_at"]) |
| |
790 | + writer = FactWriter(conn, source_id=state["source_id"] if state else None, snapshot_id=item.get("snapshot_id"), source_url=item.get("source_url"), |
| |
791 | + tier=connector.tier, connector_name=q["connector_name"], extractor="deterministic", extractor_version=connector.parser_version, |
| |
792 | + observed_at=fetched_at, run_id=q["run_id"], source_key=state["source_key"] if state else None) |
| |
793 | + ws = await writer.write(facts) |
| |
794 | + for k, v in ws.as_dict().items(): |
| |
795 | + totals[k] += v |
| |
796 | + main = facts.document_entity |
| |
797 | + if main and main.id is None: |
| |
798 | + await writer.resolver.resolve(main) |
| |
799 | + if main and main.id and item.get("doc_id"): |
| |
800 | + await execute(conn, "update documents set entity_id = coalesce(entity_id, :e), title = coalesce(:t, title) where id = :id", e=main.id, t=facts.document_title, id=item["doc_id"]) |
| |
801 | + if item.get("snapshot_id"): |
| |
802 | + await execute(conn, "update snapshots set processing_status = 'extracted' where id = :id", id=item["snapshot_id"]) |
| |
803 | + await execute(conn, "update quarantined_runs set status = 'released', resolved_at = now(), resolved_by = :who where id = :id", who=actor, id=quarantine_id) |
| |
804 | + await execute(conn, "update connector_runs set status = 'released' where id = :r", r=q["run_id"]) |
| |
805 | + await execute(conn, "update review_queue set status = 'approved', resolved_at = now() where dedupe_key = :d", d=f"quarantine:{quarantine_id}") |
| |
806 | + observed = (q["stats"] or {}).get("observed") |
| |
807 | + if observed and (q["stats"] or {}).get("full_extraction"): |
| |
808 | + baseline = (await fetch_one(conn, "select baseline from connectors where name = :n", n=q["connector_name"]) or {}).get("baseline") |
| |
809 | + nb = connector._next_baseline(baseline, observed) # noqa: SLF001 |
| |
810 | + await execute(conn, "update connectors set baseline = cast(:b as jsonb) where name = :n", b=jsonb(nb), n=q["connector_name"]) |
| |
811 | + await audit(conn, "quarantine.release", quarantine_id, {"connector": q["connector_name"], "run_id": q["run_id"], **totals}, actor=actor) |
| |
812 | + return {"id": quarantine_id, "connector": q["connector_name"], **totals} |
| |
813 | + |
| |
814 | + |
| |
815 | +async def discard_quarantine(conn: AsyncConnection, quarantine_id: str, *, actor: str = "admin", note: str | None = None) -> dict[str, Any]: |
| |
816 | + q = await fetch_one(conn, "select id, run_id, connector_name, status, facts from quarantined_runs where id = :id", id=quarantine_id) |
| |
817 | + if not q: |
| |
818 | + raise LookupError(f"quarantined run {quarantine_id} not found") |
| |
819 | + if q["status"] != "pending": |
| |
820 | + raise ValueError(f"quarantined run {quarantine_id} is already {q['status']}") |
| |
821 | + await execute(conn, "update quarantined_runs set status = 'discarded', resolved_at = now(), resolved_by = :who where id = :id", who=actor, id=quarantine_id) |
| |
822 | + await execute(conn, "update connector_runs set status = 'discarded' where id = :r", r=q["run_id"]) |
| |
823 | + await execute(conn, "update review_queue set status = 'rejected', resolved_at = now() where dedupe_key = :d", d=f"quarantine:{quarantine_id}") |
| |
824 | + for item in q["facts"] or []: |
| |
825 | + if item.get("snapshot_id"): |
| |
826 | + await execute(conn, "update snapshots set processing_status = 'discarded' where id = :id", id=item["snapshot_id"]) |
| |
827 | + await audit(conn, "quarantine.discard", quarantine_id, {"connector": q["connector_name"], "run_id": q["run_id"], "note": note}, actor=actor) |
| |
828 | + return {"id": quarantine_id, "connector": q["connector_name"], "status": "discarded"} |
| |
829 | + |
| |
830 | + |
| |
831 | +async def list_quarantine(conn: AsyncConnection, *, status: str = "pending", limit: int = 50) -> list[dict[str, Any]]: |
| |
832 | + return await fetch_all(conn, """select id, run_id, connector_name, reason, stats, status, created_at, resolved_at, resolved_by, jsonb_array_length(facts) as documents |
| |
833 | + from quarantined_runs where (cast(:s as text) = '' or status = :s) order by created_at desc limit :n""", s=status or "", n=limit) |
| |
834 | + |
| |
835 | + |
| |
836 | +__all__ = ["CANON_VERSION", "STEPS", "Report", "StepReport", "canonicalize", "discard_quarantine", "list_quarantine", "release_quarantine"] |