HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""`aia seed` — idempotent: sources, connectors (from code), curated organizations / providers / benchmarks / hardware as entities2with claims attributed to the curated registry source (tier 2, each entry's own `source_url`)."""3from __future__ import annotations45import logging6from typing import Any78from sqlalchemy.ext.asyncio import AsyncConnection910from aiatlas.connectors import registry as connector_registry11from aiatlas.db import execute, fetch_one, jsonb12from aiatlas.ids import new_id13from aiatlas.ontology.benchmarks import metric_bounds, normalize_metric14from aiatlas.ontology.taxonomy import ORG_TYPE_DEFAULT_KIND, normalize_org_kind15from aiatlas.registry import load, org_ref, organizations, provider_ref16from aiatlas.sdk.facts import EntityRef, Facts17from aiatlas.sdk.writer import FactWriter1819log = logging.getLogger(__name__)20REGISTRY_SOURCE_KEY = "ai-atlas.registry"212223async def seed(conn: AsyncConnection) -> dict[str, Any]:24 out: dict[str, Any] = {}25 out["sources"] = await _seed_sources(conn)26 out["connectors"] = await _seed_connectors(conn)27 reg_source = await fetch_one(conn, "select id, tier from sources where key = :k", k=REGISTRY_SOURCE_KEY)28 assert reg_source29 out["organizations"] = await _seed_organizations(conn, reg_source["id"])30 out["providers"] = await _seed_providers(conn, reg_source["id"])31 out["benchmarks"] = await _seed_benchmarks(conn, reg_source["id"])32 out["hardware"] = await _seed_hardware(conn, reg_source["id"])33 out["domains"] = await _seed_domains(conn)34 return out353637async def _seed_sources(conn: AsyncConnection) -> int:38 n = 039 for s in load("sources"):40 org = organizations().get(s.get("organization", ""))41 org_id = None42 if org:43 row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=org["key"])44 org_id = row["entity_id"] if row else None45 await execute(conn, """insert into sources (id, key, name, domain, organization_id, tier, kind, category, base_url, rate_limit_per_min, crawl_interval_s, enabled, priority, notes, meta)46 values (:id, :key, :name, :domain, :org, :tier, :kind, :cat, :base, :rate, :interval, :enabled, :prio, :notes, cast(:meta as jsonb))47 on conflict (key) do update set name = excluded.name, domain = excluded.domain, organization_id = coalesce(excluded.organization_id, sources.organization_id),48 tier = excluded.tier, kind = excluded.kind, category = excluded.category, base_url = excluded.base_url, rate_limit_per_min = excluded.rate_limit_per_min,49 crawl_interval_s = excluded.crawl_interval_s, priority = excluded.priority, notes = excluded.notes, updated_at = now()""",50 id=new_id("source"), key=s["key"], name=s["name"], domain=s["domain"], org=org_id, tier=s.get("tier", 2), kind=s.get("kind", "website"),51 cat=s.get("category", "lab"), base=s.get("base_url"), rate=s.get("rate_limit_per_min", 30), interval=s.get("crawl_interval_s", 86400),52 enabled=s.get("enabled", True), prio=s.get("priority", 2), notes=s.get("notes"), meta=jsonb({k: v for k, v in s.items() if k not in ("key",)}))53 n += 154 return n555657async def _seed_connectors(conn: AsyncConnection) -> int:58 n = 059 for name, cls in connector_registry().items():60 src = await fetch_one(conn, "select id from sources where key = :k", k=cls.source_key) if cls.source_key else None61 if cls.source_key and not src:62 log.warning("connector source missing in registry/sources.yaml", extra={"connector": name, "source_key": cls.source_key})63 await execute(conn, """insert into connectors (name, source_id, label, description, enabled, priority, interval_seconds, min_interval_seconds, max_interval_seconds,64 parser_version, rate_limit_per_min, expected_min_records, next_run_at, meta)65 values (:name, :src, :label, :desc, :enabled, :prio, :interval, :mini, :maxi, :pv, :rate, :emr, now(), cast(:meta as jsonb))66 on conflict (name) do update set source_id = coalesce(excluded.source_id, connectors.source_id), label = excluded.label, description = excluded.description,67 priority = excluded.priority, min_interval_seconds = excluded.min_interval_seconds, max_interval_seconds = excluded.max_interval_seconds,68 parser_version = excluded.parser_version, rate_limit_per_min = excluded.rate_limit_per_min, expected_min_records = excluded.expected_min_records,69 meta = connectors.meta || excluded.meta, updated_at = now()""",70 name=name, src=src["id"] if src else None, label=cls.label or name, desc=cls.description or None,71 enabled=bool(getattr(cls, "enabled_by_default", True)), prio=cls.priority, interval=cls.interval_seconds, mini=cls.min_interval_seconds,72 maxi=cls.max_interval_seconds, pv=cls.parser_version, rate=cls.rate_per_min, emr=cls.expected_min_records,73 meta=jsonb({"version": cls.version, "tier": cls.tier, "needs_llm": cls.needs_llm, "module": cls.__module__}))74 n += 175 return n767778def _writer(conn: AsyncConnection, source_id: str, url: str | None) -> FactWriter:79 return FactWriter(conn, source_id=source_id, snapshot_id=None, source_url=url, tier=2, connector_name="registry", extractor="curated", extractor_version="1")808182async def _seed_organizations(conn: AsyncConnection, source_id: str) -> int:83 n = 084 # parents first so `parent` relations resolve85 items = sorted(load("organizations"), key=lambda o: 0 if not o.get("parent") else 1)86 for o in items:87 facts = Facts()88 ref = org_ref(o["key"])89 facts.entities.append(ref)90 for prop in ("country", "headquarters", "founded", "website", "legal_name"):91 if o.get(prop):92 facts.claim(ref, prop, o[prop], source_url=o.get("source_url"))93 if o.get("domains"):94 facts.claim(ref, "domains", o["domains"], source_url=o.get("source_url"))95 if o.get("hf_org"):96 facts.claim(ref, "hf_org", o["hf_org"], source_url=f"https://huggingface.co/{o['hf_org']}")97 if o.get("github_org"):98 facts.claim(ref, "github_org", o["github_org"], source_url=f"https://github.com/{o['github_org']}")99 kind = normalize_org_kind(o.get("kind")) or ORG_TYPE_DEFAULT_KIND.get(o.get("type", "company"))100 if kind:101 facts.claim(ref, "org_kind", kind, source_url=o.get("source_url"))102 if o.get("parent") and o["parent"] in organizations():103 facts.relate(org_ref(o["parent"]), "owns", ref, source_url=o.get("source_url"))104 w = _writer(conn, source_id, o.get("source_url"))105 await w.write(facts)106 n += 1107 return n108109110async def _seed_providers(conn: AsyncConnection, source_id: str) -> int:111 n = 0112 for p in load("providers"):113 facts = Facts()114 ref = provider_ref(p["key"])115 facts.entities.append(ref)116 for prop in ("website", "pricing_url", "docs_url"):117 if p.get(prop):118 facts.claim(ref, prop, p[prop], source_url=p.get("website"))119 if p.get("organization") in organizations():120 facts.relate(org_ref(p["organization"]), "operates", ref, source_url=p.get("website"))121 await _writer(conn, source_id, p.get("website")).write(facts)122 n += 1123 return n124125126def benchmark_entity_ref(b: dict[str, Any]) -> EntityRef:127 return EntityRef(entity_type="benchmark", name=b["name"], identifiers={"registry_benchmark": b["key"]}, aliases=list(b.get("aliases", [])), slug_hint=b["key"])128129130def family_heads(entries: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:131 """family key → representative entry (`family_head: true`, else the entry whose key equals the family key)."""132 heads: dict[str, dict[str, Any]] = {}133 for b in entries:134 fam = b.get("family")135 if fam and b.get("family_head"):136 heads[fam] = b137 for b in entries:138 fam = b.get("family")139 if fam and fam not in heads and b["key"] == fam:140 heads[fam] = b141 return heads142143144async def _seed_benchmarks(conn: AsyncConnection, source_id: str) -> int:145 n = 0146 entries = load("benchmarks")147 heads = family_heads(entries)148 for b in entries:149 facts = Facts()150 ref = benchmark_entity_ref(b)151 facts.entities.append(ref)152 src = b.get("source_url")153 for prop in ("category", "task", "unit", "creator", "website", "paper", "known_limitations", "methodology", "family", "variant", "version", "harness",154 "comparability_note", "metric_label"):155 if b.get(prop) not in (None, ""):156 facts.claim(ref, prop, b[prop], source_url=src)157 metric = normalize_metric(b.get("metric")) or b.get("metric")158 if metric:159 facts.claim(ref, "metric", metric, source_url=src)160 if b.get("metric") and b["metric"] != metric:161 facts.claim(ref, "metric_raw", b["metric"], source_url=src)162 lo, hi = metric_bounds(metric, b.get("unit"))163 for prop, fallback in (("metric_min", lo), ("metric_max", hi)):164 value = b.get(prop, fallback)165 if value is not None:166 facts.claim(ref, prop, value, source_url=src)167 facts.claim(ref, "higher_is_better", bool(b.get("higher_is_better", True)), source_url=src)168 facts.claim(ref, "family_head", bool(b.get("family_head", False)), source_url=src)169 head = heads.get(b.get("family") or "")170 if head and head["key"] != b["key"]:171 facts.relate(ref, "variant_of", benchmark_entity_ref(head), attributes={"family": b["family"], "variant": b.get("variant")}, source_url=src)172 await _writer(conn, source_id, src).write(facts)173 n += 1174 return n175176177async def _seed_hardware(conn: AsyncConnection, source_id: str) -> int:178 n = 0179 for h in load("hardware"):180 facts = Facts()181 org = org_ref(h["manufacturer"]) if h.get("manufacturer") in organizations() else None182 ref = EntityRef(entity_type="hardware", name=h["name"], identifiers={"registry_hardware": h["key"]}, aliases=list(h.get("aliases", [])), slug_hint=h["key"], organization=org)183 facts.entities.append(ref)184 for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "tdp_watts", "runtimes", "form_factor", "price_usd"):185 if h.get(prop) not in (None, ""):186 facts.claim(ref, prop, h[prop], source_url=h.get("source_url"))187 facts.claim(ref, "manufacturer", organizations()[h["manufacturer"]]["name"] if org else h.get("manufacturer"), source_url=h.get("source_url"))188 facts.claim(ref, "spec_url", h.get("source_url"), source_url=h.get("source_url"))189 if org:190 facts.relate(org, "manufactures", ref, source_url=h.get("source_url"))191 await _writer(conn, source_id, h.get("source_url")).write(facts)192 n += 1193 return n194195196async def _seed_domains(conn: AsyncConnection) -> int:197 n = 0198 for o in load("organizations"):199 row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=o["key"])200 if not row:201 continue202 for d in o.get("domains", []):203 await execute(conn, """insert into domains (domain, organization_id, trust_tier, category) values (:d, :o, 1, :c)204 on conflict (domain) do update set organization_id = excluded.organization_id""", d=d, o=row["entity_id"], c=o.get("type", "company"))205 n += 1206 return n207208209__all__ = ["seed"]210