"""Shared API helpers: errors, pagination, date parsing, Redis response cache, admin auth, IP rate limits, serialisers. Every route returns plain dicts; `AtlasJSONResponse` renders them with orjson (datetimes → ISO-8601 UTC, Decimal → float).""" from __future__ import annotations import hmac import inspect import time from collections import defaultdict, deque from collections.abc import Awaitable, Callable from datetime import UTC, date, datetime from datetime import time as dtime from decimal import Decimal from functools import wraps from typing import Any import orjson from fastapi import Depends, HTTPException, Query, Request from fastapi.responses import JSONResponse from sqlalchemy.ext.asyncio import AsyncConnection from aiatlas.config import settings from aiatlas.db import fetch_one from aiatlas.services import cache # ------------------------------------------------------------------------------------------------------------------ errors & JSON class ApiError(HTTPException): """Always `{"detail": "..."}` — never a stack trace.""" def __init__(self, status: int, detail: str): super().__init__(status_code=status, detail=detail) def _json_default(obj: Any) -> Any: if isinstance(obj, Decimal): return float(obj) if isinstance(obj, (set, frozenset)): return sorted(obj, key=str) if isinstance(obj, bytes): return obj.decode("utf-8", "replace") return str(obj) _ORJSON_OPTS = orjson.OPT_NON_STR_KEYS | orjson.OPT_UTC_Z | orjson.OPT_NAIVE_UTC def dumps(value: Any) -> bytes: return orjson.dumps(value, default=_json_default, option=_ORJSON_OPTS) def normalize(value: Any) -> Any: """Round-trip through orjson so cached and fresh responses are byte-identical (datetimes → strings, Decimal → float).""" return orjson.loads(dumps(value)) class AtlasJSONResponse(JSONResponse): media_type = "application/json" def render(self, content: Any) -> bytes: return dumps(content) # ------------------------------------------------------------------------------------------------------------------ pagination & parsing class Pagination: def __init__(self, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0)): self.limit = limit self.offset = offset PAGINATION = Depends() # `p: Pagination = PAGINATION` — FastAPI resolves the class from the annotation def page(items: list[Any], total: int, p: Pagination) -> dict[str, Any]: return {"items": items, "total": int(total), "limit": p.limit, "offset": p.offset} def parse_date(v: str | None, name: str = "date") -> date | None: """YYYY-MM-DD (or ISO datetime) → `date`. asyncpg needs real date objects — never cast text in SQL.""" if v is None or v == "": return None s = v.strip().replace("Z", "+00:00") try: if len(s) == 10: return date.fromisoformat(s) return datetime.fromisoformat(s).astimezone(UTC).date() except ValueError as exc: raise ApiError(400, f"{name} must be YYYY-MM-DD, got {v!r}") from exc def parse_ts(v: str | None, name: str = "timestamp") -> datetime | None: """ISO-8601 (date or datetime) → aware UTC `datetime`.""" if v is None or v == "": return None s = v.strip().replace("Z", "+00:00") try: dt = datetime.fromisoformat(s) except ValueError as exc: raise ApiError(400, f"{name} must be ISO-8601, got {v!r}") from exc if dt.tzinfo is None: dt = datetime.combine(dt.date(), dt.time() or dtime.min, UTC) return dt.astimezone(UTC) def day_bounds(d: date) -> tuple[datetime, datetime]: start = datetime.combine(d, dtime.min, UTC) return start, datetime.combine(d, dtime.max, UTC).replace(microsecond=999999) def csv(v: str | None) -> list[str]: return [x.strip() for x in (v or "").split(",") if x.strip()] def num_expr(path: str) -> str: """Safe numeric cast of a JSON text attribute (`e.attributes->>'x'`).""" return f"(case when {path} ~ '^-?[0-9]+(\\.[0-9]+)?$' then ({path})::double precision end)" def attr_num(key: str, alias: str = "e") -> str: return num_expr(f"{alias}.attributes->>'{key}'") def flip_order(order_sql: str, order: str) -> str: """Apply `order=asc|desc` to the primary sort key of a canned ORDER BY fragment (keeps `nulls last`).""" if order not in ("asc", "desc"): return order_sql head, sep, tail = order_sql.partition(",") if order == "asc" and " desc" in head: head = head.replace(" desc", " asc", 1) elif order == "desc" and " asc" in head: head = head.replace(" asc", " desc", 1) return head + sep + tail # ------------------------------------------------------------------------------------------------------------------ response cache def cache_key(request: Request) -> str: items = sorted(request.query_params.multi_items()) qs = "&".join(f"{k}={v}" for k, v in items) return f"{request.url.path}?{qs}" if qs else request.url.path def cached(ttl_s: int) -> Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]: """Cache a GET route body in Redis (`aia:api:?`). The route must accept `request: Request`.""" def deco(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]: @wraps(fn) async def wrapper(*args: Any, **kwargs: Any) -> Any: request = kwargs.get("request") or next((a for a in args if isinstance(a, Request)), None) key = cache_key(request) if request is not None else None if key is not None: hit = await cache.cache_get(key) if hit is not None: return hit value = normalize(await fn(*args, **kwargs)) if key is not None: await cache.cache_set(key, value, ttl_s) return value wrapper.__signature__ = inspect.signature(fn) # type: ignore[attr-defined] return wrapper return deco # ------------------------------------------------------------------------------------------------------------------ admin auth def require_admin(request: Request) -> None: expected = settings.admin_token if not expected: raise ApiError(503, "admin API disabled: AIA_ADMIN_TOKEN is not configured") token = request.headers.get("x-aia-admin-token") or "" if not hmac.compare_digest(token, expected): raise ApiError(401, "invalid or missing x-aia-admin-token") # ------------------------------------------------------------------------------------------------------------------ rate limiting _LIMITS: dict[str, tuple[int, float]] = {"search": (60, 60.0), "views": (1, 1.0), "admin": (240, 60.0), "default": (600, 60.0)} _hits: dict[tuple[str, str], deque[float]] = defaultdict(deque) def client_ip(request: Request) -> str: fwd = request.headers.get("x-forwarded-for") if fwd: return fwd.split(",")[0].strip() return request.client.host if request.client else "unknown" def rate_limit(bucket: str) -> Callable[[Request], None]: def dep(request: Request) -> None: n, window = _LIMITS.get(bucket, _LIMITS["default"]) now = time.monotonic() q = _hits[(client_ip(request), bucket)] while q and q[0] <= now - window: q.popleft() if len(q) >= n: raise ApiError(429, f"rate limit exceeded for {bucket}: {n} per {int(window)} s") q.append(now) if len(_hits) > 20000: for k in [k for k, v in _hits.items() if not v or v[-1] < now - 120]: _hits.pop(k, None) return dep # ------------------------------------------------------------------------------------------------------------------ entity serialisers ENTITY_FIELDS = ("id", "entity_type", "canonical_name", "slug", "description", "status", "organization_id", "attributes", "quality", "counts", "first_seen_at", "last_seen_at", "updated_at", "organization_name", "organization_slug") def entity_cols(alias: str = "e", prefix: str = "") -> str: """Select-list for an entity alias (+ its organization alias `o`) with an optional column prefix.""" org = f"{alias}o" cols = [f"{alias}.{c} as {prefix}{c}" for c in ENTITY_FIELDS[:13]] cols += [f"{org}.canonical_name as {prefix}organization_name", f"{org}.slug as {prefix}organization_slug"] return ", ".join(cols) def entity_join(alias: str, on: str) -> str: return f"join entities {alias} on {alias}.id = {on} left join entities {alias}o on {alias}o.id = {alias}.organization_id" ENTITY_COLS = entity_cols("e") ENTITY_FROM = "entities e left join entities eo on eo.id = e.organization_id" SUMMARY_ATTRS: dict[str, tuple[str, ...]] = { "model": ("family", "openness", "license", "parameter_count", "active_parameter_count", "context_length", "max_output_tokens", "modalities", "release_date", "status", "knowledge_cutoff", "hf_repo", "api_model_id"), "company": ("country", "founded", "website", "org_kind"), "organization": ("country", "founded", "website", "org_kind"), "lab": ("country", "founded", "website", "org_kind"), "university": ("country", "founded", "website", "org_kind"), "provider": ("website", "pricing_url"), "benchmark": ("category", "metric", "unit"), "hardware": ("kind", "memory_gb", "memory_bandwidth_gbs", "release_date", "manufacturer"), "framework": ("latest_version", "latest_release_at", "license", "language", "metric.stars"), "repository": ("latest_version", "latest_release_at", "license", "language", "metric.stars"), "paper": ("authors", "published_at", "arxiv_id", "primary_category"), "dataset": ("license", "modality"), } COMPANY_TYPES = ("company", "organization", "lab", "university") TYPE_LABELS = {"model": "Models", "company": "Companies", "organization": "Organizations", "lab": "Labs", "university": "Universities", "provider": "Providers", "paper": "Papers", "benchmark": "Benchmarks", "hardware": "Hardware", "framework": "Frameworks", "dataset": "Datasets", "tool": "Tools", "repository": "Repositories", "regulation": "Regulation", "incident": "Incidents", "researcher": "Researchers", "agent": "Agents", "product": "Products", "runtime": "Runtimes", "conference": "Conferences"} def summary_attributes(entity_type: str, attrs: dict[str, Any] | None) -> dict[str, Any]: attrs = attrs or {} keys = SUMMARY_ATTRS.get(entity_type) if keys is None: return {k: v for k, v in attrs.items() if isinstance(v, (int, float, bool)) or (isinstance(v, str) and len(v) <= 200)} out: dict[str, Any] = {} for k in keys: if k in attrs and attrs[k] not in (None, "", [], {}): v = attrs[k] out[k] = v[:5] if k == "authors" and isinstance(v, list) else v return out STATUS_VOCAB = ("active", "preview", "deprecated", "retired", "announced", "limited-availability", "unknown") STATUS_MAP = {"available": "active", "ga": "active", "general-availability": "active", "released": "active", "live": "active", "beta": "preview", "alpha": "preview", "experimental": "preview", "coming-soon": "announced", "upcoming": "announced", "sunset": "retired", "discontinued": "retired", "archived": "retired", "legacy": "deprecated", "limited": "limited-availability", "merged": "retired"} def normalize_status(value: Any) -> str: s = str(value or "").strip().lower().replace("_", "-").replace(" ", "-") s = STATUS_MAP.get(s, s) return s if s in STATUS_VOCAB else "unknown" MAIN_ENTITY_TYPES = ("model", "company", "paper", "provider", "benchmark", "hardware", "framework", "dataset", "tool", "repository") EVENT_TYPE_LABELS: dict[str, tuple[str, int]] = { # event_type -> (human label, default importance 0–3) "NEW_MODEL": ("New model", 3), "MODEL_UPDATED": ("Model updated", 1), "PRICE_CHANGED": ("Price change", 2), "CONTEXT_CHANGED": ("Context window change", 2), "MAX_OUTPUT_CHANGED": ("Max output change", 1), "PARAMETERS_CHANGED": ("Parameter count change", 2), "KNOWLEDGE_CUTOFF_CHANGED": ("Knowledge cutoff change", 1), "RELEASE_DATE_CHANGED": ("Release date change", 1), "STATUS_CHANGED": ("Status change", 2), "LICENSE_CHANGED": ("License change", 2), "OPENNESS_CHANGED": ("Openness change", 2), "CAPABILITIES_CHANGED": ("Capabilities change", 1), "PROPERTY_CHANGED": ("Property change", 1), "DEPRECATION_ANNOUNCED": ("Deprecation announced", 2), "RETIREMENT_ANNOUNCED": ("Retirement announced", 2), "BENCHMARK_RESULT": ("Benchmark result", 1), "BENCHMARK_UPDATED": ("Benchmark updated", 1), "NEW_PAPER": ("New paper", 2), "NEW_COMPANY": ("New company", 2), "NEW_ORGANIZATION": ("New organization", 2), "NEW_LAB": ("New lab", 2), "NEW_UNIVERSITY": ("New university", 1), "NEW_PROVIDER": ("New provider", 2), "NEW_HARDWARE": ("New hardware", 2), "NEW_FRAMEWORK": ("New framework", 2), "NEW_LIBRARY": ("New library", 1), "NEW_BENCHMARK": ("New benchmark", 2), "NEW_DATASET": ("New dataset", 2), "NEW_REPOSITORY": ("New repository", 1), "NEW_TOOL": ("New tool", 1), "PROVIDER_LISTED": ("Listed by provider", 2), "PROVIDER_DELISTED": ("Delisted by provider", 2), "ANNOUNCEMENT": ("Announcement", 2), "RELEASE": ("Release", 2), "VERSION_RELEASED": ("Version released", 1), "DOCUMENT_CHANGED": ("Source document changed", 0), "ENTITY_MERGED": ("Duplicate merged", 1), "CLAIM_RETRACTED": ("Claim retracted", 1), } def event_type_label(event_type: str) -> str: known = EVENT_TYPE_LABELS.get(event_type) return known[0] if known else event_type.replace("_", " ").capitalize() def event_type_importance(event_type: str) -> int: known = EVENT_TYPE_LABELS.get(event_type) return known[1] if known else 1 async def enrich_provenance(conn: AsyncConnection, *provenances: dict[str, Any] | None) -> None: """Add `source_name` (from `sources.name`) to every provenance entry that carries a `source_id` — in place.""" 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")} if not ids: return from aiatlas.db import fetch_all rows = await fetch_all(conn, "select id, name, domain, tier from sources where id = any(cast(:ids as text[]))", ids=sorted(ids)) names = {r["id"]: r for r in rows} for p in provenances: for v in (p or {}).values(): if isinstance(v, dict) and v.get("source_id") in names: src = names[v["source_id"]] v.setdefault("source_name", src["name"]) v.setdefault("source_domain", src["domain"]) def org_of(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None: oid = row.get(f"{prefix}organization_id") if not oid: return None return {"id": oid, "slug": row.get(f"{prefix}organization_slug"), "name": row.get(f"{prefix}organization_name")} def entity_summary(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None: g = row.get if not g(f"{prefix}id"): return None etype = g(f"{prefix}entity_type") or "" desc = g(f"{prefix}description") raw_status = g(f"{prefix}status") status = normalize_status(raw_status) return {"id": g(f"{prefix}id"), "entity_type": etype, "slug": g(f"{prefix}slug"), "name": g(f"{prefix}canonical_name"), "description": (desc[:280] if isinstance(desc, str) and len(desc) > 280 else desc), "status": status, **({"status_raw": raw_status} if raw_status and raw_status != status else {}), "organization": org_of(row, prefix), "attributes": summary_attributes(etype, g(f"{prefix}attributes")), "quality": g(f"{prefix}quality") or {}, "counts": g(f"{prefix}counts") or {}, "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")} def change_event(row: dict[str, Any], prefix: str = "e_") -> dict[str, Any]: return {"id": row["id"], "event_type": row["event_type"], "category": row["category"], "property": row.get("property"), "old_value": row.get("old_value"), "new_value": row.get("new_value"), "summary": row.get("summary"), "importance": row.get("importance"), "observed_at": row.get("observed_at"), "effective_at": row.get("effective_at"), "source_url": row.get("source_url"), "connector_name": row.get("connector_name"), "entity": entity_summary(row, prefix), "meta": row.get("meta") or {}} EVENT_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, " "ev.source_url, ev.connector_name, ev.meta, " + entity_cols("e", "e_")) EVENT_FROM = "change_events ev left join entities e on e.id = ev.entity_id left join entities eo on eo.id = e.organization_id" def price_row(row: dict[str, Any]) -> dict[str, Any]: return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"), "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"), "cache_write_per_mtok": row.get("cache_write_per_mtok"), "batch_input_per_mtok": row.get("batch_input_per_mtok"), "batch_output_per_mtok": row.get("batch_output_per_mtok"), "per_image": row.get("per_image"), "per_request": row.get("per_request"), "currency": row.get("currency"), "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"), "features": row.get("features") or {}, "observed_at": row.get("observed_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"), "tier": row.get("tier")} PRICE_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, " "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, " "p.valid_to, p.source_url, p.tier, " + entity_cols("m", "m_") + ", " + entity_cols("pv", "p_")) PRICE_FROM = "prices p " + entity_join("m", "p.model_id") + " " + entity_join("pv", "p.provider_id") def result_row(row: dict[str, Any]) -> dict[str, Any]: return {"id": row["id"], "model": entity_summary(row, "m_"), "benchmark": entity_summary(row, "b_"), "score": row.get("score"), "metric": row.get("metric"), "unit": row.get("unit"), "higher_is_better": row.get("higher_is_better"), "config": row.get("config") or {}, "evaluated_at": row.get("evaluated_at"), "observed_at": row.get("observed_at"), "source_url": row.get("source_url"), "tier": row.get("tier"), "confidence": row.get("confidence"), "valid_to": row.get("valid_to")} RESULT_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, " + entity_cols("m", "m_") + ", " + entity_cols("b", "b_")) RESULT_FROM = "benchmark_results r " + entity_join("m", "r.model_id") + " " + entity_join("b", "r.benchmark_id") RESULT_ORDER = "case when r.higher_is_better then -r.score else r.score end, r.observed_at desc" def claim_row(row: dict[str, Any]) -> dict[str, Any]: return {"id": row["id"], "property": row["property"], "value": row.get("value"), "unit": row.get("unit"), "tier": row.get("tier"), "confidence": row.get("confidence"), "status": row.get("status"), "extractor": row.get("extractor"), "observed_at": row.get("observed_at"), "effective_at": row.get("effective_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"), "source_name": row.get("source_name")} CLAIM_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, " "c.source_url, s.name as source_name") CLAIM_FROM = "claims c left join sources s on s.id = c.source_id" # ------------------------------------------------------------------------------------------------------------------ entity resolution async def resolve_entity(conn: AsyncConnection, slug_or_id: str, types: tuple[str, ...] | None = None) -> dict[str, Any]: """Slug first, then id; follows `merged_into`; 404 when absent or when the type does not match the mounted alias.""" key = slug_or_id.strip() if not key or len(key) > 200: raise ApiError(404, "entity not found") 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 " f"order by case when e.slug = :k then 0 else 1 end limit 1", k=key) hops = 0 while row and row.get("merged_into") and hops < 5: 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"]) hops += 1 if not row: raise ApiError(404, "entity not found") if types and row["entity_type"] not in types: raise ApiError(404, f"entity {key!r} is a {row['entity_type']}, not one of {', '.join(types)}") return row async def resolve_id(conn: AsyncConnection, slug_or_id: str | None, types: tuple[str, ...] | None = None) -> str | None: if not slug_or_id: return None return (await resolve_entity(conn, slug_or_id, types))["id"] __all__ = [ "CLAIM_COLS", "CLAIM_FROM", "COMPANY_TYPES", "ENTITY_COLS", "ENTITY_FROM", "EVENT_COLS", "EVENT_FROM", "EVENT_TYPE_LABELS", "MAIN_ENTITY_TYPES", "PAGINATION", "PRICE_COLS", "PRICE_FROM", "RESULT_COLS", "RESULT_FROM", "RESULT_ORDER", "STATUS_VOCAB", "TYPE_LABELS", "ApiError", "AtlasJSONResponse", "Pagination", "attr_num", "cache_key", "cached", "change_event", "claim_row", "client_ip", "csv", "day_bounds", "dumps", "enrich_provenance", "entity_cols", "entity_join", "entity_summary", "event_type_importance", "event_type_label", "flip_order", "normalize", "normalize_status", "num_expr", "org_of", "page", "parse_date", "parse_ts", "price_row", "rate_limit", "require_admin", "resolve_entity", "resolve_id", "result_row", ]