1"""/models listing (filters on `entities.attributes`, facets) and /models/{slug} alias."""2from __future__ import annotations34import asyncio5from typing import Any67from fastapi import APIRouter, Query, Request89from aiatlas.api.common import (10 ENTITY_COLS,11 ENTITY_FROM,12 PAGINATION,13 ApiError,14 Pagination,15 cached,16 entity_summary,17 flip_order,18 normalize,19 num_expr,20 page,21)22from aiatlas.api.routers.entities import detail_for_type23from aiatlas.db import connection, fetch_all, fetch_val24from aiatlas.services import cache2526router = APIRouter(prefix="/api/v1/models", tags=["models"])2728PARAMS = num_expr("e.attributes->>'parameter_count'")29CONTEXT = num_expr("e.attributes->>'context_length'")30DOWNLOADS = num_expr("e.attributes->>'metric.downloads'")31SORTS = {32 "updated": "e.updated_at desc", "name": "e.canonical_name asc", "params": f"{PARAMS} desc nulls last", "context": f"{CONTEXT} desc nulls last",33 "release": "e.attributes->>'release_date' desc nulls last", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "downloads": f"{DOWNLOADS} desc nulls last",34 "first_seen": "e.first_seen_at desc",35}363738def model_filters(*, q: str | None, org: str | None, family: str | None, openness: str | None, modality: str | None, status: str | None,39 min_params: float | None, max_params: float | None, min_context: int | None, year_from: int | None, year_to: int | None,40 license: str | None) -> tuple[list[str], dict[str, Any]]:41 where = ["e.entity_type = 'model'", "e.merged_into is null"]42 p: dict[str, Any] = {}43 if q:44 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))")45 p["q"], p["qlike"] = q, f"%{q}%"46 if org:47 where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)")48 p["org"] = org49 if family:50 where.append("e.attributes->>'family' ilike :family")51 p["family"] = family52 if openness:53 vals = [v.strip() for v in openness.split(",") if v.strip()]54 if "open" in vals:55 vals += ["open-weights", "open-source"]56 where.append("e.attributes->>'openness' = any(cast(:openness as text[]))")57 p["openness"] = vals58 if modality:59 where.append("(e.attributes->'modalities' ? :modality or e.attributes->'modalities_input' ? :modality or e.attributes->'modalities_output' ? :modality)")60 p["modality"] = modality61 if status:62 where.append("e.status = any(cast(:status as text[]))")63 p["status"] = [v.strip() for v in status.split(",") if v.strip()]64 if min_params is not None:65 where.append(f"{PARAMS} >= :min_params")66 p["min_params"] = float(min_params)67 if max_params is not None:68 where.append(f"{PARAMS} <= :max_params")69 p["max_params"] = float(max_params)70 if min_context is not None:71 where.append(f"{CONTEXT} >= :min_context")72 p["min_context"] = float(min_context)73 if year_from is not None:74 where.append("left(e.attributes->>'release_date', 4) >= :yf")75 p["yf"] = str(year_from)76 if year_to is not None:77 where.append("left(e.attributes->>'release_date', 4) <= :yt")78 p["yt"] = str(year_to)79 if license:80 where.append("e.attributes->>'license' ilike :license")81 p["license"] = license82 return where, p838485async def model_facets(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:86 async def run(sql: str) -> list[dict[str, Any]]:87 async with connection() as conn:88 return await fetch_all(conn, sql, **params)8990 base = f"from {ENTITY_FROM} where {where_sql}"91 mods_from = (f"from {ENTITY_FROM} cross join lateral jsonb_array_elements_text(case when jsonb_typeof(e.attributes->'modalities') = 'array' "92 f"then e.attributes->'modalities' else '[]'::jsonb end) m where {where_sql}")93 orgs, openness, mods, fams, years, lics, status = await asyncio.gather(94 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"),95 run(f"select e.attributes->>'openness' as value, count(*) as count {base} and e.attributes ? 'openness' group by 1 order by 2 desc"),96 run(f"select m.value, count(*) as count {mods_from} group by 1 order by 2 desc limit 30"),97 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"),98 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"),99 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"),100 run(f"select e.status as value, count(*) as count {base} group by 1 order by 2 desc"),101 )102 conv = lambda rows: [{"value": r["value"], "count": int(r["count"])} for r in rows if r["value"] not in (None, "")]103 return {"organizations": [{"slug": r["slug"], "name": r["name"], "count": int(r["count"])} for r in orgs], "openness": conv(openness), "modalities": conv(mods),104 "families": conv(fams), "years": conv(years), "licenses": conv(lics), "status": conv(status)}105106107@router.get("")108@cached(300)109async def list_models(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, family: str | None = None,110 openness: str | None = None, modality: str | None = None, status: str | None = None, min_params: float | None = Query(None, ge=0),111 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),112 year_to: int | None = Query(None, ge=1950, le=2100), license: str | None = None, sort: str = "updated", order: str = "",113 facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]:114 if sort not in SORTS:115 raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")116 where, params = model_filters(q=q, org=org, family=family, openness=openness, modality=modality, status=status, min_params=min_params, max_params=max_params,117 min_context=min_context, year_from=year_from, year_to=year_to, license=license)118 order_sql = flip_order(SORTS[sort], order)119 where_sql = " and ".join(where)120 async with connection() as conn:121 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",122 lim=p.limit, off=p.offset, **params)123 total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params)124 out = page([entity_summary(r) for r in rows], int(total or 0), p)125 if facets:126 out["facets"] = await _facets_cached(where_sql, params)127 return out128129130async def _facets_cached(where_sql: str, params: dict[str, Any]) -> dict[str, Any]:131 key = "facets:models:" + where_sql + ":" + repr(sorted(params.items()))132 hit = await cache.cache_get(key)133 if hit is not None:134 return hit135 value = normalize(await model_facets(where_sql, params))136 await cache.cache_set(key, value, 600)137 return value138139140@router.get("/{slug}")141@cached(300)142async def get_model(request: Request, slug: str) -> dict[str, Any]:143 return await detail_for_type(slug, ("model",))144