SPB Git

spb/ai-atlas

Public
8.4 KB · 119 lines python
1"""/compare?ids=a,b,… — side-by-side of 2–6 entities of the same type with per-type dimension sets and provenance."""2from __future__ import annotations34from typing import Any56from fastapi import APIRouter, Query, Request78from aiatlas.api.common import (9    COMPANY_TYPES,10    PRICE_COLS,11    PRICE_FROM,12    RESULT_COLS,13    RESULT_FROM,14    ApiError,15    cached,16    csv,17    enrich_provenance,18    entity_summary,19    price_row,20    resolve_entity,21    result_row,22)23from aiatlas.db import connection, fetch_all2425router = APIRouter(prefix="/api/v1/compare", tags=["compare"])2627D = lambda key, label, kind="text", unit=None, source="attr": {"key": key, "label": label, "kind": kind, **({"unit": unit} if unit else {}), "source": source}2829DIMENSIONS: dict[str, list[dict[str, Any]]] = {30    "model": [D("parameter_count", "Parameters", "number", "params"), D("active_parameter_count", "Active parameters", "number", "params"),31              D("context_length", "Context window", "number", "tokens"), D("max_output_tokens", "Max output", "number", "tokens"), D("openness", "Openness"),32              D("license", "License"), D("modalities", "Modalities", "list"), D("release_date", "Release date", "date"), D("knowledge_cutoff", "Knowledge cutoff", "date"),33              D("status", "Status"), D("family", "Family"), D("architecture", "Architecture"),34              D("best_input_per_mtok", "Best input price", "number", "USD / 1M tokens", "prices"), D("best_output_per_mtok", "Best output price", "number", "USD / 1M tokens", "prices"),35              D("provider_count", "Providers", "number", None, "prices")],36    "provider": [D("website", "Website"), D("pricing_url", "Pricing page"), D("model_count", "Models priced", "number", None, "prices"),37                 D("min_input_per_mtok", "Min input price", "number", "USD / 1M tokens", "prices"), D("min_output_per_mtok", "Min output price", "number", "USD / 1M tokens", "prices"),38                 D("features", "Features", "list", None, "prices")],39    "hardware": [D("kind", "Kind"), D("manufacturer", "Manufacturer"), D("memory_gb", "Memory", "number", "GB"), D("memory_bandwidth_gbs", "Memory bandwidth", "number", "GB/s"),40                 D("tdp_watts", "TDP", "number", "W"), D("architecture", "Architecture"), D("memory_type", "Memory type"), D("runtimes", "Runtimes", "list"), D("release_date", "Release", "date")],41    "framework": [D("latest_version", "Latest version"), D("latest_release_at", "Latest release", "date"), D("license", "License"), D("language", "Language"),42                  D("metric.stars", "GitHub stars", "number"), D("metric.forks", "Forks", "number"), D("repository_url", "Repository")],43    "company": [D("country", "Country"), D("founded", "Founded", "date"), D("headquarters", "Headquarters"), D("org_kind", "Kind"), D("website", "Website"),44                D("model_count", "Models", "number", None, "graph"), D("paper_count", "Papers", "number", None, "graph")],45    "benchmark": [D("category", "Category"), D("metric", "Metric"), D("unit", "Unit"), D("task", "Task"), D("result_count", "Results", "number", None, "results")],46}474849@router.get("")50@cached(300)51async def compare(request: Request, ids: str = Query(..., description="2–6 slugs or ids, comma-separated")) -> dict[str, Any]:52    keys = csv(ids)53    if not 2 <= len(keys) <= 6:54        raise ApiError(400, "ids must list between 2 and 6 entities")55    async with connection() as conn:56        rows = [await resolve_entity(conn, k) for k in keys]57        types = {r["entity_type"] for r in rows}58        etype = rows[0]["entity_type"]59        if etype in COMPANY_TYPES:60            etype = "company"61            types = {"company"}62        if len(types) != 1:63            raise ApiError(400, f"all entities must share one type, got {', '.join(sorted(types))}")64        dims = [dict(d) for d in DIMENSIONS.get(etype, [])]65        eids = [r["id"] for r in rows]66        items: list[dict[str, Any]] = []67        extra: dict[str, dict[str, Any]] = {eid: {} for eid in eids}68        prices_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids}69        results_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids}70        if etype == "model":71            for p in await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where p.model_id = any(cast(:ids as text[])) and p.valid_to is null order by p.input_per_mtok nulls last", ids=eids):72                prices_by[p["m_id"]].append(price_row(p))73            for eid, plist in prices_by.items():74                ins = [x["input_per_mtok"] for x in plist if x["input_per_mtok"] is not None]75                outs = [x["output_per_mtok"] for x in plist if x["output_per_mtok"] is not None]76                extra[eid] = {"best_input_per_mtok": min(ins) if ins else None, "best_output_per_mtok": min(outs) if outs else None,77                              "provider_count": len({x["provider"]["id"] for x in plist if x["provider"]})}78            res = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where r.model_id = any(cast(:ids as text[])) and r.valid_to is null order by b.canonical_name, r.observed_at desc", ids=eids)79            per_bench: dict[str, dict[str, float]] = {}80            bench_meta: dict[str, dict[str, Any]] = {}81            for r in res:82                results_by[r["m_id"]].append(result_row(r))83                per_bench.setdefault(r["b_id"], {}).setdefault(r["m_id"], r["score"])84                bench_meta.setdefault(r["b_id"], {"slug": r["b_slug"], "name": r["b_canonical_name"], "unit": r.get("unit") or r.get("b_attributes", {}).get("unit"), "higher_is_better": r["higher_is_better"]})85            for bid, scores in per_bench.items():86                if all(eid in scores for eid in eids):87                    meta = bench_meta[bid]88                    dims.append({"key": f"bench:{meta['slug']}", "label": meta["name"], "kind": "number", "unit": meta["unit"], "source": "results", "higher_is_better": meta["higher_is_better"]})89                    for eid in eids:90                        extra[eid][f"bench:{meta['slug']}"] = scores[eid]91        elif etype == "provider":92            agg = await fetch_all(conn, """select provider_id, count(distinct model_id) as model_count, min(nullif(input_per_mtok, 0)) as min_input_per_mtok,93                                           min(nullif(output_per_mtok, 0)) as min_output_per_mtok, jsonb_agg(distinct k.key) filter (where k.key is not null) as features94                                           from prices p left join lateral jsonb_object_keys(p.features) k(key) on true95                                           where provider_id = any(cast(:ids as text[])) and valid_to is null group by 1""", ids=eids)96            for a in agg:97                extra[a["provider_id"]] = {"model_count": int(a["model_count"]), "min_input_per_mtok": a["min_input_per_mtok"], "min_output_per_mtok": a["min_output_per_mtok"], "features": a["features"] or []}98        elif etype == "company":99            agg = await fetch_all(conn, """select e.id, (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null) as model_count,100                                           (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) as paper_count101                                           from entities e where e.id = any(cast(:ids as text[]))""", ids=eids)102            for a in agg:103                extra[a["id"]] = {"model_count": int(a["model_count"]), "paper_count": int(a["paper_count"])}104        elif etype == "benchmark":105            agg = await fetch_all(conn, "select benchmark_id, count(*) as n from benchmark_results where benchmark_id = any(cast(:ids as text[])) and valid_to is null group by 1", ids=eids)106            for a in agg:107                extra[a["benchmark_id"]] = {"result_count": int(a["n"])}108        await enrich_provenance(conn, *(r.get("provenance") for r in rows))109        for r in rows:110            attrs = r.get("attributes") or {}111            prov = r.get("provenance") or {}112            values = {d["key"]: (extra[r["id"]].get(d["key"]) if d.get("source") != "attr" else attrs.get(d["key"])) for d in dims}113            item: dict[str, Any] = {"entity": entity_summary(r), "values": values, "provenance": {k: prov[k] for k in values if k in prov}}114            if etype == "model":115                item["prices"] = prices_by[r["id"]]116                item["results"] = results_by[r["id"]]117            items.append(item)118    return {"entity_type": etype, "dimensions": dims, "items": items}119
Files · c509d7e03d5aed8bc61fb24bc1c8bb3de1f4b9e2
Menu

Theme