1"""Shared API helpers: errors, pagination, date parsing, Redis response cache, admin auth, IP rate limits, serialisers.23Every route returns plain dicts; `AtlasJSONResponse` renders them with orjson (datetimes → ISO-8601 UTC, Decimal → float)."""4from __future__ import annotations56import hmac7import inspect8import time9from collections import defaultdict, deque10from collections.abc import Awaitable, Callable11from datetime import UTC, date, datetime12from datetime import time as dtime13from decimal import Decimal14from functools import wraps15from typing import Any1617import orjson18from fastapi import Depends, HTTPException, Query, Request19from fastapi.responses import JSONResponse20from sqlalchemy.ext.asyncio import AsyncConnection2122from aiatlas.config import settings23from aiatlas.db import fetch_one24from aiatlas.services import cache2526# ------------------------------------------------------------------------------------------------------------------ errors & JSON272829class ApiError(HTTPException):30 """Always `{"detail": "..."}` — never a stack trace."""3132 def __init__(self, status: int, detail: str):33 super().__init__(status_code=status, detail=detail)343536def _json_default(obj: Any) -> Any:37 if isinstance(obj, Decimal):38 return float(obj)39 if isinstance(obj, (set, frozenset)):40 return sorted(obj, key=str)41 if isinstance(obj, bytes):42 return obj.decode("utf-8", "replace")43 return str(obj)444546_ORJSON_OPTS = orjson.OPT_NON_STR_KEYS | orjson.OPT_UTC_Z | orjson.OPT_NAIVE_UTC474849def dumps(value: Any) -> bytes:50 return orjson.dumps(value, default=_json_default, option=_ORJSON_OPTS)515253def normalize(value: Any) -> Any:54 """Round-trip through orjson so cached and fresh responses are byte-identical (datetimes → strings, Decimal → float)."""55 return orjson.loads(dumps(value))565758class AtlasJSONResponse(JSONResponse):59 media_type = "application/json"6061 def render(self, content: Any) -> bytes:62 return dumps(content)636465# ------------------------------------------------------------------------------------------------------------------ pagination & parsing666768class Pagination:69 def __init__(self, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0)):70 self.limit = limit71 self.offset = offset727374PAGINATION = Depends() # `p: Pagination = PAGINATION` — FastAPI resolves the class from the annotation757677def page(items: list[Any], total: int, p: Pagination) -> dict[str, Any]:78 return {"items": items, "total": int(total), "limit": p.limit, "offset": p.offset}798081def parse_date(v: str | None, name: str = "date") -> date | None:82 """YYYY-MM-DD (or ISO datetime) → `date`. asyncpg needs real date objects — never cast text in SQL."""83 if v is None or v == "":84 return None85 s = v.strip().replace("Z", "+00:00")86 try:87 if len(s) == 10:88 return date.fromisoformat(s)89 return datetime.fromisoformat(s).astimezone(UTC).date()90 except ValueError as exc:91 raise ApiError(400, f"{name} must be YYYY-MM-DD, got {v!r}") from exc929394def parse_ts(v: str | None, name: str = "timestamp") -> datetime | None:95 """ISO-8601 (date or datetime) → aware UTC `datetime`."""96 if v is None or v == "":97 return None98 s = v.strip().replace("Z", "+00:00")99 try:100 dt = datetime.fromisoformat(s)101 except ValueError as exc:102 raise ApiError(400, f"{name} must be ISO-8601, got {v!r}") from exc103 if dt.tzinfo is None:104 dt = datetime.combine(dt.date(), dt.time() or dtime.min, UTC)105 return dt.astimezone(UTC)106107108def day_bounds(d: date) -> tuple[datetime, datetime]:109 start = datetime.combine(d, dtime.min, UTC)110 return start, datetime.combine(d, dtime.max, UTC).replace(microsecond=999999)111112113def csv(v: str | None) -> list[str]:114 return [x.strip() for x in (v or "").split(",") if x.strip()]115116117def num_expr(path: str) -> str:118 """Safe numeric cast of a JSON text attribute (`e.attributes->>'x'`)."""119 return f"(case when {path} ~ '^-?[0-9]+(\\.[0-9]+)?$' then ({path})::double precision end)"120121122def attr_num(key: str, alias: str = "e") -> str:123 return num_expr(f"{alias}.attributes->>'{key}'")124125126def flip_order(order_sql: str, order: str) -> str:127 """Apply `order=asc|desc` to the primary sort key of a canned ORDER BY fragment (keeps `nulls last`)."""128 if order not in ("asc", "desc"):129 return order_sql130 head, sep, tail = order_sql.partition(",")131 if order == "asc" and " desc" in head:132 head = head.replace(" desc", " asc", 1)133 elif order == "desc" and " asc" in head:134 head = head.replace(" asc", " desc", 1)135 return head + sep + tail136137138# ------------------------------------------------------------------------------------------------------------------ response cache139140141def cache_key(request: Request) -> str:142 items = sorted(request.query_params.multi_items())143 qs = "&".join(f"{k}={v}" for k, v in items)144 return f"{request.url.path}?{qs}" if qs else request.url.path145146147def cached(ttl_s: int) -> Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]:148 """Cache a GET route body in Redis (`aia:api:<path>?<sorted query>`). The route must accept `request: Request`."""149150 def deco(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]:151 @wraps(fn)152 async def wrapper(*args: Any, **kwargs: Any) -> Any:153 request = kwargs.get("request") or next((a for a in args if isinstance(a, Request)), None)154 key = cache_key(request) if request is not None else None155 if key is not None:156 hit = await cache.cache_get(key)157 if hit is not None:158 return hit159 value = normalize(await fn(*args, **kwargs))160 if key is not None:161 await cache.cache_set(key, value, ttl_s)162 return value163164 wrapper.__signature__ = inspect.signature(fn) # type: ignore[attr-defined]165 return wrapper166167 return deco168169170# ------------------------------------------------------------------------------------------------------------------ admin auth171172173def require_admin(request: Request) -> None:174 expected = settings.admin_token175 if not expected:176 raise ApiError(503, "admin API disabled: AIA_ADMIN_TOKEN is not configured")177 token = request.headers.get("x-aia-admin-token") or ""178 if not hmac.compare_digest(token, expected):179 raise ApiError(401, "invalid or missing x-aia-admin-token")180181182# ------------------------------------------------------------------------------------------------------------------ rate limiting183184_LIMITS: dict[str, tuple[int, float]] = {"search": (60, 60.0), "views": (1, 1.0), "admin": (240, 60.0), "default": (600, 60.0)}185_hits: dict[tuple[str, str], deque[float]] = defaultdict(deque)186187188def client_ip(request: Request) -> str:189 fwd = request.headers.get("x-forwarded-for")190 if fwd:191 return fwd.split(",")[0].strip()192 return request.client.host if request.client else "unknown"193194195def rate_limit(bucket: str) -> Callable[[Request], None]:196 def dep(request: Request) -> None:197 n, window = _LIMITS.get(bucket, _LIMITS["default"])198 now = time.monotonic()199 q = _hits[(client_ip(request), bucket)]200 while q and q[0] <= now - window:201 q.popleft()202 if len(q) >= n:203 raise ApiError(429, f"rate limit exceeded for {bucket}: {n} per {int(window)} s")204 q.append(now)205 if len(_hits) > 20000:206 for k in [k for k, v in _hits.items() if not v or v[-1] < now - 120]:207 _hits.pop(k, None)208209 return dep210211212# ------------------------------------------------------------------------------------------------------------------ entity serialisers213214ENTITY_FIELDS = ("id", "entity_type", "canonical_name", "slug", "description", "status", "organization_id", "attributes", "quality", "counts",215 "first_seen_at", "last_seen_at", "updated_at", "organization_name", "organization_slug")216217218def entity_cols(alias: str = "e", prefix: str = "") -> str:219 """Select-list for an entity alias (+ its organization alias `<alias>o`) with an optional column prefix."""220 org = f"{alias}o"221 cols = [f"{alias}.{c} as {prefix}{c}" for c in ENTITY_FIELDS[:13]]222 cols += [f"{org}.canonical_name as {prefix}organization_name", f"{org}.slug as {prefix}organization_slug"]223 return ", ".join(cols)224225226def entity_join(alias: str, on: str) -> str:227 return f"join entities {alias} on {alias}.id = {on} left join entities {alias}o on {alias}o.id = {alias}.organization_id"228229230ENTITY_COLS = entity_cols("e")231ENTITY_FROM = "entities e left join entities eo on eo.id = e.organization_id"232233SUMMARY_ATTRS: dict[str, tuple[str, ...]] = {234 "model": ("family", "openness", "license", "parameter_count", "active_parameter_count", "context_length", "max_output_tokens", "modalities",235 "release_date", "status", "knowledge_cutoff", "hf_repo", "api_model_id"),236 "company": ("country", "founded", "website", "org_kind"),237 "organization": ("country", "founded", "website", "org_kind"),238 "lab": ("country", "founded", "website", "org_kind"),239 "university": ("country", "founded", "website", "org_kind"),240 "provider": ("website", "pricing_url"),241 "benchmark": ("category", "metric", "unit"),242 "hardware": ("kind", "memory_gb", "memory_bandwidth_gbs", "release_date", "manufacturer"),243 "framework": ("latest_version", "latest_release_at", "license", "language", "metric.stars"),244 "repository": ("latest_version", "latest_release_at", "license", "language", "metric.stars"),245 "paper": ("authors", "published_at", "arxiv_id", "primary_category"),246 "dataset": ("license", "modality"),247}248COMPANY_TYPES = ("company", "organization", "lab", "university")249TYPE_LABELS = {"model": "Models", "company": "Companies", "organization": "Organizations", "lab": "Labs", "university": "Universities",250 "provider": "Providers", "paper": "Papers", "benchmark": "Benchmarks", "hardware": "Hardware", "framework": "Frameworks",251 "dataset": "Datasets", "tool": "Tools", "repository": "Repositories", "regulation": "Regulation", "incident": "Incidents",252 "researcher": "Researchers", "agent": "Agents", "product": "Products", "runtime": "Runtimes", "conference": "Conferences"}253254255def summary_attributes(entity_type: str, attrs: dict[str, Any] | None) -> dict[str, Any]:256 attrs = attrs or {}257 keys = SUMMARY_ATTRS.get(entity_type)258 if keys is None:259 return {k: v for k, v in attrs.items() if isinstance(v, (int, float, bool)) or (isinstance(v, str) and len(v) <= 200)}260 out: dict[str, Any] = {}261 for k in keys:262 if k in attrs and attrs[k] not in (None, "", [], {}):263 v = attrs[k]264 out[k] = v[:5] if k == "authors" and isinstance(v, list) else v265 return out266267268STATUS_VOCAB = ("active", "preview", "deprecated", "retired", "announced", "limited-availability", "unknown")269STATUS_MAP = {"available": "active", "ga": "active", "general-availability": "active", "released": "active", "live": "active", "beta": "preview",270 "alpha": "preview", "experimental": "preview", "coming-soon": "announced", "upcoming": "announced", "sunset": "retired", "discontinued": "retired",271 "archived": "retired", "legacy": "deprecated", "limited": "limited-availability", "merged": "retired"}272273274def normalize_status(value: Any) -> str:275 s = str(value or "").strip().lower().replace("_", "-").replace(" ", "-")276 s = STATUS_MAP.get(s, s)277 return s if s in STATUS_VOCAB else "unknown"278279280MAIN_ENTITY_TYPES = ("model", "company", "paper", "provider", "benchmark", "hardware", "framework", "dataset", "tool", "repository")281282EVENT_TYPE_LABELS: dict[str, tuple[str, int]] = { # event_type -> (human label, default importance 0–3)283 "NEW_MODEL": ("New model", 3), "MODEL_UPDATED": ("Model updated", 1), "PRICE_CHANGED": ("Price change", 2), "CONTEXT_CHANGED": ("Context window change", 2),284 "MAX_OUTPUT_CHANGED": ("Max output change", 1), "PARAMETERS_CHANGED": ("Parameter count change", 2), "KNOWLEDGE_CUTOFF_CHANGED": ("Knowledge cutoff change", 1),285 "RELEASE_DATE_CHANGED": ("Release date change", 1), "STATUS_CHANGED": ("Status change", 2), "LICENSE_CHANGED": ("License change", 2),286 "OPENNESS_CHANGED": ("Openness change", 2), "CAPABILITIES_CHANGED": ("Capabilities change", 1), "PROPERTY_CHANGED": ("Property change", 1),287 "DEPRECATION_ANNOUNCED": ("Deprecation announced", 2), "RETIREMENT_ANNOUNCED": ("Retirement announced", 2), "BENCHMARK_RESULT": ("Benchmark result", 1),288 "BENCHMARK_UPDATED": ("Benchmark updated", 1), "NEW_PAPER": ("New paper", 2), "NEW_COMPANY": ("New company", 2), "NEW_ORGANIZATION": ("New organization", 2),289 "NEW_LAB": ("New lab", 2), "NEW_UNIVERSITY": ("New university", 1), "NEW_PROVIDER": ("New provider", 2), "NEW_HARDWARE": ("New hardware", 2),290 "NEW_FRAMEWORK": ("New framework", 2), "NEW_LIBRARY": ("New library", 1), "NEW_BENCHMARK": ("New benchmark", 2), "NEW_DATASET": ("New dataset", 2),291 "NEW_REPOSITORY": ("New repository", 1), "NEW_TOOL": ("New tool", 1), "PROVIDER_LISTED": ("Listed by provider", 2), "PROVIDER_DELISTED": ("Delisted by provider", 2),292 "ANNOUNCEMENT": ("Announcement", 2), "RELEASE": ("Release", 2), "VERSION_RELEASED": ("Version released", 1), "DOCUMENT_CHANGED": ("Source document changed", 0),293 "ENTITY_MERGED": ("Duplicate merged", 1), "CLAIM_RETRACTED": ("Claim retracted", 1),294}295296297def event_type_label(event_type: str) -> str:298 known = EVENT_TYPE_LABELS.get(event_type)299 return known[0] if known else event_type.replace("_", " ").capitalize()300301302def event_type_importance(event_type: str) -> int:303 known = EVENT_TYPE_LABELS.get(event_type)304 return known[1] if known else 1305306307async def enrich_provenance(conn: AsyncConnection, *provenances: dict[str, Any] | None) -> None:308 """Add `source_name` (from `sources.name`) to every provenance entry that carries a `source_id` — in place."""309 ids = {v.get("source_id") for p in provenances if p for v in p.values() if isinstance(v, dict) and v.get("source_id")}310 if not ids:311 return312 from aiatlas.db import fetch_all313314 rows = await fetch_all(conn, "select id, name, domain, tier from sources where id = any(cast(:ids as text[]))", ids=sorted(ids))315 names = {r["id"]: r for r in rows}316 for p in provenances:317 for v in (p or {}).values():318 if isinstance(v, dict) and v.get("source_id") in names:319 src = names[v["source_id"]]320 v.setdefault("source_name", src["name"])321 v.setdefault("source_domain", src["domain"])322323324def org_of(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None:325 oid = row.get(f"{prefix}organization_id")326 if not oid:327 return None328 return {"id": oid, "slug": row.get(f"{prefix}organization_slug"), "name": row.get(f"{prefix}organization_name")}329330331def entity_summary(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None:332 g = row.get333 if not g(f"{prefix}id"):334 return None335 etype = g(f"{prefix}entity_type") or ""336 desc = g(f"{prefix}description")337 raw_status = g(f"{prefix}status")338 status = normalize_status(raw_status)339 return {"id": g(f"{prefix}id"), "entity_type": etype, "slug": g(f"{prefix}slug"), "name": g(f"{prefix}canonical_name"),340 "description": (desc[:280] if isinstance(desc, str) and len(desc) > 280 else desc), "status": status,341 **({"status_raw": raw_status} if raw_status and raw_status != status else {}),342 "organization": org_of(row, prefix), "attributes": summary_attributes(etype, g(f"{prefix}attributes")),343 "quality": g(f"{prefix}quality") or {}, "counts": g(f"{prefix}counts") or {},344 "first_seen_at": g(f"{prefix}first_seen_at"), "last_seen_at": g(f"{prefix}last_seen_at"), "updated_at": g(f"{prefix}updated_at")}345346347def change_event(row: dict[str, Any], prefix: str = "e_") -> dict[str, Any]:348 return {"id": row["id"], "event_type": row["event_type"], "category": row["category"], "property": row.get("property"),349 "old_value": row.get("old_value"), "new_value": row.get("new_value"), "summary": row.get("summary"), "importance": row.get("importance"),350 "observed_at": row.get("observed_at"), "effective_at": row.get("effective_at"), "source_url": row.get("source_url"),351 "connector_name": row.get("connector_name"), "entity": entity_summary(row, prefix), "meta": row.get("meta") or {}}352353354EVENT_COLS = ("ev.id, ev.event_type, ev.category, ev.property, ev.old_value, ev.new_value, ev.summary, ev.importance, ev.observed_at, ev.effective_at, "355 "ev.source_url, ev.connector_name, ev.meta, " + entity_cols("e", "e_"))356EVENT_FROM = "change_events ev left join entities e on e.id = ev.entity_id left join entities eo on eo.id = e.organization_id"357358359def price_row(row: dict[str, Any]) -> dict[str, Any]:360 return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"),361 "input_per_mtok": row.get("input_per_mtok"), "output_per_mtok": row.get("output_per_mtok"), "cached_input_per_mtok": row.get("cached_input_per_mtok"),362 "cache_write_per_mtok": row.get("cache_write_per_mtok"), "batch_input_per_mtok": row.get("batch_input_per_mtok"),363 "batch_output_per_mtok": row.get("batch_output_per_mtok"), "per_image": row.get("per_image"), "per_request": row.get("per_request"),364 "currency": row.get("currency"), "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"),365 "features": row.get("features") or {}, "observed_at": row.get("observed_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"),366 "source_url": row.get("source_url"), "tier": row.get("tier")}367368369PRICE_COLS = ("p.id, p.provider_model_id, p.input_per_mtok, p.output_per_mtok, p.cached_input_per_mtok, p.cache_write_per_mtok, p.batch_input_per_mtok, "370 "p.batch_output_per_mtok, p.per_image, p.per_request, p.currency, p.context_length, p.max_output_tokens, p.features, p.observed_at, p.valid_from, "371 "p.valid_to, p.source_url, p.tier, " + entity_cols("m", "m_") + ", " + entity_cols("pv", "p_"))372PRICE_FROM = "prices p " + entity_join("m", "p.model_id") + " " + entity_join("pv", "p.provider_id")373374375def result_row(row: dict[str, Any]) -> dict[str, Any]:376 return {"id": row["id"], "model": entity_summary(row, "m_"), "benchmark": entity_summary(row, "b_"), "score": row.get("score"), "metric": row.get("metric"),377 "unit": row.get("unit"), "higher_is_better": row.get("higher_is_better"), "config": row.get("config") or {}, "evaluated_at": row.get("evaluated_at"),378 "observed_at": row.get("observed_at"), "source_url": row.get("source_url"), "tier": row.get("tier"), "confidence": row.get("confidence"),379 "valid_to": row.get("valid_to")}380381382RESULT_COLS = ("r.id, r.score, r.metric, r.unit, r.higher_is_better, r.config, r.evaluated_at, r.observed_at, r.source_url, r.tier, r.confidence, r.valid_to, "383 + entity_cols("m", "m_") + ", " + entity_cols("b", "b_"))384RESULT_FROM = "benchmark_results r " + entity_join("m", "r.model_id") + " " + entity_join("b", "r.benchmark_id")385RESULT_ORDER = "case when r.higher_is_better then -r.score else r.score end, r.observed_at desc"386387388def claim_row(row: dict[str, Any]) -> dict[str, Any]:389 return {"id": row["id"], "property": row["property"], "value": row.get("value"), "unit": row.get("unit"), "tier": row.get("tier"),390 "confidence": row.get("confidence"), "status": row.get("status"), "extractor": row.get("extractor"), "observed_at": row.get("observed_at"),391 "effective_at": row.get("effective_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"),392 "source_name": row.get("source_name")}393394395CLAIM_COLS = ("c.id, c.property, c.value, c.unit, c.tier, c.confidence, c.status, c.extractor, c.observed_at, c.effective_at, c.valid_from, c.valid_to, "396 "c.source_url, s.name as source_name")397CLAIM_FROM = "claims c left join sources s on s.id = c.source_id"398399400# ------------------------------------------------------------------------------------------------------------------ entity resolution401402403async def resolve_entity(conn: AsyncConnection, slug_or_id: str, types: tuple[str, ...] | None = None) -> dict[str, Any]:404 """Slug first, then id; follows `merged_into`; 404 when absent or when the type does not match the mounted alias."""405 key = slug_or_id.strip()406 if not key or len(key) > 200:407 raise ApiError(404, "entity not found")408 row = await fetch_one(conn, f"select {ENTITY_COLS}, e.provenance, e.merged_into from {ENTITY_FROM} where e.slug = :k or e.id = :k "409 f"order by case when e.slug = :k then 0 else 1 end limit 1", k=key)410 hops = 0411 while row and row.get("merged_into") and hops < 5:412 row = await fetch_one(conn, f"select {ENTITY_COLS}, e.provenance, e.merged_into from {ENTITY_FROM} where e.id = :k", k=row["merged_into"])413 hops += 1414 if not row:415 raise ApiError(404, "entity not found")416 if types and row["entity_type"] not in types:417 raise ApiError(404, f"entity {key!r} is a {row['entity_type']}, not one of {', '.join(types)}")418 return row419420421async def resolve_id(conn: AsyncConnection, slug_or_id: str | None, types: tuple[str, ...] | None = None) -> str | None:422 if not slug_or_id:423 return None424 return (await resolve_entity(conn, slug_or_id, types))["id"]425426427__all__ = [428 "CLAIM_COLS",429 "CLAIM_FROM",430 "COMPANY_TYPES",431 "ENTITY_COLS",432 "ENTITY_FROM",433 "EVENT_COLS",434 "EVENT_FROM",435 "EVENT_TYPE_LABELS",436 "MAIN_ENTITY_TYPES",437 "PAGINATION",438 "PRICE_COLS",439 "PRICE_FROM",440 "RESULT_COLS",441 "RESULT_FROM",442 "RESULT_ORDER",443 "STATUS_VOCAB",444 "TYPE_LABELS",445 "ApiError",446 "AtlasJSONResponse",447 "Pagination",448 "attr_num",449 "cache_key",450 "cached",451 "change_event",452 "claim_row",453 "client_ip",454 "csv",455 "day_bounds",456 "dumps",457 "enrich_provenance",458 "entity_cols",459 "entity_join",460 "entity_summary",461 "event_type_importance",462 "event_type_label",463 "flip_order",464 "normalize",465 "normalize_status",466 "num_expr",467 "org_of",468 "page",469 "parse_date",470 "parse_ts",471 "price_row",472 "rate_limit",473 "require_admin",474 "resolve_entity",475 "resolve_id",476 "result_row",477]478