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%
11.4 KB · 152 lines python
Raw Blame History
1"""/families · /families/{slug} — model families (API 1.1). Canonical `model_family` entities first; pre-canonicalisation data is served2from the legacy `attributes.family` label (items flagged `canonical: false`)."""3from __future__ import annotations45from collections import defaultdict6from typing import Any78from fastapi import APIRouter, Query, Request910from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, ApiError, cached, entity_summary11from aiatlas.api.detail import LINEAGE_PREDICATES12from aiatlas.db import connection, fetch_all, fetch_one, fetch_val13from aiatlas.ontology.licenses import LICENSES, normalize_license14from aiatlas.services.frontier import all_primary_groups, rank_rows1516router = APIRouter(prefix="/api/v1/families", tags=["families"])17MEMBER_SQL = "((f.id is not null and e.family_id = f.id) or (f.id is null and e.attributes->>'family' = f.label))"181920def _f(v: Any) -> float | None:21    if v is None or isinstance(v, bool):22        return None23    try:24        return float(v)25    except (TypeError, ValueError):26        return None272829def _ranks(groups: dict[str, dict[str, Any]]) -> dict[str, dict[str, dict[str, Any]]]:30    """{model_id: {benchmark slug: {rank, score, metric, config_key}}} over every benchmark's primary group."""31    out: dict[str, dict[str, dict[str, Any]]] = defaultdict(dict)32    for g in groups.values():33        for r in rank_rows(g["rows"], g["higher_is_better"]):34            out[r["model_id"]][g["benchmark"]["slug"]] = {"rank": r["rank"], "score": r["score"], "metric": g["metric"], "config_key": g["config_key"],35                                                          "higher_is_better": g["higher_is_better"], "benchmark": g["benchmark"]["slug"]}36    return out373839def _aggregate(members: list[dict[str, Any]], ranks: dict[str, dict[str, dict[str, Any]]]) -> dict[str, Any]:40    params = [p for p in (_f((m["attributes"] or {}).get("parameter_count")) for m in members) if p]41    dates = sorted(str((m["attributes"] or {}).get("release_date") or "")[:10] for m in members if (m["attributes"] or {}).get("release_date"))42    mods: set[str] = set()43    lics: dict[str, int] = defaultdict(int)44    best: dict[str, dict[str, Any]] = {}45    for m in members:46        a = m["attributes"] or {}47        for k in ("modalities", "modalities_input", "modalities_output"):48            if isinstance(a.get(k), list):49                mods |= {str(x).lower() for x in a[k]}50        key = a.get("license_key") or normalize_license(a.get("license"))51        if key or a.get("license"):52            lics[key or str(a.get("license"))] += 153        for b, rk in ranks.get(m["id"], {}).items():54            if b not in best or rk["rank"] < best[b]["rank"]:55                best[b] = {**rk, "model": m["slug"], "model_name": m["canonical_name"]}56    return {"model_count": len(members), "first_release": dates[0] if dates else None, "last_release": dates[-1] if dates else None,57            "param_range": {"min": min(params), "max": max(params)} if params else None, "modalities": sorted(mods),58            "licenses": [{"key": k, "label": LICENSES[k].label if k in LICENSES else k, "models": n} for k, n in sorted(lics.items(), key=lambda kv: -kv[1])],59            "benchmark_best": dict(sorted(best.items()))}606162async def _families(conn: Any) -> list[dict[str, Any]]:63    """Canonical model_family entities + legacy labels not yet backed by an entity."""64    fam = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'model_family' and e.merged_into is null order by e.canonical_name")65    labels = await fetch_all(conn, """select e.attributes->>'family' as label, min(e.organization_id) as organization_id, count(*) as n from entities e66                                      where e.entity_type = 'model' and e.merged_into is null and e.family_id is null and e.attributes ? 'family' group by 1 order by 1""")67    out = [{"id": r["id"], "slug": r["slug"], "name": r["canonical_name"], "label": r["canonical_name"], "canonical": True, "summary": entity_summary(r), "organization_id": r["organization_id"]} for r in fam]68    names = {x["name"].lower() for x in out}69    for r in labels:70        if r["label"] and r["label"].lower() not in names:71            out.append({"id": None, "slug": r["label"].lower().replace(" ", "-").replace(".", "-"), "name": r["label"], "label": r["label"], "canonical": False, "summary": None, "organization_id": r["organization_id"]})72    return out737475@router.get("")76@cached(300)77async def list_families(request: Request, q: str | None = Query(None, max_length=120), org: str | None = None, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0),78                        sort: str = Query("models", pattern="^(models|name|last_release)$")) -> dict[str, Any]:79    async with connection() as conn:80        fams = await _families(conn)81        members = await fetch_all(conn, f"""select e.id, e.slug, e.canonical_name, e.attributes, e.family_id, e.attributes->>'family' as label, e.organization_id, eo.slug as org_slug, eo.canonical_name as org_name82                                            from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and (e.family_id is not null or e.attributes ? 'family')""")83        groups = await all_primary_groups(conn)84    ranks = _ranks(groups)85    by_fid: dict[str, list[dict[str, Any]]] = defaultdict(list)86    by_label: dict[str, list[dict[str, Any]]] = defaultdict(list)87    for m in members:88        if m["family_id"]:89            by_fid[m["family_id"]].append(m)90        elif m["label"]:91            by_label[m["label"].lower()].append(m)92    items = []93    for f in fams:94        mem = by_fid.get(f["id"], []) if f["id"] else by_label.get(f["name"].lower(), [])95        if not mem and not f["canonical"]:96            continue97        if q and q.lower() not in f["name"].lower():98            continue99        orgs = defaultdict(int)100        for m in mem:101            if m["organization_id"]:102                orgs[(m["organization_id"], m["org_slug"], m["org_name"])] += 1103        top_org = max(orgs.items(), key=lambda kv: kv[1])[0] if orgs else None104        if org and not (top_org and (top_org[1] == org or top_org[0] == org or (top_org[2] or "").lower() == org.lower())):105            continue106        items.append({"id": f["id"], "slug": f["slug"], "name": f["name"], "canonical": f["canonical"], "entity_type": "model_family",107                      "organization": {"id": top_org[0], "slug": top_org[1], "name": top_org[2]} if top_org else None, **_aggregate(mem, ranks)})108    key = {"models": lambda x: (-x["model_count"], x["name"]), "name": lambda x: x["name"].lower(), "last_release": lambda x: (x["last_release"] or "", x["name"])}[sort]109    items.sort(key=key, reverse=(sort == "last_release"))110    return {"items": items[offset:offset + limit], "total": len(items), "limit": limit, "offset": offset,111            "note": "canonical: true = model_family entity; false = legacy attributes.family label awaiting canonicalisation. benchmark_best = best rank of any member in each benchmark's primary group."}112113114@router.get("/{slug}")115@cached(300)116async def family_detail(request: Request, slug: str, limit: int = Query(100, ge=1, le=500)) -> dict[str, Any]:117    async with connection() as conn:118        fam = await fetch_one(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'model_family' and e.merged_into is null and (e.slug = :s or e.id = :s or e.canonical_name ilike :s) limit 1", s=slug)119        label = fam["canonical_name"] if fam else slug120        if fam:121            members = await fetch_all(conn, f"select {ENTITY_COLS}, e.family_id from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and (e.family_id = :fid or e.attributes->>'family' ilike :label) "122                                            f"order by e.attributes->>'release_date' asc nulls last, e.canonical_name limit :lim", fid=fam["id"], label=label, lim=limit)123        else:124            members = await fetch_all(conn, f"select {ENTITY_COLS}, e.family_id from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and e.family_id is null "125                                            f"and (e.attributes->>'family' ilike :label or replace(replace(lower(e.attributes->>'family'), ' ', '-'), '.', '-') = lower(:slug)) "126                                            f"order by e.attributes->>'release_date' asc nulls last, e.canonical_name limit :lim", label=label, slug=slug, lim=limit)127        if not fam and not members:128            raise ApiError(404, f"family {slug!r} not found")129        ids = [m["id"] for m in members]130        arts = await fetch_val(conn, "select count(*) from entities where entity_type = 'artifact' and merged_into is null and canonical_id = any(cast(:ids as text[]))", ids=ids) if ids else 0131        providers = await fetch_all(conn, f"select distinct {ENTITY_COLS} from prices p join entities e on e.id = p.provider_id left join entities eo on eo.id = e.organization_id where p.model_id = any(cast(:ids as text[])) and p.valid_to is null order by e.canonical_name", ids=ids) if ids else []132        edges = await fetch_all(conn, "select r.subject_id, r.predicate, r.object_id from relations r where r.valid_to is null and r.predicate = any(cast(:p as text[])) and r.subject_id = any(cast(:ids as text[])) and r.object_id = any(cast(:ids as text[]))",133                                p=[*LINEAGE_PREDICATES, "superseded_by"], ids=ids) if ids else []134        events = await fetch_all(conn, """select ev.entity_id, ev.event_type, ev.summary, ev.occurred_at, ev.importance from change_events ev where ev.entity_id = any(cast(:ids as text[]))135                                          and ev.event_type in ('NEW_MODEL','RELEASE','DEPRECATION_ANNOUNCED','RETIREMENT_ANNOUNCED','STATUS_CHANGED') order by ev.occurred_at asc limit 500""", ids=ids) if ids else []136        groups = await all_primary_groups(conn)137    ranks = _ranks(groups)138    mem_summ = []139    for m in members:140        a = m["attributes"] or {}141        rk = ranks.get(m["id"], {})142        mem_summ.append({"model": entity_summary(m), "key_facts": {k: a.get(k) for k in ("release_date", "parameter_count", "active_parameter_count", "context_length", "openness", "license", "modalities", "status") if a.get(k) not in (None, "", [])},143                         "benchmark_ranks": {b: v["rank"] for b, v in sorted(rk.items())},144                         "benchmark_best": [{"benchmark": b, "rank": v["rank"], "score": v["score"], "metric": v["metric"], "config_key": v["config_key"]} for b, v in sorted(rk.items(), key=lambda kv: kv[1]["rank"])]})145    timeline = [{"date": str((m["attributes"] or {}).get("release_date") or "")[:10] or None, "kind": "release", "model": {"id": m["id"], "slug": m["slug"], "name": m["canonical_name"]}} for m in members]146    timeline += [{"date": e["occurred_at"], "kind": e["event_type"], "summary": e["summary"], "model_id": e["entity_id"]} for e in events if e["event_type"] != "NEW_MODEL"]147    timeline.sort(key=lambda x: str(x["date"] or ""))148    return {"id": fam["id"] if fam else None, "slug": fam["slug"] if fam else slug, "name": label, "canonical": bool(fam), "entity_type": "model_family",149            "summary": entity_summary(fam) if fam else None, **_aggregate(members, ranks), "members": mem_summ, "artifacts_count": int(arts or 0),150            "providers": [entity_summary(p) for p in providers], "lineage": [{"source": e["subject_id"], "target": e["object_id"], "predicate": e["predicate"]} for e in edges],151            "timeline": timeline, "note": None if fam else "served from the legacy attributes.family label (no model_family entity yet)"}152