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 logging9import os10import time11from collections import defaultdict, deque12from collections.abc import Awaitable, Callable13from datetime import UTC, date, datetime14from datetime import time as dtime15from decimal import Decimal16from functools import wraps17from typing import Any1819import orjson20from fastapi import Depends, HTTPException, Query, Request21from fastapi.responses import JSONResponse22from sqlalchemy.ext.asyncio import AsyncConnection2324from aiatlas.config import settings25from aiatlas.db import execute, fetch_one, jsonb, transaction26from aiatlas.services import cache2728log = logging.getLogger("aiatlas.api")2930# ------------------------------------------------------------------------------------------------------------------ errors & JSON313233class ApiError(HTTPException):34 """Always `{"detail": "..."}` — never a stack trace."""3536 def __init__(self, status: int, detail: str):37 super().__init__(status_code=status, detail=detail)383940def _json_default(obj: Any) -> Any:41 if isinstance(obj, Decimal):42 return float(obj)43 if isinstance(obj, (set, frozenset)):44 return sorted(obj, key=str)45 if isinstance(obj, bytes):46 return obj.decode("utf-8", "replace")47 return str(obj)484950_ORJSON_OPTS = orjson.OPT_NON_STR_KEYS | orjson.OPT_UTC_Z | orjson.OPT_NAIVE_UTC515253def dumps(value: Any) -> bytes:54 return orjson.dumps(value, default=_json_default, option=_ORJSON_OPTS)555657def normalize(value: Any) -> Any:58 """Round-trip through orjson so cached and fresh responses are byte-identical (datetimes → strings, Decimal → float)."""59 return orjson.loads(dumps(value))606162class AtlasJSONResponse(JSONResponse):63 media_type = "application/json"6465 def render(self, content: Any) -> bytes:66 return dumps(content)676869# ------------------------------------------------------------------------------------------------------------------ pagination & parsing707172class Pagination:73 def __init__(self, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0)):74 self.limit = limit75 self.offset = offset767778PAGINATION = Depends() # `p: Pagination = PAGINATION` — FastAPI resolves the class from the annotation798081def page(items: list[Any], total: int, p: Pagination) -> dict[str, Any]:82 return {"items": items, "total": int(total), "limit": p.limit, "offset": p.offset}838485def parse_date(v: str | None, name: str = "date") -> date | None:86 """YYYY-MM-DD (or ISO datetime) → `date`. asyncpg needs real date objects — never cast text in SQL."""87 if v is None or v == "":88 return None89 s = v.strip().replace("Z", "+00:00")90 try:91 if len(s) == 10:92 return date.fromisoformat(s)93 return datetime.fromisoformat(s).astimezone(UTC).date()94 except ValueError as exc:95 raise ApiError(400, f"{name} must be YYYY-MM-DD, got {v!r}") from exc969798def parse_ts(v: str | None, name: str = "timestamp") -> datetime | None:99 """ISO-8601 (date or datetime) → aware UTC `datetime`."""100 if v is None or v == "":101 return None102 s = v.strip().replace("Z", "+00:00")103 try:104 dt = datetime.fromisoformat(s)105 except ValueError as exc:106 raise ApiError(400, f"{name} must be ISO-8601, got {v!r}") from exc107 if dt.tzinfo is None:108 dt = datetime.combine(dt.date(), dt.time() or dtime.min, UTC)109 return dt.astimezone(UTC)110111112def day_bounds(d: date) -> tuple[datetime, datetime]:113 start = datetime.combine(d, dtime.min, UTC)114 return start, datetime.combine(d, dtime.max, UTC).replace(microsecond=999999)115116117def csv(v: str | None) -> list[str]:118 return [x.strip() for x in (v or "").split(",") if x.strip()]119120121def num_expr(path: str) -> str:122 """Safe numeric cast of a JSON text attribute (`e.attributes->>'x'`)."""123 return f"(case when {path} ~ '^-?[0-9]+(\\.[0-9]+)?$' then ({path})::double precision end)"124125126def attr_num(key: str, alias: str = "e") -> str:127 return num_expr(f"{alias}.attributes->>'{key}'")128129130def flip_order(order_sql: str, order: str) -> str:131 """Apply `order=asc|desc` to the primary sort key of a canned ORDER BY fragment (keeps `nulls last`)."""132 if order not in ("asc", "desc"):133 return order_sql134 head, sep, tail = order_sql.partition(",")135 if order == "asc" and " desc" in head:136 head = head.replace(" desc", " asc", 1)137 elif order == "desc" and " asc" in head:138 head = head.replace(" asc", " desc", 1)139 return head + sep + tail140141142# ------------------------------------------------------------------------------------------------------------------ response cache143144145def cache_key(request: Request) -> str:146 items = sorted(request.query_params.multi_items())147 qs = "&".join(f"{k}={v}" for k, v in items)148 return f"{request.url.path}?{qs}" if qs else request.url.path149150151def cached(ttl_s: int) -> Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]:152 """Cache a GET route body in Redis (`aia:api:<path>?<sorted query>`). The route must accept `request: Request`."""153154 def deco(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]:155 @wraps(fn)156 async def wrapper(*args: Any, **kwargs: Any) -> Any:157 request = kwargs.get("request") or next((a for a in args if isinstance(a, Request)), None)158 key = cache_key(request) if request is not None else None159 if key is not None:160 hit = await cache.cache_get(key)161 if hit is not None:162 return hit163 value = normalize(await fn(*args, **kwargs))164 if key is not None:165 await cache.cache_set(key, value, ttl_s)166 return value167168 wrapper.__signature__ = inspect.signature(fn) # type: ignore[attr-defined]169 return wrapper170171 return deco172173174# ------------------------------------------------------------------------------------------------------------------ rate limiting175176_LIMITS: dict[str, tuple[int, float]] = {"search": (60, 60.0), "views": (1, 1.0), "admin": (240, 60.0), "admin-auth-failed": (10, 60.0), "default": (600, 60.0)}177_hits: dict[tuple[str, str], deque[float]] = defaultdict(deque)178179180def trust_proxy() -> bool:181 """`x-forwarded-for` is only honoured behind a trusted reverse proxy (`settings.trust_proxy` or env `AIA_TRUST_PROXY=1`)."""182 v = getattr(settings, "trust_proxy", None)183 if v is None:184 v = os.environ.get("AIA_TRUST_PROXY", "")185 return str(v).strip().lower() in ("1", "true", "yes", "on")186187188def client_ip(request: Request) -> str:189 if trust_proxy():190 fwd = request.headers.get("x-forwarded-for")191 if fwd:192 return fwd.split(",")[0].strip()193 return request.client.host if request.client else "unknown"194195196def _bucket_hit(ip: str, bucket: str, *, record: bool = True) -> bool:197 """Sliding-window counter. Returns True when the call is allowed (and records it), False when the bucket is full."""198 n, window = _LIMITS.get(bucket, _LIMITS["default"])199 now = time.monotonic()200 q = _hits[(ip, bucket)]201 while q and q[0] <= now - window:202 q.popleft()203 if len(q) >= n:204 return False205 if record:206 q.append(now)207 if len(_hits) > 20000:208 for k in [k for k, v in _hits.items() if not v or v[-1] < now - 120]:209 _hits.pop(k, None)210 return True211212213def rate_limit(bucket: str) -> Callable[[Request], None]:214 def dep(request: Request) -> None:215 n, window = _LIMITS.get(bucket, _LIMITS["default"])216 if not _bucket_hit(client_ip(request), bucket):217 raise ApiError(429, f"rate limit exceeded for {bucket}: {n} per {int(window)} s")218219 return dep220221222# ------------------------------------------------------------------------------------------------------------------ admin auth223224225def require_admin(request: Request) -> None:226 """Token check. Failed attempts are counted per IP (10/min → 429) so the token cannot be brute-forced; must run AFTER `rate_limit('admin')`."""227 ip = client_ip(request)228 if not _bucket_hit(ip, "admin-auth-failed", record=False):229 raise ApiError(429, "too many failed admin authentications; retry in a minute")230 expected = settings.admin_token231 if not expected:232 raise ApiError(503, "admin API disabled: AIA_ADMIN_TOKEN is not configured")233 token = request.headers.get("x-aia-admin-token") or ""234 if not hmac.compare_digest(token, expected):235 _bucket_hit(ip, "admin-auth-failed")236 raise ApiError(401, "invalid or missing x-aia-admin-token")237238239ADMIN_DEPENDENCIES = [Depends(rate_limit("admin")), Depends(require_admin)] # order matters: limit first, then auth240241242async def audit(action: str, target: str | None = None, payload: dict[str, Any] | None = None, ip: str | None = None, *, actor: str = "admin") -> None:243 """One row in `admin_audit_log` per admin action (best effort — never breaks the request)."""244 try:245 async with transaction() as conn:246 await execute(conn, "insert into admin_audit_log (actor, action, target, payload, ip) values (:a, :ac, :t, cast(:p as jsonb), :ip)",247 a=actor, ac=action, t=target, p=jsonb(payload or {}), ip=ip)248 except Exception: # noqa: BLE001249 log.warning("audit log write failed", extra={"action": action, "target": target})250251252# ------------------------------------------------------------------------------------------------------------------ entity serialisers253254ENTITY_FIELDS = ("id", "entity_type", "canonical_name", "slug", "description", "status", "organization_id", "attributes", "quality", "counts",255 "first_seen_at", "last_seen_at", "updated_at", "organization_name", "organization_slug")256257258def entity_cols(alias: str = "e", prefix: str = "") -> str:259 """Select-list for an entity alias (+ its organization alias `<alias>o`) with an optional column prefix."""260 org = f"{alias}o"261 cols = [f"{alias}.{c} as {prefix}{c}" for c in ENTITY_FIELDS[:13]]262 cols += [f"{org}.canonical_name as {prefix}organization_name", f"{org}.slug as {prefix}organization_slug"]263 return ", ".join(cols)264265266def entity_join(alias: str, on: str) -> str:267 return f"join entities {alias} on {alias}.id = {on} left join entities {alias}o on {alias}o.id = {alias}.organization_id"268269270ENTITY_COLS = entity_cols("e")271ENTITY_FROM = "entities e left join entities eo on eo.id = e.organization_id"272273SUMMARY_ATTRS: dict[str, tuple[str, ...]] = {274 "model": ("family", "openness", "license", "parameter_count", "active_parameter_count", "context_length", "max_output_tokens", "modalities",275 "release_date", "status", "knowledge_cutoff", "hf_repo", "api_model_id"),276 "company": ("country", "founded", "website", "org_kind"),277 "organization": ("country", "founded", "website", "org_kind"),278 "lab": ("country", "founded", "website", "org_kind"),279 "university": ("country", "founded", "website", "org_kind"),280 "provider": ("website", "pricing_url"),281 "benchmark": ("category", "metric", "unit"),282 "hardware": ("kind", "memory_gb", "memory_bandwidth_gbs", "release_date", "manufacturer"),283 "framework": ("latest_version", "latest_release_at", "license", "language", "metric.stars"),284 "repository": ("latest_version", "latest_release_at", "license", "language", "metric.stars"),285 "paper": ("authors", "published_at", "arxiv_id", "primary_category"),286 "dataset": ("license", "modality"),287}288COMPANY_TYPES = ("company", "organization", "lab", "university")289TYPE_LABELS = {"model": "Models", "company": "Companies", "organization": "Organizations", "lab": "Labs", "university": "Universities",290 "provider": "Providers", "paper": "Papers", "benchmark": "Benchmarks", "hardware": "Hardware", "framework": "Frameworks",291 "dataset": "Datasets", "tool": "Tools", "repository": "Repositories", "regulation": "Regulation", "incident": "Incidents",292 "researcher": "Researchers", "agent": "Agents", "product": "Products", "runtime": "Runtimes", "conference": "Conferences",293 "artifact": "Artifacts", "model_family": "Model families", "license": "Licenses"}294295# Canonical model universe (API 1.1): model releases only — artifacts (checkpoints, quantisations, conversions) are `entity_type = 'artifact'`,296# folded evaluation variants carry `merged_into`. `include=artifacts` on /models restores the pre-1.1 universe.297MODEL_UNIVERSE = "e.entity_type = 'model' and e.merged_into is null"298MODEL_OR_ARTIFACT_UNIVERSE = "e.entity_type in ('model', 'artifact') and e.merged_into is null"299OPEN_CATEGORIES = ("open-weights", "open-source") # "open-*"300DOWNLOADABLE_CATEGORIES = ("open-weights", "open-source", "restricted-weights", "restricted") # weights can be downloaded (terms may restrict)301ARTIFACT_KINDS = ("checkpoint", "quantization", "conversion", "packaging")302303304def openness_values(raw: str | None) -> list[str]:305 """Expand the `openness=` filter vocabulary: `open` → open-weights + open-source (+ legacy `open`); `restricted` ↔ `restricted-weights`."""306 vals = [v.strip() for v in (raw or "").split(",") if v.strip()]307 out: list[str] = []308 for v in vals:309 if v == "open":310 out += [*OPEN_CATEGORIES, "open"]311 elif v in ("restricted", "restricted-weights"):312 out += ["restricted", "restricted-weights"]313 elif v == "downloadable":314 out += list(DOWNLOADABLE_CATEGORIES)315 else:316 out.append(v)317 return list(dict.fromkeys(out))318319320def is_open(openness: Any) -> bool:321 return str(openness or "") in OPEN_CATEGORIES322323324def summary_attributes(entity_type: str, attrs: dict[str, Any] | None) -> dict[str, Any]:325 attrs = attrs or {}326 keys = SUMMARY_ATTRS.get(entity_type)327 if keys is None:328 return {k: v for k, v in attrs.items() if isinstance(v, (int, float, bool)) or (isinstance(v, str) and len(v) <= 200)}329 out: dict[str, Any] = {}330 for k in keys:331 if k in attrs and attrs[k] not in (None, "", [], {}):332 v = attrs[k]333 out[k] = v[:5] if k == "authors" and isinstance(v, list) else v334 return out335336337STATUS_VOCAB = ("active", "preview", "deprecated", "retired", "announced", "limited-availability", "unknown")338STATUS_MAP = {"available": "active", "ga": "active", "general-availability": "active", "released": "active", "live": "active", "beta": "preview",339 "alpha": "preview", "experimental": "preview", "coming-soon": "announced", "upcoming": "announced", "sunset": "retired", "discontinued": "retired",340 "archived": "retired", "legacy": "deprecated", "limited": "limited-availability", "merged": "retired"}341342343def normalize_status(value: Any) -> str:344 s = str(value or "").strip().lower().replace("_", "-").replace(" ", "-")345 s = STATUS_MAP.get(s, s)346 return s if s in STATUS_VOCAB else "unknown"347348349MAIN_ENTITY_TYPES = ("model", "company", "paper", "provider", "benchmark", "hardware", "framework", "dataset", "tool", "repository")350351EVENT_TYPE_LABELS: dict[str, tuple[str, int]] = { # event_type -> (human label, default importance 0–3)352 "NEW_MODEL": ("New model", 3), "MODEL_UPDATED": ("Model updated", 1), "PRICE_CHANGED": ("Price change", 2), "CONTEXT_CHANGED": ("Context window change", 2),353 "MAX_OUTPUT_CHANGED": ("Max output change", 1), "PARAMETERS_CHANGED": ("Parameter count change", 2), "KNOWLEDGE_CUTOFF_CHANGED": ("Knowledge cutoff change", 1),354 "RELEASE_DATE_CHANGED": ("Release date change", 1), "STATUS_CHANGED": ("Status change", 2), "LICENSE_CHANGED": ("License change", 2),355 "OPENNESS_CHANGED": ("Openness change", 2), "CAPABILITIES_CHANGED": ("Capabilities change", 1), "PROPERTY_CHANGED": ("Property change", 1),356 "DEPRECATION_ANNOUNCED": ("Deprecation announced", 2), "RETIREMENT_ANNOUNCED": ("Retirement announced", 2), "BENCHMARK_RESULT": ("Benchmark result", 1),357 "BENCHMARK_UPDATED": ("Benchmark updated", 1), "NEW_PAPER": ("New paper", 2), "NEW_COMPANY": ("New company", 2), "NEW_ORGANIZATION": ("New organization", 2),358 "NEW_LAB": ("New lab", 2), "NEW_UNIVERSITY": ("New university", 1), "NEW_PROVIDER": ("New provider", 2), "NEW_HARDWARE": ("New hardware", 2),359 "NEW_FRAMEWORK": ("New framework", 2), "NEW_LIBRARY": ("New library", 1), "NEW_BENCHMARK": ("New benchmark", 2), "NEW_DATASET": ("New dataset", 2),360 "NEW_REPOSITORY": ("New repository", 1), "NEW_TOOL": ("New tool", 1), "PROVIDER_LISTED": ("Listed by provider", 2), "PROVIDER_DELISTED": ("Delisted by provider", 2),361 "ANNOUNCEMENT": ("Announcement", 2), "RELEASE": ("Release", 2), "VERSION_RELEASED": ("Version released", 1), "DOCUMENT_CHANGED": ("Source document changed", 0),362 "ENTITY_MERGED": ("Duplicate merged", 1), "CLAIM_RETRACTED": ("Claim retracted", 1),363}364365366def event_type_label(event_type: str) -> str:367 known = EVENT_TYPE_LABELS.get(event_type)368 return known[0] if known else event_type.replace("_", " ").capitalize()369370371def event_type_importance(event_type: str) -> int:372 known = EVENT_TYPE_LABELS.get(event_type)373 return known[1] if known else 1374375376async def enrich_provenance(conn: AsyncConnection, *provenances: dict[str, Any] | None) -> None:377 """Add `source_name` (from `sources.name`) to every provenance entry that carries a `source_id` — in place."""378 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")}379 if not ids:380 return381 from aiatlas.db import fetch_all382383 rows = await fetch_all(conn, "select id, name, domain, tier from sources where id = any(cast(:ids as text[]))", ids=sorted(ids))384 names = {r["id"]: r for r in rows}385 for p in provenances:386 for v in (p or {}).values():387 if isinstance(v, dict) and v.get("source_id") in names:388 src = names[v["source_id"]]389 v.setdefault("source_name", src["name"])390 v.setdefault("source_domain", src["domain"])391392393def org_of(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None:394 oid = row.get(f"{prefix}organization_id")395 if not oid:396 return None397 return {"id": oid, "slug": row.get(f"{prefix}organization_slug"), "name": row.get(f"{prefix}organization_name")}398399400def entity_summary(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None:401 g = row.get402 if not g(f"{prefix}id"):403 return None404 etype = g(f"{prefix}entity_type") or ""405 desc = g(f"{prefix}description")406 raw_status = g(f"{prefix}status")407 status = normalize_status(raw_status)408 return {"id": g(f"{prefix}id"), "entity_type": etype, "slug": g(f"{prefix}slug"), "name": g(f"{prefix}canonical_name"),409 "description": (desc[:280] if isinstance(desc, str) and len(desc) > 280 else desc), "status": status,410 **({"status_raw": raw_status} if raw_status and raw_status != status else {}),411 "organization": org_of(row, prefix), "attributes": summary_attributes(etype, g(f"{prefix}attributes")),412 "quality": g(f"{prefix}quality") or {}, "counts": g(f"{prefix}counts") or {},413 "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")}414415416def change_event(row: dict[str, Any], prefix: str = "e_") -> dict[str, Any]:417 return {"id": row["id"], "event_type": row["event_type"], "category": row["category"], "property": row.get("property"),418 "old_value": row.get("old_value"), "new_value": row.get("new_value"), "summary": row.get("summary"), "importance": row.get("importance"),419 "observed_at": row.get("observed_at"), "effective_at": row.get("effective_at"), "source_url": row.get("source_url"),420 "connector_name": row.get("connector_name"), "entity": entity_summary(row, prefix), "meta": row.get("meta") or {}}421422423EVENT_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, "424 "ev.source_url, ev.connector_name, ev.meta, " + entity_cols("e", "e_"))425EVENT_FROM = "change_events ev left join entities e on e.id = ev.entity_id left join entities eo on eo.id = e.organization_id"426427428def price_row(row: dict[str, Any]) -> dict[str, Any]:429 return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"),430 "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"),431 "cache_write_per_mtok": row.get("cache_write_per_mtok"), "batch_input_per_mtok": row.get("batch_input_per_mtok"),432 "batch_output_per_mtok": row.get("batch_output_per_mtok"), "per_image": row.get("per_image"), "per_request": row.get("per_request"),433 "currency": row.get("currency"), "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"),434 "features": row.get("features") or {}, "observed_at": row.get("observed_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"),435 "source_url": row.get("source_url"), "tier": row.get("tier")}436437438PRICE_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, "439 "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, "440 "p.valid_to, p.source_url, p.tier, " + entity_cols("m", "m_") + ", " + entity_cols("pv", "p_"))441PRICE_FROM = "prices p " + entity_join("m", "p.model_id") + " " + entity_join("pv", "p.provider_id")442443_NATIVE_UNIT_SUFFIXES = ("_per_mtok", "_per_1k_requests", "_per_mtok_hour", "_per_image", "_per_request", "_per_second", "_per_minute", "_per_hour", "_per_char")444445446def deployment_row(row: dict[str, Any]) -> dict[str, Any]:447 """`Deployment` (API 1.1): one model × provider × provider_model_id offer. Prices in USD per 1M tokens unless the key says otherwise;448 provider-specific priced features (`flex_input_per_mtok`, `search_grounding_per_1k_requests`…) are kept verbatim under `native_units`."""449 feats = dict(row.get("features") or {})450 native = {k: v for k, v in feats.items() if any(k.endswith(s) for s in _NATIVE_UNIT_SUFFIXES) or k == "per_request"}451 other = {k: v for k, v in feats.items() if k not in native}452 return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"),453 "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"),454 "prices": {"input": row.get("input_per_mtok"), "cached_input": row.get("cached_input_per_mtok"), "cache_write": row.get("cache_write_per_mtok"),455 "output": row.get("output_per_mtok"), "batch_input": row.get("batch_input_per_mtok"), "batch_output": row.get("batch_output_per_mtok"),456 "per_image": row.get("per_image"), "per_request": row.get("per_request"), "currency": row.get("currency") or "USD", "unit": "USD per 1M tokens",457 "native_units": native},458 "features": other, "status": "active" if row.get("valid_to") is None else "delisted", "observed_at": row.get("observed_at"),459 "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"), "tier": row.get("tier")}460461462def result_row(row: dict[str, Any]) -> dict[str, Any]:463 return {"id": row["id"], "model": entity_summary(row, "m_"), "benchmark": entity_summary(row, "b_"), "score": row.get("score"), "metric": row.get("metric"),464 "unit": row.get("unit"), "higher_is_better": row.get("higher_is_better"), "config": row.get("config") or {}, "evaluated_at": row.get("evaluated_at"),465 "observed_at": row.get("observed_at"), "source_url": row.get("source_url"), "tier": row.get("tier"), "confidence": row.get("confidence"),466 "valid_to": row.get("valid_to")}467468469RESULT_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, "470 + entity_cols("m", "m_") + ", " + entity_cols("b", "b_"))471RESULT_FROM = "benchmark_results r " + entity_join("m", "r.model_id") + " " + entity_join("b", "r.benchmark_id")472RESULT_ORDER = "case when r.higher_is_better then -r.score else r.score end, r.observed_at desc"473474475def claim_row(row: dict[str, Any]) -> dict[str, Any]:476 return {"id": row["id"], "property": row["property"], "value": row.get("value"), "unit": row.get("unit"), "tier": row.get("tier"),477 "confidence": row.get("confidence"), "status": row.get("status"), "extractor": row.get("extractor"), "observed_at": row.get("observed_at"),478 "effective_at": row.get("effective_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"),479 "source_name": row.get("source_name")}480481482CLAIM_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, "483 "c.source_url, s.name as source_name")484CLAIM_FROM = "claims c left join sources s on s.id = c.source_id"485486487# ------------------------------------------------------------------------------------------------------------------ entity resolution488489490ENTITY_EXTRA_COLS = "e.provenance, e.merged_into, e.family_id, e.canonical_id, e.artifact_kind, e.identity_confidence"491492493async def resolve_entity(conn: AsyncConnection, slug_or_id: str, types: tuple[str, ...] | None = None, *, aliases: bool = False) -> dict[str, Any]:494 """Slug first, then id (then alias when `aliases=True`); follows `merged_into` and records the hop in `redirected_from`495 (`{slug, id}` of the row the caller asked for) so the web layer can 301; 404 when absent or when the type does not match the mounted alias."""496 key = slug_or_id.strip()497 if not key or len(key) > 200:498 raise ApiError(404, "entity not found")499 sel = f"select {ENTITY_COLS}, {ENTITY_EXTRA_COLS} from {ENTITY_FROM}"500 row = await fetch_one(conn, f"{sel} where e.slug = :k or e.id = :k order by case when e.slug = :k then 0 else 1 end limit 1", k=key)501 if not row and aliases:502 row = await fetch_one(conn, f"{sel} where e.id = (select a.entity_id from entity_aliases a join entities x on x.id = a.entity_id "503 f"where (a.alias_norm = :n or lower(a.alias) = lower(:k)) {'and x.entity_type = any(cast(:types as text[]))' if types else ''} "504 f"order by x.merged_into is not null, x.updated_at desc limit 1)", k=key, n=_alias_norm(key), types=list(types or ()))505 asked = row506 hops = 0507 while row and row.get("merged_into") and hops < 5:508 row = await fetch_one(conn, f"{sel} where e.id = :k", k=row["merged_into"])509 hops += 1510 if not row:511 raise ApiError(404, "entity not found")512 if types and row["entity_type"] not in types:513 raise ApiError(404, f"entity {key!r} is a {row['entity_type']}, not one of {', '.join(types)}")514 if asked is not None and asked["id"] != row["id"]:515 row["redirected_from"] = {"slug": asked["slug"], "id": asked["id"], "entity_type": asked["entity_type"]}516 return row517518519def _alias_norm(s: str) -> str:520 try:521 from aiatlas.ids import normalize_alias522523 return normalize_alias(s)524 except Exception: # noqa: BLE001525 return "".join(ch for ch in s.lower() if ch.isalnum())526527528async def resolve_id(conn: AsyncConnection, slug_or_id: str | None, types: tuple[str, ...] | None = None) -> str | None:529 if not slug_or_id:530 return None531 return (await resolve_entity(conn, slug_or_id, types))["id"]532533534__all__ = [535 "ADMIN_DEPENDENCIES",536 "ARTIFACT_KINDS",537 "CLAIM_COLS",538 "CLAIM_FROM",539 "COMPANY_TYPES",540 "DOWNLOADABLE_CATEGORIES",541 "ENTITY_COLS",542 "ENTITY_EXTRA_COLS",543 "ENTITY_FROM",544 "EVENT_COLS",545 "EVENT_FROM",546 "EVENT_TYPE_LABELS",547 "MAIN_ENTITY_TYPES",548 "MODEL_OR_ARTIFACT_UNIVERSE",549 "MODEL_UNIVERSE",550 "OPEN_CATEGORIES",551 "PAGINATION",552 "PRICE_COLS",553 "PRICE_FROM",554 "RESULT_COLS",555 "RESULT_FROM",556 "RESULT_ORDER",557 "STATUS_VOCAB",558 "TYPE_LABELS",559 "ApiError",560 "AtlasJSONResponse",561 "Pagination",562 "attr_num",563 "audit",564 "cache_key",565 "cached",566 "change_event",567 "claim_row",568 "client_ip",569 "csv",570 "day_bounds",571 "deployment_row",572 "dumps",573 "enrich_provenance",574 "entity_cols",575 "entity_join",576 "entity_summary",577 "event_type_importance",578 "event_type_label",579 "flip_order",580 "is_open",581 "normalize",582 "normalize_status",583 "num_expr",584 "openness_values",585 "org_of",586 "page",587 "parse_date",588 "parse_ts",589 "price_row",590 "rate_limit",591 "require_admin",592 "resolve_entity",593 "resolve_id",594 "result_row",595 "trust_proxy",596]597