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%
15.8 KB · 261 lines python
Raw Blame History
1"""/models listing (canonical universe, filters on `entities.attributes`, facets), /models/{slug} alias, /models/{a}/diff/{b}.23API 1.1: the default universe is CANONICAL MODELS (`entity_type = 'model' and merged_into is null`). `include=artifacts` adds4checkpoints / quantisations / conversions as rows with `entity_type: 'artifact'` and a `canonical` summary."""5from __future__ import annotations67import asyncio8from typing import Any910from fastapi import APIRouter, Query, Request1112from aiatlas.api.common import (13    ENTITY_COLS,14    ENTITY_FROM,15    MODEL_OR_ARTIFACT_UNIVERSE,16    MODEL_UNIVERSE,17    PAGINATION,18    ApiError,19    Pagination,20    attr_num,21    cached,22    entity_cols,23    entity_summary,24    flip_order,25    normalize,26    num_expr,27    openness_values,28    page,29)30from aiatlas.api.routers.entities import detail_for_type31from aiatlas.db import connection, fetch_all, fetch_val32from aiatlas.ontology.licenses import LICENSES, normalize_license33from aiatlas.services import cache3435router = APIRouter(prefix="/api/v1/models", tags=["models"])3637PARAMS = num_expr("e.attributes->>'parameter_count'")38CONTEXT = num_expr("e.attributes->>'context_length'")39DOWNLOADS = num_expr("e.attributes->>'metric.downloads'")40SORTS = {41    "updated": "e.updated_at desc", "name": "e.canonical_name asc", "params": f"{PARAMS} desc nulls last", "context": f"{CONTEXT} desc nulls last",42    "release": "e.attributes->>'release_date' desc nulls last", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "downloads": f"{DOWNLOADS} desc nulls last",43    "first_seen": "e.first_seen_at desc",44    "cheapest": "(select min(p.output_per_mtok) from prices p where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0) asc nulls last",45}46FAMILY_JOIN = "left join entities fam on fam.id = e.family_id"47CANONICAL_JOIN = "left join entities can on can.id = e.canonical_id left join entities cano on cano.id = can.organization_id"484950def license_match_sql(param: str = "license") -> str:51    """Match a canonical licence key (`attributes.license_key`) OR any raw label the ontology maps to it OR a plain ilike on the raw label."""52    return f"(e.attributes->>'license_key' = :{param} or lower(e.attributes->>'license') = any(cast(:{param}_raw as text[])) or e.attributes->>'license' ilike :{param})"535455def license_params(value: str) -> dict[str, Any]:56    key = value if value in LICENSES else (normalize_license(value) or value)57    info = LICENSES.get(key)58    raw = {key.lower(), value.lower()}59    if info:60        raw |= {a.lower() for a in info.aliases}61        if info.spdx:62            raw.add(info.spdx.lower())63    return {"license": key, "license_raw": sorted(raw)}646566def model_filters(*, q: str | None, org: str | None, family: str | None, openness: str | None, modality: str | None, status: str | None,67                  min_params: float | None, max_params: float | None, min_context: int | None, year_from: int | None, year_to: int | None,68                  license: str | None, include_artifacts: bool = False, reasoning: bool | None = None, trust: str | None = None) -> tuple[list[str], dict[str, Any]]:69    where = [MODEL_OR_ARTIFACT_UNIVERSE if include_artifacts else MODEL_UNIVERSE]70    p: dict[str, Any] = {}71    if q:72        where.append("(e.canonical_name ilike :qlike or e.search @@ plainto_tsquery('simple', :q) or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")73        p["q"], p["qlike"] = q, f"%{q}%"74    if org:75        where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)")76        p["org"] = org77    if family:78        # family slug / id (model_family entity via family_id) — falls back to the legacy `attributes.family` label79        where.append("(fam.slug = :family or fam.id = :family or fam.canonical_name ilike :family or e.attributes->>'family' ilike :family)")80        p["family"] = family81    if openness:82        where.append("e.attributes->>'openness' = any(cast(:openness as text[]))")83        p["openness"] = openness_values(openness)84    if modality:85        where.append("(e.attributes->'modalities' ? :modality or e.attributes->'modalities_input' ? :modality or e.attributes->'modalities_output' ? :modality)")86        p["modality"] = modality87    if status:88        where.append("e.status = any(cast(:status as text[]))")89        p["status"] = [v.strip() for v in status.split(",") if v.strip()]90    if min_params is not None:91        where.append(f"{PARAMS} >= :min_params")92        p["min_params"] = float(min_params)93    if max_params is not None:94        where.append(f"{PARAMS} <= :max_params")95        p["max_params"] = float(max_params)96    if min_context is not None:97        where.append(f"{CONTEXT} >= :min_context")98        p["min_context"] = float(min_context)99    if year_from is not None:100        where.append("left(e.attributes->>'release_date', 4) >= :yf")101        p["yf"] = str(year_from)102    if year_to is not None:103        where.append("left(e.attributes->>'release_date', 4) <= :yt")104        p["yt"] = str(year_to)105    if license:106        where.append(license_match_sql())107        p.update(license_params(license))108    if reasoning is not None:109        where.append("e.attributes->>'reasoning' = :reasoning")110        p["reasoning"] = "true" if reasoning else "false"111    if trust:112        where.append("e.identity_confidence = any(cast(:trust as text[]))")113        p["trust"] = [v.strip() for v in trust.split(",") if v.strip()]114    return where, p115116117def _canonical_licenses(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:118    """Fold raw licence labels into canonical keys (Apache 2.0 / apache-2.0 → Apache-2.0); unknown labels are kept raw under `raw: true`."""119    agg: dict[str, dict[str, Any]] = {}120    for r in rows:121        raw = r["value"]122        if raw in (None, ""):123            continue124        key = normalize_license(raw)125        k = key or str(raw)126        item = agg.setdefault(k, {"value": k, "label": LICENSES[key].label if key else str(raw), "category": LICENSES[key].category if key else "unknown",127                                  "count": 0, "raw_labels": [], **({} if key else {"raw": True})})128        item["count"] += int(r["count"])129        if str(raw) != k:130            item["raw_labels"].append(str(raw))131    return sorted(agg.values(), key=lambda x: (-x["count"], x["value"]))132133134async def model_facets(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:135    async def run(sql: str) -> list[dict[str, Any]]:136        async with connection() as conn:137            return await fetch_all(conn, sql, **params)138139    base = f"from {ENTITY_FROM} {FAMILY_JOIN} where {where_sql}"140    mods_from = (f"from {ENTITY_FROM} {FAMILY_JOIN} cross join lateral jsonb_array_elements_text(case when jsonb_typeof(e.attributes->'modalities') = 'array' "141                 f"then e.attributes->'modalities' else '[]'::jsonb end) m where {where_sql}")142    orgs, openness, mods, fams, years, lics, status, trust = await asyncio.gather(143        run(f"select eo.slug, eo.canonical_name as name, count(*) as count {base} and eo.id is not null group by 1, 2 order by 3 desc, 2 limit 60"),144        run(f"select e.attributes->>'openness' as value, count(*) as count {base} and e.attributes ? 'openness' group by 1 order by 2 desc"),145        run(f"select m.value, count(*) as count {mods_from} group by 1 order by 2 desc limit 30"),146        run(f"select coalesce(fam.slug, e.attributes->>'family') as value, coalesce(fam.canonical_name, e.attributes->>'family') as label, fam.id is not null as canonical, "147            f"count(*) as count {base} and (fam.id is not null or e.attributes ? 'family') group by 1, 2, 3 order by 4 desc, 2 limit 60"),148        run(f"select left(e.attributes->>'release_date', 4) as value, count(*) as count {base} and e.attributes ? 'release_date' group by 1 order by 1 desc limit 30"),149        run(f"select coalesce(e.attributes->>'license_key', e.attributes->>'license') as value, count(*) as count {base} and (e.attributes ? 'license' or e.attributes ? 'license_key') group by 1 order by 2 desc, 1 limit 80"),150        run(f"select e.status as value, count(*) as count {base} group by 1 order by 2 desc"),151        run(f"select e.identity_confidence as value, count(*) as count {base} group by 1 order by 2 desc"),152    )153    conv = lambda rows: [{"value": r["value"], "count": int(r["count"])} for r in rows if r["value"] not in (None, "")]154    return {"organizations": [{"slug": r["slug"], "name": r["name"], "count": int(r["count"])} for r in orgs], "openness": conv(openness), "modalities": conv(mods),155            "families": [{"value": r["value"], "label": r["label"], "canonical": bool(r["canonical"]), "count": int(r["count"])} for r in fams if r["value"]],156            "years": conv(years), "licenses": _canonical_licenses(lics), "status": conv(status),157            "trust": [{**x, "label": {"high": "Identity confirmed", "medium": "Identity probable", "low": "Identity uncertain"}.get(x["value"], x["value"])} for x in conv(trust)],158            "definitions": {"trust": "identity_confidence of the row: how sure AI Atlas is that this entry is one real model release (high | medium | low)",159                            "licenses": "canonical licence keys from the ontology; raw labels that could not be classified are flagged raw: true",160                            "families": "model_family entities (canonical: true) or legacy attribute labels (canonical: false)"}}161162163@router.get("")164@cached(300)165async def list_models(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, family: str | None = None,166                      openness: str | None = None, modality: str | None = None, status: str | None = None, min_params: float | None = Query(None, ge=0),167                      max_params: float | None = Query(None, ge=0), min_context: int | None = Query(None, ge=0), year_from: int | None = Query(None, ge=1950, le=2100),168                      year_to: int | None = Query(None, ge=1950, le=2100), license: str | None = None, sort: str = "updated", order: str = "",169                      facets: int = Query(0, ge=0, le=1), include: str | None = Query(None, description="`artifacts` restores the pre-1.1 universe (models + artifacts)"),170                      reasoning: int | None = Query(None, ge=0, le=1), trust: str | None = Query(None, description="identity_confidence: high,medium,low")) -> dict[str, Any]:171    if sort not in SORTS:172        raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")173    include_artifacts = "artifacts" in {x.strip() for x in (include or "").split(",")}174    where, params = model_filters(q=q, org=org, family=family, openness=openness, modality=modality, status=status, min_params=min_params, max_params=max_params,175                                  min_context=min_context, year_from=year_from, year_to=year_to, license=license, include_artifacts=include_artifacts,176                                  reasoning=None if reasoning is None else bool(reasoning), trust=trust)177    order_sql = flip_order(SORTS[sort], order)178    where_sql = " and ".join(where)179    async with connection() as conn:180        rows = await fetch_all(conn, f"""select {ENTITY_COLS}, e.family_id, e.canonical_id, e.artifact_kind, e.identity_confidence,181                                                fam.slug as family_slug, fam.canonical_name as family_name, {entity_cols("can", "c_")},182                                                bp.input_per_mtok as bp_input, bp.output_per_mtok as bp_output, bp.provider_slug as bp_provider_slug, bp.provider_name as bp_provider_name,183                                                bp.provider_id as bp_provider_id, bp.providers as bp_providers184                                         from {ENTITY_FROM} {FAMILY_JOIN} {CANONICAL_JOIN}185                                         left join lateral (select p.input_per_mtok, p.output_per_mtok, p.provider_id, pv.slug as provider_slug, pv.canonical_name as provider_name,186                                                                   (select count(distinct q.provider_id) from prices q where q.model_id = e.id and q.valid_to is null) as providers187                                                            from prices p join entities pv on pv.id = p.provider_id188                                                            where p.model_id = e.id and p.valid_to is null and p.output_per_mtok > 0 order by p.output_per_mtok asc, p.input_per_mtok asc nulls last limit 1) bp on true189                                         where {where_sql} order by {order_sql}, e.id limit :lim offset :off""",190                               lim=p.limit, off=p.offset, **params)191        total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} {FAMILY_JOIN} where {where_sql}", **params)192    items = []193    for r in rows:194        s = entity_summary(r) or {}195        if r.get("family_id"):196            s["family"] = {"id": r["family_id"], "slug": r["family_slug"], "name": r["family_name"]}197        if r["entity_type"] == "artifact":198            s["artifact_kind"] = r.get("artifact_kind")199            s["canonical"] = entity_summary(r, "c_")200        s["identity_confidence"] = r.get("identity_confidence")201        s["best_price"] = ({"input_per_mtok": r["bp_input"], "output_per_mtok": r["bp_output"], "unit": "USD per 1M tokens",202                            "provider": {"id": r["bp_provider_id"], "slug": r["bp_provider_slug"], "name": r["bp_provider_name"]}, "providers": int(r["bp_providers"] or 0),203                            "note": "cheapest current offer by output price"} if r.get("bp_output") is not None else None)204        items.append(s)205    out = page(items, int(total or 0), p)206    out["universe"] = "models+artifacts" if include_artifacts else "canonical models"207    if facets:208        out["facets"] = await _facets_cached(where_sql, params)209    return out210211212async def _facets_cached(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:213    key = "facets:models:v2:" + where_sql + ":" + repr(sorted(params.items()))214    hit = await cache.cache_get(key)215    if hit is not None:216        return hit217    value = normalize(await model_facets(where_sql, params))218    await cache.cache_set(key, value, 600)219    return value220221222@router.get("/{slug}")223@cached(300)224async def get_model(request: Request, slug: str) -> dict[str, Any]:225    """Accepts canonical models AND artifacts (an artifact detail carries `canonical` + `artifact_kind`); folded variants follow `merged_into`226    and report `redirected_from` so the web layer can 301."""227    return await detail_for_type(slug, ("model", "artifact"))228229230@router.get("/{a}/diff/{b}")231@cached(300)232async def model_diff(request: Request, a: str, b: str) -> dict[str, Any]:233    """Only the dimensions where two models differ, with a delta (numeric % change, list added/removed)."""234    from aiatlas.api.routers.compare import compare_entities235236    body = await compare_entities([a, b], diff_only=True)237    items = body["items"]238    if len(items) != 2:239        raise ApiError(400, "two models are required")240    va, vb = items[0]["values"], items[1]["values"]241    dims = []242    for d in body["dimensions"]:243        x, y = va.get(d["key"]), vb.get(d["key"])244        dims.append({**d, "a": x, "b": y, "delta": _delta(x, y, d.get("kind"))})245    return {"a": items[0]["entity"], "b": items[1]["entity"], "dimensions": dims, "comparability": body.get("comparability"),246            "note": "Only dimensions with differing observed values; numeric delta = (b − a) / a; lists show added/removed elements."}247248249def _delta(x: Any, y: Any, kind: str | None) -> dict[str, Any] | None:250    if isinstance(x, list) or isinstance(y, list):251        sx, sy = {str(v) for v in (x or [])}, {str(v) for v in (y or [])}252        return {"added": sorted(sy - sx), "removed": sorted(sx - sy)}253    try:254        fx, fy = float(x), float(y)255    except (TypeError, ValueError):256        return None257    return {"absolute": fy - fx, "percent": round((fy - fx) / fx * 100, 2) if fx else None}258259260__all__ = ["CONTEXT", "PARAMS", "attr_num", "license_match_sql", "license_params", "model_filters", "router"]261