SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
3.6 KB · 106 lines python
Raw Blame History
1"""Curated registries (YAML under /registry) and their loaders."""2from __future__ import annotations34from functools import lru_cache5from pathlib import Path6from typing import Any78import yaml910REGISTRY_DIR = Path(__file__).resolve().parents[3] / "registry"111213@lru_cache14def load(name: str) -> list[dict[str, Any]]:15    """`registry/<name>.yaml` merged with fragment files `registry/<name>.d/*.yaml` (same top-level key). Later fragments may not16    redefine a key — duplicates are an error."""17    paths = [REGISTRY_DIR / f"{name}.yaml", *sorted((REGISTRY_DIR / f"{name}.d").glob("*.yaml"))]18    items: list[dict[str, Any]] = []19    for path in paths:20        if not path.exists():21            continue22        with path.open("r", encoding="utf-8") as fh:23            data = yaml.safe_load(fh) or {}24        items.extend(data.get(name) or [])25    keys = [i["key"] for i in items]26    dupes = {k for k in keys if keys.count(k) > 1}27    if dupes:28        raise ValueError(f"duplicate keys in {name}.yaml: {sorted(dupes)}")29    return items303132@lru_cache33def organizations() -> dict[str, dict[str, Any]]:34    return {o["key"]: o for o in load("organizations")}353637@lru_cache38def providers() -> dict[str, dict[str, Any]]:39    return {p["key"]: p for p in load("providers")}404142@lru_cache43def org_by_hf(hf_org: str) -> dict[str, Any] | None:44    for o in organizations().values():45        if o.get("hf_org", "").lower() == hf_org.lower():46            return o47    return None484950@lru_cache51def org_by_github(gh_org: str) -> dict[str, Any] | None:52    for o in organizations().values():53        if o.get("github_org", "").lower() == gh_org.lower():54            return o55    return None565758@lru_cache59def org_by_domain(domain: str) -> dict[str, Any] | None:60    d = domain.lower()61    d = d.removeprefix("www.")62    for o in organizations().values():63        for od in o.get("domains", []):64            if d == od or d.endswith("." + od):65                return o66    return None676869def org_ref(key: str):  # type: ignore[no-untyped-def]70    """EntityRef for a registry organization (identifiers make resolution deterministic)."""71    from aiatlas.sdk.facts import EntityRef7273    o = organizations()[key]74    ids: dict[str, str] = {"registry_org": key}75    if o.get("domains"):76        ids["domain"] = o["domains"][0]77    if o.get("hf_org"):78        ids["hf_org"] = o["hf_org"]79    if o.get("github_org"):80        ids["github_org"] = o["github_org"]81    return EntityRef(entity_type=o.get("type", "company") if o.get("type") in ("company", "lab", "organization", "university") else "company",82                     name=o["name"], identifiers=ids, aliases=list(o.get("aliases", [])), slug_hint=key)838485def provider_ref(key: str):  # type: ignore[no-untyped-def]86    from aiatlas.sdk.facts import EntityRef8788    p = providers()[key]89    ids = {"registry_provider": key}90    if p.get("openrouter_slug"):91        ids["openrouter_provider"] = p["openrouter_slug"]92    org = org_ref(p["organization"]) if p.get("organization") in organizations() else None93    # providers share names with their companies ("Anthropic" the company vs "Anthropic API" the provider): keep slugs distinct and readable94    slug = key if key not in organizations() else f"{key}-api"95    return EntityRef(entity_type="provider", name=p["name"], identifiers=ids, aliases=list(p.get("aliases", [])), slug_hint=slug, organization=org)969798def provider_by_openrouter(slug: str) -> str | None:99    for k, p in providers().items():100        if p.get("openrouter_slug") == slug:101            return k102    return None103104105__all__ = ["REGISTRY_DIR", "load", "org_by_domain", "org_by_github", "org_by_hf", "org_ref", "organizations", "provider_by_openrouter", "provider_ref", "providers"]106