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.registry import load, org_ref, organizations, provider_ref14from aiatlas.sdk.facts import EntityRef, Facts15from aiatlas.sdk.writer import FactWriter1617log = logging.getLogger(__name__)18REGISTRY_SOURCE_KEY = "ai-atlas.registry"192021async def seed(conn: AsyncConnection) -> dict[str, Any]:22 out: dict[str, Any] = {}23 out["sources"] = await _seed_sources(conn)24 out["connectors"] = await _seed_connectors(conn)25 reg_source = await fetch_one(conn, "select id, tier from sources where key = :k", k=REGISTRY_SOURCE_KEY)26 assert reg_source27 out["organizations"] = await _seed_organizations(conn, reg_source["id"])28 out["providers"] = await _seed_providers(conn, reg_source["id"])29 out["benchmarks"] = await _seed_benchmarks(conn, reg_source["id"])30 out["hardware"] = await _seed_hardware(conn, reg_source["id"])31 out["domains"] = await _seed_domains(conn)32 return out333435async def _seed_sources(conn: AsyncConnection) -> int:36 n = 037 for s in load("sources"):38 org = organizations().get(s.get("organization", ""))39 org_id = None40 if org:41 row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=org["key"])42 org_id = row["entity_id"] if row else None43 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)44 values (:id, :key, :name, :domain, :org, :tier, :kind, :cat, :base, :rate, :interval, :enabled, :prio, :notes, cast(:meta as jsonb))45 on conflict (key) do update set name = excluded.name, domain = excluded.domain, organization_id = coalesce(excluded.organization_id, sources.organization_id),46 tier = excluded.tier, kind = excluded.kind, category = excluded.category, base_url = excluded.base_url, rate_limit_per_min = excluded.rate_limit_per_min,47 crawl_interval_s = excluded.crawl_interval_s, priority = excluded.priority, notes = excluded.notes, updated_at = now()""",48 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"),49 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),50 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",)}))51 n += 152 return n535455async def _seed_connectors(conn: AsyncConnection) -> int:56 n = 057 for name, cls in connector_registry().items():58 src = await fetch_one(conn, "select id from sources where key = :k", k=cls.source_key) if cls.source_key else None59 if cls.source_key and not src:60 log.warning("connector source missing in registry/sources.yaml", extra={"connector": name, "source_key": cls.source_key})61 await execute(conn, """insert into connectors (name, source_id, label, description, enabled, priority, interval_seconds, min_interval_seconds, max_interval_seconds,62 parser_version, rate_limit_per_min, expected_min_records, next_run_at, meta)63 values (:name, :src, :label, :desc, :enabled, :prio, :interval, :mini, :maxi, :pv, :rate, :emr, now(), cast(:meta as jsonb))64 on conflict (name) do update set source_id = coalesce(excluded.source_id, connectors.source_id), label = excluded.label, description = excluded.description,65 priority = excluded.priority, min_interval_seconds = excluded.min_interval_seconds, max_interval_seconds = excluded.max_interval_seconds,66 parser_version = excluded.parser_version, rate_limit_per_min = excluded.rate_limit_per_min, expected_min_records = excluded.expected_min_records,67 meta = connectors.meta || excluded.meta, updated_at = now()""",68 name=name, src=src["id"] if src else None, label=cls.label or name, desc=cls.description or None,69 enabled=bool(getattr(cls, "enabled_by_default", True)), prio=cls.priority, interval=cls.interval_seconds, mini=cls.min_interval_seconds,70 maxi=cls.max_interval_seconds, pv=cls.parser_version, rate=cls.rate_per_min, emr=cls.expected_min_records,71 meta=jsonb({"version": cls.version, "tier": cls.tier, "needs_llm": cls.needs_llm, "module": cls.__module__}))72 n += 173 return n747576def _writer(conn: AsyncConnection, source_id: str, url: str | None) -> FactWriter:77 return FactWriter(conn, source_id=source_id, snapshot_id=None, source_url=url, tier=2, connector_name="registry", extractor="curated", extractor_version="1")787980async def _seed_organizations(conn: AsyncConnection, source_id: str) -> int:81 n = 082 # parents first so `parent` relations resolve83 items = sorted(load("organizations"), key=lambda o: 0 if not o.get("parent") else 1)84 for o in items:85 facts = Facts()86 ref = org_ref(o["key"])87 facts.entities.append(ref)88 for prop in ("country", "headquarters", "founded", "website", "legal_name"):89 if o.get(prop):90 facts.claim(ref, prop, o[prop], source_url=o.get("source_url"))91 if o.get("domains"):92 facts.claim(ref, "domains", o["domains"], source_url=o.get("source_url"))93 if o.get("hf_org"):94 facts.claim(ref, "hf_org", o["hf_org"], source_url=f"https://huggingface.co/{o['hf_org']}")95 if o.get("github_org"):96 facts.claim(ref, "github_org", o["github_org"], source_url=f"https://github.com/{o['github_org']}")97 facts.claim(ref, "org_kind", o.get("type", "company"), source_url=o.get("source_url"))98 if o.get("parent") and o["parent"] in organizations():99 facts.relate(org_ref(o["parent"]), "owns", ref, source_url=o.get("source_url"))100 w = _writer(conn, source_id, o.get("source_url"))101 await w.write(facts)102 n += 1103 return n104105106async def _seed_providers(conn: AsyncConnection, source_id: str) -> int:107 n = 0108 for p in load("providers"):109 facts = Facts()110 ref = provider_ref(p["key"])111 facts.entities.append(ref)112 for prop in ("website", "pricing_url", "docs_url"):113 if p.get(prop):114 facts.claim(ref, prop, p[prop], source_url=p.get("website"))115 if p.get("organization") in organizations():116 facts.relate(org_ref(p["organization"]), "operates", ref, source_url=p.get("website"))117 await _writer(conn, source_id, p.get("website")).write(facts)118 n += 1119 return n120121122async def _seed_benchmarks(conn: AsyncConnection, source_id: str) -> int:123 n = 0124 for b in load("benchmarks"):125 facts = Facts()126 ref = EntityRef(entity_type="benchmark", name=b["name"], identifiers={"registry_benchmark": b["key"]}, aliases=list(b.get("aliases", [])), slug_hint=b["key"])127 facts.entities.append(ref)128 for prop in ("category", "task", "metric", "unit", "creator", "website", "paper", "known_limitations", "methodology"):129 if b.get(prop) not in (None, ""):130 facts.claim(ref, prop, b[prop], source_url=b.get("source_url"))131 await _writer(conn, source_id, b.get("source_url")).write(facts)132 n += 1133 return n134135136async def _seed_hardware(conn: AsyncConnection, source_id: str) -> int:137 n = 0138 for h in load("hardware"):139 facts = Facts()140 org = org_ref(h["manufacturer"]) if h.get("manufacturer") in organizations() else None141 ref = EntityRef(entity_type="hardware", name=h["name"], identifiers={"registry_hardware": h["key"]}, aliases=list(h.get("aliases", [])), slug_hint=h["key"], organization=org)142 facts.entities.append(ref)143 for prop in ("kind", "architecture", "release_date", "memory_gb", "memory_type", "memory_bandwidth_gbs", "tdp_watts", "runtimes", "form_factor", "price_usd"):144 if h.get(prop) not in (None, ""):145 facts.claim(ref, prop, h[prop], source_url=h.get("source_url"))146 facts.claim(ref, "manufacturer", organizations()[h["manufacturer"]]["name"] if org else h.get("manufacturer"), source_url=h.get("source_url"))147 facts.claim(ref, "spec_url", h.get("source_url"), source_url=h.get("source_url"))148 if org:149 facts.relate(org, "manufactures", ref, source_url=h.get("source_url"))150 await _writer(conn, source_id, h.get("source_url")).write(facts)151 n += 1152 return n153154155async def _seed_domains(conn: AsyncConnection) -> int:156 n = 0157 for o in load("organizations"):158 row = await fetch_one(conn, "select entity_id from entity_identifiers where scheme = 'registry_org' and value = :v", v=o["key"])159 if not row:160 continue161 for d in o.get("domains", []):162 await execute(conn, """insert into domains (domain, organization_id, trust_tier, category) values (:d, :o, 1, :c)163 on conflict (domain) do update set organization_id = excluded.organization_id""", d=d, o=row["entity_id"], c=o.get("type", "company"))164 n += 1165 return n166167168__all__ = ["seed"]169