"""/models listing (filters on `entities.attributes`, facets) and /models/{slug} alias.""" from __future__ import annotations import asyncio from typing import Any from fastapi import APIRouter, Query, Request from aiatlas.api.common import ( ENTITY_COLS, ENTITY_FROM, PAGINATION, ApiError, Pagination, cached, entity_summary, flip_order, normalize, num_expr, page, ) from aiatlas.api.routers.entities import detail_for_type from aiatlas.db import connection, fetch_all, fetch_val from aiatlas.services import cache router = APIRouter(prefix="/api/v1/models", tags=["models"]) PARAMS = num_expr("e.attributes->>'parameter_count'") CONTEXT = num_expr("e.attributes->>'context_length'") DOWNLOADS = num_expr("e.attributes->>'metric.downloads'") SORTS = { "updated": "e.updated_at desc", "name": "e.canonical_name asc", "params": f"{PARAMS} desc nulls last", "context": f"{CONTEXT} desc nulls last", "release": "e.attributes->>'release_date' desc nulls last", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "downloads": f"{DOWNLOADS} desc nulls last", "first_seen": "e.first_seen_at desc", } def model_filters(*, q: str | None, org: str | None, family: str | None, openness: str | None, modality: str | None, status: str | None, min_params: float | None, max_params: float | None, min_context: int | None, year_from: int | None, year_to: int | None, license: str | None) -> tuple[list[str], dict[str, Any]]: where = ["e.entity_type = 'model'", "e.merged_into is null"] p: dict[str, Any] = {} if q: 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))") p["q"], p["qlike"] = q, f"%{q}%" if org: where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)") p["org"] = org if family: where.append("e.attributes->>'family' ilike :family") p["family"] = family if openness: vals = [v.strip() for v in openness.split(",") if v.strip()] if "open" in vals: vals += ["open-weights", "open-source"] where.append("e.attributes->>'openness' = any(cast(:openness as text[]))") p["openness"] = vals if modality: where.append("(e.attributes->'modalities' ? :modality or e.attributes->'modalities_input' ? :modality or e.attributes->'modalities_output' ? :modality)") p["modality"] = modality if status: where.append("e.status = any(cast(:status as text[]))") p["status"] = [v.strip() for v in status.split(",") if v.strip()] if min_params is not None: where.append(f"{PARAMS} >= :min_params") p["min_params"] = float(min_params) if max_params is not None: where.append(f"{PARAMS} <= :max_params") p["max_params"] = float(max_params) if min_context is not None: where.append(f"{CONTEXT} >= :min_context") p["min_context"] = float(min_context) if year_from is not None: where.append("left(e.attributes->>'release_date', 4) >= :yf") p["yf"] = str(year_from) if year_to is not None: where.append("left(e.attributes->>'release_date', 4) <= :yt") p["yt"] = str(year_to) if license: where.append("e.attributes->>'license' ilike :license") p["license"] = license return where, p async def model_facets(where_sql: str, params: dict[str, Any]) -> dict[str, Any]: async def run(sql: str) -> list[dict[str, Any]]: async with connection() as conn: return await fetch_all(conn, sql, **params) base = f"from {ENTITY_FROM} where {where_sql}" mods_from = (f"from {ENTITY_FROM} cross join lateral jsonb_array_elements_text(case when jsonb_typeof(e.attributes->'modalities') = 'array' " f"then e.attributes->'modalities' else '[]'::jsonb end) m where {where_sql}") orgs, openness, mods, fams, years, lics, status = await asyncio.gather( 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"), run(f"select e.attributes->>'openness' as value, count(*) as count {base} and e.attributes ? 'openness' group by 1 order by 2 desc"), run(f"select m.value, count(*) as count {mods_from} group by 1 order by 2 desc limit 30"), run(f"select e.attributes->>'family' as value, count(*) as count {base} and e.attributes ? 'family' group by 1 order by 2 desc, 1 limit 60"), 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"), run(f"select e.attributes->>'license' as value, count(*) as count {base} and e.attributes ? 'license' group by 1 order by 2 desc, 1 limit 40"), run(f"select e.status as value, count(*) as count {base} group by 1 order by 2 desc"), ) conv = lambda rows: [{"value": r["value"], "count": int(r["count"])} for r in rows if r["value"] not in (None, "")] return {"organizations": [{"slug": r["slug"], "name": r["name"], "count": int(r["count"])} for r in orgs], "openness": conv(openness), "modalities": conv(mods), "families": conv(fams), "years": conv(years), "licenses": conv(lics), "status": conv(status)} @router.get("") @cached(300) async def list_models(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, family: str | None = None, openness: str | None = None, modality: str | None = None, status: str | None = None, min_params: float | None = Query(None, ge=0), 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), year_to: int | None = Query(None, ge=1950, le=2100), license: str | None = None, sort: str = "updated", order: str = "", facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]: if sort not in SORTS: raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") where, params = model_filters(q=q, org=org, family=family, openness=openness, modality=modality, status=status, min_params=min_params, max_params=max_params, min_context=min_context, year_from=year_from, year_to=year_to, license=license) order_sql = flip_order(SORTS[sort], order) where_sql = " and ".join(where) async with connection() as conn: rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {order_sql}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) out = page([entity_summary(r) for r in rows], int(total or 0), p) if facets: out["facets"] = await _facets_cached(where_sql, params) return out async def _facets_cached(where_sql: str, params: dict[str, Any]) -> dict[str, Any]: key = "facets:models:" + where_sql + ":" + repr(sorted(params.items())) hit = await cache.cache_get(key) if hit is not None: return hit value = normalize(await model_facets(where_sql, params)) await cache.cache_set(key, value, 600) return value @router.get("/{slug}") @cached(300) async def get_model(request: Request, slug: str) -> dict[str, Any]: return await detail_for_type(slug, ("model",))