"""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 logging import os 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 execute, fetch_one, jsonb, transaction from aiatlas.services import cache log = logging.getLogger("aiatlas.api") # ------------------------------------------------------------------------------------------------------------------ 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 # ------------------------------------------------------------------------------------------------------------------ rate limiting _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)} _hits: dict[tuple[str, str], deque[float]] = defaultdict(deque) def trust_proxy() -> bool: """`x-forwarded-for` is only honoured behind a trusted reverse proxy (`settings.trust_proxy` or env `AIA_TRUST_PROXY=1`).""" v = getattr(settings, "trust_proxy", None) if v is None: v = os.environ.get("AIA_TRUST_PROXY", "") return str(v).strip().lower() in ("1", "true", "yes", "on") def client_ip(request: Request) -> str: if trust_proxy(): fwd = request.headers.get("x-forwarded-for") if fwd: return fwd.split(",")[0].strip() return request.client.host if request.client else "unknown" def _bucket_hit(ip: str, bucket: str, *, record: bool = True) -> bool: """Sliding-window counter. Returns True when the call is allowed (and records it), False when the bucket is full.""" n, window = _LIMITS.get(bucket, _LIMITS["default"]) now = time.monotonic() q = _hits[(ip, bucket)] while q and q[0] <= now - window: q.popleft() if len(q) >= n: return False if record: 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 True def rate_limit(bucket: str) -> Callable[[Request], None]: def dep(request: Request) -> None: n, window = _LIMITS.get(bucket, _LIMITS["default"]) if not _bucket_hit(client_ip(request), bucket): raise ApiError(429, f"rate limit exceeded for {bucket}: {n} per {int(window)} s") return dep # ------------------------------------------------------------------------------------------------------------------ admin auth def require_admin(request: Request) -> None: """Token check. Failed attempts are counted per IP (10/min → 429) so the token cannot be brute-forced; must run AFTER `rate_limit('admin')`.""" ip = client_ip(request) if not _bucket_hit(ip, "admin-auth-failed", record=False): raise ApiError(429, "too many failed admin authentications; retry in a minute") 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): _bucket_hit(ip, "admin-auth-failed") raise ApiError(401, "invalid or missing x-aia-admin-token") ADMIN_DEPENDENCIES = [Depends(rate_limit("admin")), Depends(require_admin)] # order matters: limit first, then auth async def audit(action: str, target: str | None = None, payload: dict[str, Any] | None = None, ip: str | None = None, *, actor: str = "admin") -> None: """One row in `admin_audit_log` per admin action (best effort — never breaks the request).""" try: async with transaction() as conn: await execute(conn, "insert into admin_audit_log (actor, action, target, payload, ip) values (:a, :ac, :t, cast(:p as jsonb), :ip)", a=actor, ac=action, t=target, p=jsonb(payload or {}), ip=ip) except Exception: # noqa: BLE001 log.warning("audit log write failed", extra={"action": action, "target": target}) # ------------------------------------------------------------------------------------------------------------------ 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", "license_key", "parameter_count", "active_parameter_count", "context_length", "max_output_tokens", "modalities", "release_date", "status", "knowledge_cutoff", "hf_repo", "api_model_id", "reasoning", "tool_calling"), "artifact": ("openness", "license", "license_key", "parameter_count", "quant_format", "file_size_gb", "weights_dtype", "hf_repo", "release_date", "base_model"), "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": ("kind", "latest_version", "latest_release_at", "license", "license_key", "language", "metric.stars", "repository_url", "homepage"), "library": ("kind", "latest_version", "latest_release_at", "license", "license_key", "language", "metric.stars", "repository_url", "homepage"), "runtime": ("kind", "latest_version", "latest_release_at", "license", "license_key", "language", "metric.stars", "repository_url"), "repository": ("latest_version", "latest_release_at", "license", "language", "metric.stars"), "paper": ("authors", "published_at", "arxiv_id", "primary_category"), "dataset": ("license", "license_key", "modality", "modalities", "publisher", "size", "access", "hf_repo", "hf_dataset", "task", "languages", "release_date"), } def _fill_license_key(etype: str, attrs: dict[str, Any], out: dict[str, Any]) -> None: """Canonical `license_key` derived on read when the writer has not stored it yet (raw label → ontology key; unknown labels stay absent).""" if "license_key" not in out and attrs.get("license"): from aiatlas.ontology.licenses import normalize_license key = normalize_license(str(attrs["license"])) if key: out["license_key"] = key 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", "artifact": "Artifacts", "model_family": "Model families", "license": "Licenses"} # Canonical model universe (API 1.1): model releases only — artifacts (checkpoints, quantisations, conversions) are `entity_type = 'artifact'`, # folded evaluation variants carry `merged_into`. `include=artifacts` on /models restores the pre-1.1 universe. MODEL_UNIVERSE = "e.entity_type = 'model' and e.merged_into is null" MODEL_OR_ARTIFACT_UNIVERSE = "e.entity_type in ('model', 'artifact') and e.merged_into is null" OPEN_CATEGORIES = ("open-weights", "open-source") # "open-*" DOWNLOADABLE_CATEGORIES = ("open-weights", "open-source", "restricted-weights", "restricted") # weights can be downloaded (terms may restrict) ARTIFACT_KINDS = ("checkpoint", "quantization", "conversion", "packaging") def openness_values(raw: str | None) -> list[str]: """Expand the `openness=` filter vocabulary: `open` → open-weights + open-source (+ legacy `open`); `restricted` ↔ `restricted-weights`.""" vals = [v.strip() for v in (raw or "").split(",") if v.strip()] out: list[str] = [] for v in vals: if v == "open": out += [*OPEN_CATEGORIES, "open"] elif v in ("restricted", "restricted-weights"): out += ["restricted", "restricted-weights"] elif v == "downloadable": out += list(DOWNLOADABLE_CATEGORIES) else: out.append(v) return list(dict.fromkeys(out)) def is_open(openness: Any) -> bool: return str(openness or "") in OPEN_CATEGORIES 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 if "license_key" in keys: _fill_license_key(entity_type, attrs, out) if "kind" in keys and out.get("kind") and entity_type in ("framework", "library", "runtime"): from aiatlas.ontology.taxonomy import normalize_framework_kind canon = normalize_framework_kind(out["kind"]) if canon and canon != out["kind"]: out["kind_raw"], out["kind"] = out["kind"], canon 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") _NATIVE_UNIT_SUFFIXES = ("_per_mtok", "_per_1k_requests", "_per_mtok_hour", "_per_image", "_per_request", "_per_second", "_per_minute", "_per_hour", "_per_char") def deployment_row(row: dict[str, Any]) -> dict[str, Any]: """`Deployment` (API 1.1): one model × provider × provider_model_id offer. Prices in USD per 1M tokens unless the key says otherwise; provider-specific priced features (`flex_input_per_mtok`, `search_grounding_per_1k_requests`…) are kept verbatim under `native_units`.""" feats = dict(row.get("features") or {}) native = {k: v for k, v in feats.items() if any(k.endswith(s) for s in _NATIVE_UNIT_SUFFIXES) or k == "per_request"} other = {k: v for k, v in feats.items() if k not in native} return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"), "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"), "prices": {"input": row.get("input_per_mtok"), "cached_input": row.get("cached_input_per_mtok"), "cache_write": row.get("cache_write_per_mtok"), "output": row.get("output_per_mtok"), "batch_input": row.get("batch_input_per_mtok"), "batch_output": row.get("batch_output_per_mtok"), "per_image": row.get("per_image"), "per_request": row.get("per_request"), "currency": row.get("currency") or "USD", "unit": "USD per 1M tokens", "native_units": native}, "features": other, "status": "active" if row.get("valid_to") is None else "delisted", "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")} 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 ENTITY_EXTRA_COLS = "e.provenance, e.merged_into, e.family_id, e.canonical_id, e.artifact_kind, e.identity_confidence" async def resolve_entity(conn: AsyncConnection, slug_or_id: str, types: tuple[str, ...] | None = None, *, aliases: bool = False) -> dict[str, Any]: """Slug first, then id (then alias when `aliases=True`); follows `merged_into` and records the hop in `redirected_from` (`{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.""" key = slug_or_id.strip() if not key or len(key) > 200: raise ApiError(404, "entity not found") sel = f"select {ENTITY_COLS}, {ENTITY_EXTRA_COLS} from {ENTITY_FROM}" 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) if not row and aliases: 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 " f"where (a.alias_norm = :n or lower(a.alias) = lower(:k)) {'and x.entity_type = any(cast(:types as text[]))' if types else ''} " 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 ())) asked = row hops = 0 while row and row.get("merged_into") and hops < 5: row = await fetch_one(conn, f"{sel} 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)}") if asked is not None and asked["id"] != row["id"]: row["redirected_from"] = {"slug": asked["slug"], "id": asked["id"], "entity_type": asked["entity_type"]} return row def _alias_norm(s: str) -> str: try: from aiatlas.ids import normalize_alias return normalize_alias(s) except Exception: # noqa: BLE001 return "".join(ch for ch in s.lower() if ch.isalnum()) 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__ = [ "ADMIN_DEPENDENCIES", "ARTIFACT_KINDS", "CLAIM_COLS", "CLAIM_FROM", "COMPANY_TYPES", "DOWNLOADABLE_CATEGORIES", "ENTITY_COLS", "ENTITY_EXTRA_COLS", "ENTITY_FROM", "EVENT_COLS", "EVENT_FROM", "EVENT_TYPE_LABELS", "MAIN_ENTITY_TYPES", "MODEL_OR_ARTIFACT_UNIVERSE", "MODEL_UNIVERSE", "OPEN_CATEGORIES", "PAGINATION", "PRICE_COLS", "PRICE_FROM", "RESULT_COLS", "RESULT_FROM", "RESULT_ORDER", "STATUS_VOCAB", "TYPE_LABELS", "ApiError", "AtlasJSONResponse", "Pagination", "attr_num", "audit", "cache_key", "cached", "change_event", "claim_row", "client_ip", "csv", "day_bounds", "deployment_row", "dumps", "enrich_provenance", "entity_cols", "entity_join", "entity_summary", "event_type_importance", "event_type_label", "flip_order", "is_open", "normalize", "normalize_status", "num_expr", "openness_values", "org_of", "page", "parse_date", "parse_ts", "price_row", "rate_limit", "require_admin", "resolve_entity", "resolve_id", "result_row", "trust_proxy", ]