API v1: FastAPI routers per docs/API.md (public + admin), detail builder, hardware fit, entity merge, indexes, tests
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
29 changed files +3,317 −4
modified
docs/API.md
+31 −0
@@ -117,3 +117,34 @@ EntitySummary & { | ||
| 117 | 117 | | `GET /admin/infrastructure` | heartbeats, archive size, DB size (`pg_database_size`), table sizes, node hostname, uptime | |
| 118 | 118 | | `POST /admin/cache/flush` | `{ flushed: n }` | |
| 119 | 119 | | `POST /admin/stats/recompute` · `POST /admin/quality/recompute` | maintenance | |
| 120 | + | |
| 121 | +## Implementation notes (2026-09-11, `src/aiatlas/api`) | |
| 122 | + | |
| 123 | +Deviations and additions relative to the table above — everything else is implemented as written. | |
| 124 | + | |
| 125 | +- **Validation errors**: parameters rejected by FastAPI/pydantic typing or bounds (`limit=999`, non-integer `depth`…) return **422** with | |
| 126 | + `{ detail: "<field>: <message>", errors: [...] }`; semantic errors raised by the routes (bad `sort`, bad date, wrong `scope`, `ids` count) return **400**. | |
| 127 | +- **`/companies`** (listing, alias, `/compare` type `company`) covers the organization-like types `company | organization | lab | university`. | |
| 128 | + `/frameworks/{slug}` also accepts `library | runtime`; `/tools/{slug}` accepts `tool | agent | application | product | mcp_server`. | |
| 129 | +- **`/explore/{type}`** accepts plural forms (`/explore/models`, `/explore/companies`) and adds `entity_type` + `label` to the page. | |
| 130 | +- **`EntitySummary.description`** is truncated at 280 characters in summaries (full text in `EntityDetail`). | |
| 131 | +- **Cursor feeds** (`/changes`, `/entities/{slug}/timeline`) add `next_before` (the `observed_at` of the last item, or `null` at the end). | |
| 132 | + `/changes` also honours `offset` and `entity=` for convenience; `total` is capped at 10 000. | |
| 133 | +- **`/changes/daily`** adds `total`, `labels`, `previous_day`, `next_day`; sections are capped at `per_section` (default 30) items each. | |
| 134 | +- **`/timeline`** month groups carry `count`; the response includes `total` (events returned). | |
| 135 | +- **`/search`**: `query` carries `semantic: bool` (embedding used or not). When the embedding endpoint times out (2 s) the API backs off to | |
| 136 | + FTS-only for 5 minutes (`aia:api:search:embed-degraded`) so the LLM never adds latency to consecutive queries. | |
| 137 | +- **`hardware_fit`** (model detail) is *omitted* when `parameter_count` is unknown (nothing is estimated from thin air); when present the | |
| 138 | + detail also carries `hardware_fit_assumptions: string[]`. `/hardware/fit` adds `estimated: true`, `counts`, optional `openness=` filter. | |
| 139 | +- **`/prices/index`** series rows add `max_input` and `offers`; response adds `days` and `note`. `/prices/history` requires `model` or `provider`. | |
| 140 | +- **`/benchmarks/{slug}/results`** and `/history` include the `benchmark` summary; result rows include `valid_to`. | |
| 141 | +- **`/diff`** adds `scope` (resolved description) and extra `counts` (`events`, `price_rows_opened/closed`, `claims_superseded`, `entities_at_a/b`). | |
| 142 | +- **`/sources`** items add `snapshots`, `claims`, `base_url`, `robots_policy`, `priority`, `notes`; response adds `tiers` legend. | |
| 143 | +- **`/methodology`** adds `quality_version`, `expected_fields`, `principles`; `event_types` are live counts from `change_events`. | |
| 144 | +- **`/trending`** accepts `type=`; `/sitemap` `limit` ≤ 5000 (the one listing not capped at 200). | |
| 145 | +- **Rate limits** are in-process sliding windows (search 60/min, views 1/s, admin 240/min per IP) — fine for the single uvicorn worker behind Next.js. | |
| 146 | +- **Admin**: `GET /admin/review?status=all` lists every status; approving a `conflict` with `resolution.keep_claim_id` promotes that claim; | |
| 147 | + `PATCH /admin/connectors/{name}` also resets `health`/`circuit_open_until` when toggling; `/admin/connectors` flags `run_now_pending` and | |
| 148 | + connectors present in code but not in the table; `/admin/snapshots/{id}` never returns `raw_path`/`text_path` (`has_raw`/`has_text` booleans instead). | |
| 149 | +- **Migration `0002_api_indexes`** adds partial indexes for the read paths (type + release_date/family/openness/updated/first_seen, importance-ordered | |
| 150 | + events, current leaderboards/prices, page views, review kinds). | |
added
migrations/versions/0002_api_indexes.py
+63 −0
@@ -0,0 +1,63 @@ | ||
| 1 | +"""API read-path indexes: attribute-driven listings (release_date, family, openness), importance-ordered change feeds, | |
| 2 | +current benchmark leaderboards, page-view trending, review queue by kind. | |
| 3 | + | |
| 4 | +Revision ID: 0002 | |
| 5 | +""" | |
| 6 | +from __future__ import annotations | |
| 7 | + | |
| 8 | +from alembic import op | |
| 9 | + | |
| 10 | +revision = "0002" | |
| 11 | +down_revision = "0001" | |
| 12 | +branch_labels = None | |
| 13 | +depends_on = None | |
| 14 | + | |
| 15 | +SQL = r""" | |
| 16 | +create index if not exists entities_type_release_idx on entities (entity_type, (attributes->>'release_date') desc nulls last) where merged_into is null; | |
| 17 | +create index if not exists entities_type_family_idx on entities (entity_type, (attributes->>'family')) where merged_into is null; | |
| 18 | +create index if not exists entities_type_openness_idx on entities (entity_type, (attributes->>'openness')) where merged_into is null; | |
| 19 | +create index if not exists entities_type_updated_idx on entities (entity_type, updated_at desc) where merged_into is null; | |
| 20 | +create index if not exists entities_type_first_seen_idx on entities (entity_type, first_seen_at desc) where merged_into is null; | |
| 21 | +create index if not exists entities_slug_id_idx on entities (id) include (slug); | |
| 22 | +create index if not exists change_events_importance_idx on change_events (importance desc, observed_at desc); | |
| 23 | +create index if not exists change_events_type_observed_idx on change_events (observed_at desc) where event_type <> 'DOCUMENT_CHANGED'; | |
| 24 | +create index if not exists benchmark_results_current_idx on benchmark_results (benchmark_id, higher_is_better, score) where valid_to is null; | |
| 25 | +create index if not exists benchmark_results_model_current_idx on benchmark_results (model_id) where valid_to is null; | |
| 26 | +create index if not exists prices_current_model_idx on prices (model_id, provider_id) where valid_to is null; | |
| 27 | +create index if not exists prices_valid_range_idx on prices (valid_from, valid_to); | |
| 28 | +create index if not exists page_views_day_idx on page_views (day desc, views desc); | |
| 29 | +create index if not exists review_queue_kind_idx on review_queue (kind, status, created_at desc); | |
| 30 | +create index if not exists documents_source_idx on documents (source_id, last_fetched_at desc); | |
| 31 | +create index if not exists claims_entity_valid_idx on claims (entity_id, valid_from, valid_to); | |
| 32 | +create index if not exists relations_predicate_idx on relations (predicate, subject_id, object_id) where valid_to is null; | |
| 33 | +""" | |
| 34 | + | |
| 35 | + | |
| 36 | +def upgrade() -> None: | |
| 37 | + for statement in _split(SQL): | |
| 38 | + op.execute(statement) | |
| 39 | + | |
| 40 | + | |
| 41 | +def downgrade() -> None: | |
| 42 | + raise RuntimeError("AI Atlas migrations are forward-only: historical data is never disposable") | |
| 43 | + | |
| 44 | + | |
| 45 | +def _split(sql: str) -> list[str]: | |
| 46 | + out: list[str] = [] | |
| 47 | + buf: list[str] = [] | |
| 48 | + in_dollar = False | |
| 49 | + for line in sql.splitlines(): | |
| 50 | + stripped = line.strip() | |
| 51 | + if stripped.count("$$") % 2 == 1: | |
| 52 | + in_dollar = not in_dollar | |
| 53 | + buf.append(line) | |
| 54 | + if not in_dollar and stripped.endswith(";"): | |
| 55 | + body = [ln for ln in buf if not ln.strip().startswith("--") or in_dollar] | |
| 56 | + stmt = "\n".join(body).strip() | |
| 57 | + if stmt: | |
| 58 | + out.append(stmt) | |
| 59 | + buf = [] | |
| 60 | + tail = "\n".join(ln for ln in buf if not ln.strip().startswith("--")).strip() | |
| 61 | + if tail: | |
| 62 | + out.append(tail) | |
| 63 | + return out | |
added
src/aiatlas/api/__init__.py
+1 −0
@@ -0,0 +1 @@ | ||
| 1 | +"""AI Atlas public + admin API (FastAPI, `/api/v1`). Contract: docs/API.md.""" | |
added
src/aiatlas/api/common.py
+412 −0
@@ -0,0 +1,412 @@ | ||
| 1 | +"""Shared API helpers: errors, pagination, date parsing, Redis response cache, admin auth, IP rate limits, serialisers. | |
| 2 | + | |
| 3 | +Every route returns plain dicts; `AtlasJSONResponse` renders them with orjson (datetimes → ISO-8601 UTC, Decimal → float).""" | |
| 4 | +from __future__ import annotations | |
| 5 | + | |
| 6 | +import hmac | |
| 7 | +import inspect | |
| 8 | +import time | |
| 9 | +from collections import defaultdict, deque | |
| 10 | +from collections.abc import Awaitable, Callable | |
| 11 | +from datetime import UTC, date, datetime | |
| 12 | +from datetime import time as dtime | |
| 13 | +from decimal import Decimal | |
| 14 | +from functools import wraps | |
| 15 | +from typing import Any | |
| 16 | + | |
| 17 | +import orjson | |
| 18 | +from fastapi import Depends, HTTPException, Query, Request | |
| 19 | +from fastapi.responses import JSONResponse | |
| 20 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 21 | + | |
| 22 | +from aiatlas.config import settings | |
| 23 | +from aiatlas.db import fetch_one | |
| 24 | +from aiatlas.services import cache | |
| 25 | + | |
| 26 | +# ------------------------------------------------------------------------------------------------------------------ errors & JSON | |
| 27 | + | |
| 28 | + | |
| 29 | +class ApiError(HTTPException): | |
| 30 | + """Always `{"detail": "..."}` — never a stack trace.""" | |
| 31 | + | |
| 32 | + def __init__(self, status: int, detail: str): | |
| 33 | + super().__init__(status_code=status, detail=detail) | |
| 34 | + | |
| 35 | + | |
| 36 | +def _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) | |
| 44 | + | |
| 45 | + | |
| 46 | +_ORJSON_OPTS = orjson.OPT_NON_STR_KEYS | orjson.OPT_UTC_Z | orjson.OPT_NAIVE_UTC | |
| 47 | + | |
| 48 | + | |
| 49 | +def dumps(value: Any) -> bytes: | |
| 50 | + return orjson.dumps(value, default=_json_default, option=_ORJSON_OPTS) | |
| 51 | + | |
| 52 | + | |
| 53 | +def 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)) | |
| 56 | + | |
| 57 | + | |
| 58 | +class AtlasJSONResponse(JSONResponse): | |
| 59 | + media_type = "application/json" | |
| 60 | + | |
| 61 | + def render(self, content: Any) -> bytes: | |
| 62 | + return dumps(content) | |
| 63 | + | |
| 64 | + | |
| 65 | +# ------------------------------------------------------------------------------------------------------------------ pagination & parsing | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | +class Pagination: | |
| 70 | + def __init__(self, limit: int = Query(50, ge=1, le=200), offset: int = Query(0, ge=0)): | |
| 71 | + self.limit = limit | |
| 72 | + self.offset = offset | |
| 73 | + | |
| 74 | + | |
| 75 | +PAGINATION = Depends() # `p: Pagination = PAGINATION` — FastAPI resolves the class from the annotation | |
| 76 | + | |
| 77 | + | |
| 78 | +def page(items: list[Any], total: int, p: Pagination) -> dict[str, Any]: | |
| 79 | + return {"items": items, "total": int(total), "limit": p.limit, "offset": p.offset} | |
| 80 | + | |
| 81 | + | |
| 82 | +def parse_date(v: str | None, name: str = "date") -> date | None: | |
| 83 | + """YYYY-MM-DD (or ISO datetime) → `date`. asyncpg needs real date objects — never cast text in SQL.""" | |
| 84 | + if v is None or v == "": | |
| 85 | + return None | |
| 86 | + s = v.strip().replace("Z", "+00:00") | |
| 87 | + try: | |
| 88 | + if len(s) == 10: | |
| 89 | + return date.fromisoformat(s) | |
| 90 | + return datetime.fromisoformat(s).astimezone(UTC).date() | |
| 91 | + except ValueError as exc: | |
| 92 | + raise ApiError(400, f"{name} must be YYYY-MM-DD, got {v!r}") from exc | |
| 93 | + | |
| 94 | + | |
| 95 | +def parse_ts(v: str | None, name: str = "timestamp") -> datetime | None: | |
| 96 | + """ISO-8601 (date or datetime) → aware UTC `datetime`.""" | |
| 97 | + if v is None or v == "": | |
| 98 | + return None | |
| 99 | + s = v.strip().replace("Z", "+00:00") | |
| 100 | + try: | |
| 101 | + dt = datetime.fromisoformat(s) | |
| 102 | + except ValueError as exc: | |
| 103 | + raise ApiError(400, f"{name} must be ISO-8601, got {v!r}") from exc | |
| 104 | + if dt.tzinfo is None: | |
| 105 | + dt = datetime.combine(dt.date(), dt.time() or dtime.min, UTC) | |
| 106 | + return dt.astimezone(UTC) | |
| 107 | + | |
| 108 | + | |
| 109 | +def day_bounds(d: date) -> tuple[datetime, datetime]: | |
| 110 | + start = datetime.combine(d, dtime.min, UTC) | |
| 111 | + return start, datetime.combine(d, dtime.max, UTC).replace(microsecond=999999) | |
| 112 | + | |
| 113 | + | |
| 114 | +def csv(v: str | None) -> list[str]: | |
| 115 | + return [x.strip() for x in (v or "").split(",") if x.strip()] | |
| 116 | + | |
| 117 | + | |
| 118 | +def num_expr(path: str) -> str: | |
| 119 | + """Safe numeric cast of a JSON text attribute (`e.attributes->>'x'`).""" | |
| 120 | + return f"(case when {path} ~ '^-?[0-9]+(\\.[0-9]+)?$' then ({path})::double precision end)" | |
| 121 | + | |
| 122 | + | |
| 123 | +def attr_num(key: str, alias: str = "e") -> str: | |
| 124 | + return num_expr(f"{alias}.attributes->>'{key}'") | |
| 125 | + | |
| 126 | + | |
| 127 | +def flip_order(order_sql: str, order: str) -> str: | |
| 128 | + """Apply `order=asc|desc` to the primary sort key of a canned ORDER BY fragment (keeps `nulls last`).""" | |
| 129 | + if order not in ("asc", "desc"): | |
| 130 | + return order_sql | |
| 131 | + head, sep, tail = order_sql.partition(",") | |
| 132 | + if order == "asc" and " desc" in head: | |
| 133 | + head = head.replace(" desc", " asc", 1) | |
| 134 | + elif order == "desc" and " asc" in head: | |
| 135 | + head = head.replace(" asc", " desc", 1) | |
| 136 | + return head + sep + tail | |
| 137 | + | |
| 138 | + | |
| 139 | +# ------------------------------------------------------------------------------------------------------------------ response cache | |
| 140 | + | |
| 141 | + | |
| 142 | +def cache_key(request: Request) -> str: | |
| 143 | + items = sorted(request.query_params.multi_items()) | |
| 144 | + qs = "&".join(f"{k}={v}" for k, v in items) | |
| 145 | + return f"{request.url.path}?{qs}" if qs else request.url.path | |
| 146 | + | |
| 147 | + | |
| 148 | +def cached(ttl_s: int) -> Callable[[Callable[..., Awaitable[Any]]], Callable[..., Awaitable[Any]]]: | |
| 149 | + """Cache a GET route body in Redis (`aia:api:<path>?<sorted query>`). The route must accept `request: Request`.""" | |
| 150 | + | |
| 151 | + def deco(fn: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]: | |
| 152 | + @wraps(fn) | |
| 153 | + async def wrapper(*args: Any, **kwargs: Any) -> Any: | |
| 154 | + request = kwargs.get("request") or next((a for a in args if isinstance(a, Request)), None) | |
| 155 | + key = cache_key(request) if request is not None else None | |
| 156 | + if key is not None: | |
| 157 | + hit = await cache.cache_get(key) | |
| 158 | + if hit is not None: | |
| 159 | + return hit | |
| 160 | + value = normalize(await fn(*args, **kwargs)) | |
| 161 | + if key is not None: | |
| 162 | + await cache.cache_set(key, value, ttl_s) | |
| 163 | + return value | |
| 164 | + | |
| 165 | + wrapper.__signature__ = inspect.signature(fn) # type: ignore[attr-defined] | |
| 166 | + return wrapper | |
| 167 | + | |
| 168 | + return deco | |
| 169 | + | |
| 170 | + | |
| 171 | +# ------------------------------------------------------------------------------------------------------------------ admin auth | |
| 172 | + | |
| 173 | + | |
| 174 | +def require_admin(request: Request) -> None: | |
| 175 | + expected = settings.admin_token | |
| 176 | + if not expected: | |
| 177 | + raise ApiError(503, "admin API disabled: AIA_ADMIN_TOKEN is not configured") | |
| 178 | + token = request.headers.get("x-aia-admin-token") or "" | |
| 179 | + if not hmac.compare_digest(token, expected): | |
| 180 | + raise ApiError(401, "invalid or missing x-aia-admin-token") | |
| 181 | + | |
| 182 | + | |
| 183 | +# ------------------------------------------------------------------------------------------------------------------ rate limiting | |
| 184 | + | |
| 185 | +_LIMITS: dict[str, tuple[int, float]] = {"search": (60, 60.0), "views": (1, 1.0), "admin": (240, 60.0), "default": (600, 60.0)} | |
| 186 | +_hits: dict[tuple[str, str], deque[float]] = defaultdict(deque) | |
| 187 | + | |
| 188 | + | |
| 189 | +def client_ip(request: Request) -> str: | |
| 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" | |
| 194 | + | |
| 195 | + | |
| 196 | +def rate_limit(bucket: str) -> Callable[[Request], None]: | |
| 197 | + def dep(request: Request) -> None: | |
| 198 | + n, window = _LIMITS.get(bucket, _LIMITS["default"]) | |
| 199 | + now = time.monotonic() | |
| 200 | + q = _hits[(client_ip(request), bucket)] | |
| 201 | + while q and q[0] <= now - window: | |
| 202 | + q.popleft() | |
| 203 | + if len(q) >= n: | |
| 204 | + raise ApiError(429, f"rate limit exceeded for {bucket}: {n} per {int(window)} s") | |
| 205 | + q.append(now) | |
| 206 | + if len(_hits) > 20000: | |
| 207 | + for k in [k for k, v in _hits.items() if not v or v[-1] < now - 120]: | |
| 208 | + _hits.pop(k, None) | |
| 209 | + | |
| 210 | + return dep | |
| 211 | + | |
| 212 | + | |
| 213 | +# ------------------------------------------------------------------------------------------------------------------ entity serialisers | |
| 214 | + | |
| 215 | +ENTITY_FIELDS = ("id", "entity_type", "canonical_name", "slug", "description", "status", "organization_id", "attributes", "quality", "counts", | |
| 216 | + "first_seen_at", "last_seen_at", "updated_at", "organization_name", "organization_slug") | |
| 217 | + | |
| 218 | + | |
| 219 | +def entity_cols(alias: str = "e", prefix: str = "") -> str: | |
| 220 | + """Select-list for an entity alias (+ its organization alias `<alias>o`) with an optional column prefix.""" | |
| 221 | + org = f"{alias}o" | |
| 222 | + cols = [f"{alias}.{c} as {prefix}{c}" for c in ENTITY_FIELDS[:13]] | |
| 223 | + cols += [f"{org}.canonical_name as {prefix}organization_name", f"{org}.slug as {prefix}organization_slug"] | |
| 224 | + return ", ".join(cols) | |
| 225 | + | |
| 226 | + | |
| 227 | +def entity_join(alias: str, on: str) -> str: | |
| 228 | + return f"join entities {alias} on {alias}.id = {on} left join entities {alias}o on {alias}o.id = {alias}.organization_id" | |
| 229 | + | |
| 230 | + | |
| 231 | +ENTITY_COLS = entity_cols("e") | |
| 232 | +ENTITY_FROM = "entities e left join entities eo on eo.id = e.organization_id" | |
| 233 | + | |
| 234 | +SUMMARY_ATTRS: dict[str, tuple[str, ...]] = { | |
| 235 | + "model": ("family", "openness", "license", "parameter_count", "active_parameter_count", "context_length", "max_output_tokens", "modalities", | |
| 236 | + "release_date", "status", "knowledge_cutoff", "hf_repo", "api_model_id"), | |
| 237 | + "company": ("country", "founded", "website", "org_kind"), | |
| 238 | + "organization": ("country", "founded", "website", "org_kind"), | |
| 239 | + "lab": ("country", "founded", "website", "org_kind"), | |
| 240 | + "university": ("country", "founded", "website", "org_kind"), | |
| 241 | + "provider": ("website", "pricing_url"), | |
| 242 | + "benchmark": ("category", "metric", "unit"), | |
| 243 | + "hardware": ("kind", "memory_gb", "memory_bandwidth_gbs", "release_date", "manufacturer"), | |
| 244 | + "framework": ("latest_version", "latest_release_at", "license", "language", "metric.stars"), | |
| 245 | + "repository": ("latest_version", "latest_release_at", "license", "language", "metric.stars"), | |
| 246 | + "paper": ("authors", "published_at", "arxiv_id", "primary_category"), | |
| 247 | + "dataset": ("license", "modality"), | |
| 248 | +} | |
| 249 | +COMPANY_TYPES = ("company", "organization", "lab", "university") | |
| 250 | +TYPE_LABELS = {"model": "Models", "company": "Companies", "organization": "Organizations", "lab": "Labs", "university": "Universities", | |
| 251 | + "provider": "Providers", "paper": "Papers", "benchmark": "Benchmarks", "hardware": "Hardware", "framework": "Frameworks", | |
| 252 | + "dataset": "Datasets", "tool": "Tools", "repository": "Repositories", "regulation": "Regulation", "incident": "Incidents", | |
| 253 | + "researcher": "Researchers", "agent": "Agents", "product": "Products", "runtime": "Runtimes", "conference": "Conferences"} | |
| 254 | + | |
| 255 | + | |
| 256 | +def summary_attributes(entity_type: str, attrs: dict[str, Any] | None) -> dict[str, Any]: | |
| 257 | + attrs = attrs or {} | |
| 258 | + keys = SUMMARY_ATTRS.get(entity_type) | |
| 259 | + if keys is None: | |
| 260 | + return {k: v for k, v in attrs.items() if isinstance(v, (int, float, bool)) or (isinstance(v, str) and len(v) <= 200)} | |
| 261 | + out: dict[str, Any] = {} | |
| 262 | + for k in keys: | |
| 263 | + if k in attrs and attrs[k] not in (None, "", [], {}): | |
| 264 | + v = attrs[k] | |
| 265 | + out[k] = v[:5] if k == "authors" and isinstance(v, list) else v | |
| 266 | + return out | |
| 267 | + | |
| 268 | + | |
| 269 | +def org_of(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None: | |
| 270 | + oid = row.get(f"{prefix}organization_id") | |
| 271 | + if not oid: | |
| 272 | + return None | |
| 273 | + return {"id": oid, "slug": row.get(f"{prefix}organization_slug"), "name": row.get(f"{prefix}organization_name")} | |
| 274 | + | |
| 275 | + | |
| 276 | +def entity_summary(row: dict[str, Any], prefix: str = "") -> dict[str, Any] | None: | |
| 277 | + g = row.get | |
| 278 | + if not g(f"{prefix}id"): | |
| 279 | + return None | |
| 280 | + etype = g(f"{prefix}entity_type") or "" | |
| 281 | + desc = g(f"{prefix}description") | |
| 282 | + return {"id": g(f"{prefix}id"), "entity_type": etype, "slug": g(f"{prefix}slug"), "name": g(f"{prefix}canonical_name"), | |
| 283 | + "description": (desc[:280] if isinstance(desc, str) and len(desc) > 280 else desc), "status": g(f"{prefix}status"), | |
| 284 | + "organization": org_of(row, prefix), "attributes": summary_attributes(etype, g(f"{prefix}attributes")), | |
| 285 | + "quality": g(f"{prefix}quality") or {}, "counts": g(f"{prefix}counts") or {}, | |
| 286 | + "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")} | |
| 287 | + | |
| 288 | + | |
| 289 | +def change_event(row: dict[str, Any], prefix: str = "e_") -> dict[str, Any]: | |
| 290 | + return {"id": row["id"], "event_type": row["event_type"], "category": row["category"], "property": row.get("property"), | |
| 291 | + "old_value": row.get("old_value"), "new_value": row.get("new_value"), "summary": row.get("summary"), "importance": row.get("importance"), | |
| 292 | + "observed_at": row.get("observed_at"), "effective_at": row.get("effective_at"), "source_url": row.get("source_url"), | |
| 293 | + "connector_name": row.get("connector_name"), "entity": entity_summary(row, prefix), "meta": row.get("meta") or {}} | |
| 294 | + | |
| 295 | + | |
| 296 | +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, " | |
| 297 | + "ev.source_url, ev.connector_name, ev.meta, " + entity_cols("e", "e_")) | |
| 298 | +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" | |
| 299 | + | |
| 300 | + | |
| 301 | +def price_row(row: dict[str, Any]) -> dict[str, Any]: | |
| 302 | + return {"id": row["id"], "model": entity_summary(row, "m_"), "provider": entity_summary(row, "p_"), "provider_model_id": row.get("provider_model_id"), | |
| 303 | + "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"), | |
| 304 | + "cache_write_per_mtok": row.get("cache_write_per_mtok"), "batch_input_per_mtok": row.get("batch_input_per_mtok"), | |
| 305 | + "batch_output_per_mtok": row.get("batch_output_per_mtok"), "per_image": row.get("per_image"), "per_request": row.get("per_request"), | |
| 306 | + "currency": row.get("currency"), "context_length": row.get("context_length"), "max_output_tokens": row.get("max_output_tokens"), | |
| 307 | + "features": row.get("features") or {}, "observed_at": row.get("observed_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), | |
| 308 | + "source_url": row.get("source_url"), "tier": row.get("tier")} | |
| 309 | + | |
| 310 | + | |
| 311 | +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, " | |
| 312 | + "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, " | |
| 313 | + "p.valid_to, p.source_url, p.tier, " + entity_cols("m", "m_") + ", " + entity_cols("pv", "p_")) | |
| 314 | +PRICE_FROM = "prices p " + entity_join("m", "p.model_id") + " " + entity_join("pv", "p.provider_id") | |
| 315 | + | |
| 316 | + | |
| 317 | +def result_row(row: dict[str, Any]) -> dict[str, Any]: | |
| 318 | + return {"id": row["id"], "model": entity_summary(row, "m_"), "benchmark": entity_summary(row, "b_"), "score": row.get("score"), "metric": row.get("metric"), | |
| 319 | + "unit": row.get("unit"), "higher_is_better": row.get("higher_is_better"), "config": row.get("config") or {}, "evaluated_at": row.get("evaluated_at"), | |
| 320 | + "observed_at": row.get("observed_at"), "source_url": row.get("source_url"), "tier": row.get("tier"), "confidence": row.get("confidence"), | |
| 321 | + "valid_to": row.get("valid_to")} | |
| 322 | + | |
| 323 | + | |
| 324 | +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, " | |
| 325 | + + entity_cols("m", "m_") + ", " + entity_cols("b", "b_")) | |
| 326 | +RESULT_FROM = "benchmark_results r " + entity_join("m", "r.model_id") + " " + entity_join("b", "r.benchmark_id") | |
| 327 | +RESULT_ORDER = "case when r.higher_is_better then -r.score else r.score end, r.observed_at desc" | |
| 328 | + | |
| 329 | + | |
| 330 | +def claim_row(row: dict[str, Any]) -> dict[str, Any]: | |
| 331 | + return {"id": row["id"], "property": row["property"], "value": row.get("value"), "unit": row.get("unit"), "tier": row.get("tier"), | |
| 332 | + "confidence": row.get("confidence"), "status": row.get("status"), "extractor": row.get("extractor"), "observed_at": row.get("observed_at"), | |
| 333 | + "effective_at": row.get("effective_at"), "valid_from": row.get("valid_from"), "valid_to": row.get("valid_to"), "source_url": row.get("source_url"), | |
| 334 | + "source_name": row.get("source_name")} | |
| 335 | + | |
| 336 | + | |
| 337 | +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, " | |
| 338 | + "c.source_url, s.name as source_name") | |
| 339 | +CLAIM_FROM = "claims c left join sources s on s.id = c.source_id" | |
| 340 | + | |
| 341 | + | |
| 342 | +# ------------------------------------------------------------------------------------------------------------------ entity resolution | |
| 343 | + | |
| 344 | + | |
| 345 | +async def resolve_entity(conn: AsyncConnection, slug_or_id: str, types: tuple[str, ...] | None = None) -> dict[str, Any]: | |
| 346 | + """Slug first, then id; follows `merged_into`; 404 when absent or when the type does not match the mounted alias.""" | |
| 347 | + key = slug_or_id.strip() | |
| 348 | + if not key or len(key) > 200: | |
| 349 | + raise ApiError(404, "entity not found") | |
| 350 | + 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 " | |
| 351 | + f"order by case when e.slug = :k then 0 else 1 end limit 1", k=key) | |
| 352 | + hops = 0 | |
| 353 | + while row and row.get("merged_into") and hops < 5: | |
| 354 | + 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"]) | |
| 355 | + hops += 1 | |
| 356 | + if not row: | |
| 357 | + raise ApiError(404, "entity not found") | |
| 358 | + if types and row["entity_type"] not in types: | |
| 359 | + raise ApiError(404, f"entity {key!r} is a {row['entity_type']}, not one of {', '.join(types)}") | |
| 360 | + return row | |
| 361 | + | |
| 362 | + | |
| 363 | +async def resolve_id(conn: AsyncConnection, slug_or_id: str | None, types: tuple[str, ...] | None = None) -> str | None: | |
| 364 | + if not slug_or_id: | |
| 365 | + return None | |
| 366 | + return (await resolve_entity(conn, slug_or_id, types))["id"] | |
| 367 | + | |
| 368 | + | |
| 369 | +__all__ = [ | |
| 370 | + "CLAIM_COLS", | |
| 371 | + "CLAIM_FROM", | |
| 372 | + "COMPANY_TYPES", | |
| 373 | + "ENTITY_COLS", | |
| 374 | + "ENTITY_FROM", | |
| 375 | + "EVENT_COLS", | |
| 376 | + "EVENT_FROM", | |
| 377 | + "PAGINATION", | |
| 378 | + "PRICE_COLS", | |
| 379 | + "PRICE_FROM", | |
| 380 | + "RESULT_COLS", | |
| 381 | + "RESULT_FROM", | |
| 382 | + "RESULT_ORDER", | |
| 383 | + "TYPE_LABELS", | |
| 384 | + "ApiError", | |
| 385 | + "AtlasJSONResponse", | |
| 386 | + "Pagination", | |
| 387 | + "attr_num", | |
| 388 | + "cache_key", | |
| 389 | + "cached", | |
| 390 | + "change_event", | |
| 391 | + "claim_row", | |
| 392 | + "client_ip", | |
| 393 | + "csv", | |
| 394 | + "day_bounds", | |
| 395 | + "dumps", | |
| 396 | + "entity_cols", | |
| 397 | + "entity_join", | |
| 398 | + "entity_summary", | |
| 399 | + "flip_order", | |
| 400 | + "normalize", | |
| 401 | + "num_expr", | |
| 402 | + "org_of", | |
| 403 | + "page", | |
| 404 | + "parse_date", | |
| 405 | + "parse_ts", | |
| 406 | + "price_row", | |
| 407 | + "rate_limit", | |
| 408 | + "require_admin", | |
| 409 | + "resolve_entity", | |
| 410 | + "resolve_id", | |
| 411 | + "result_row", | |
| 412 | +] | |
added
src/aiatlas/api/detail.py
+258 −0
@@ -0,0 +1,258 @@ | ||
| 1 | +"""`EntityDetail` builder (docs/API.md): shared by `/entities/{slug}` and the type-scoped aliases.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 8 | + | |
| 9 | +from aiatlas.api.common import ( | |
| 10 | + COMPANY_TYPES, | |
| 11 | + ENTITY_COLS, | |
| 12 | + ENTITY_FROM, | |
| 13 | + EVENT_COLS, | |
| 14 | + EVENT_FROM, | |
| 15 | + PRICE_COLS, | |
| 16 | + PRICE_FROM, | |
| 17 | + RESULT_COLS, | |
| 18 | + RESULT_FROM, | |
| 19 | + RESULT_ORDER, | |
| 20 | + change_event, | |
| 21 | + entity_summary, | |
| 22 | + price_row, | |
| 23 | + result_row, | |
| 24 | +) | |
| 25 | +from aiatlas.db import connection, fetch_all | |
| 26 | +from aiatlas.services import hardware_fit as hf | |
| 27 | + | |
| 28 | +LINEAGE_PREDICATES = ("derived_from", "fine_tuned_from", "distilled_from", "merged_from", "quantized_from") | |
| 29 | +RELATION_GROUP_LIMIT = 24 | |
| 30 | + | |
| 31 | + | |
| 32 | +async def relations_grouped(conn: AsyncConnection, entity_id: str) -> list[dict[str, Any]]: | |
| 33 | + rows = await fetch_all(conn, f""" | |
| 34 | + with rel as ( | |
| 35 | + select r.predicate, 'out' as direction, r.object_id as other_id, r.observed_at from relations r where r.subject_id = :id and r.valid_to is null | |
| 36 | + union all | |
| 37 | + select r.predicate, 'in' as direction, r.subject_id as other_id, r.observed_at from relations r where r.object_id = :id and r.valid_to is null), | |
| 38 | + ranked as (select rel.*, row_number() over (partition by predicate, direction order by observed_at desc) as rn, | |
| 39 | + count(*) over (partition by predicate, direction) as total from rel) | |
| 40 | + select k.predicate, k.direction, k.total, {ENTITY_COLS} from ranked k join entities e on e.id = k.other_id left join entities eo on eo.id = e.organization_id | |
| 41 | + where k.rn <= :lim and e.merged_into is null order by k.predicate, k.direction, k.rn""", id=entity_id, lim=RELATION_GROUP_LIMIT) | |
| 42 | + groups: dict[tuple[str, str], dict[str, Any]] = {} | |
| 43 | + for r in rows: | |
| 44 | + g = groups.setdefault((r["predicate"], r["direction"]), {"predicate": r["predicate"], "direction": r["direction"], "items": [], "total": int(r["total"])}) | |
| 45 | + g["items"].append(entity_summary(r)) | |
| 46 | + return list(groups.values()) | |
| 47 | + | |
| 48 | + | |
| 49 | +async def sources_of(conn: AsyncConnection, entity_id: str, limit: int = 60) -> list[dict[str, Any]]: | |
| 50 | + """Documents attached to the entity + documents whose snapshots back its claims (deduped by URL — `documents.url` is unique).""" | |
| 51 | + rows = await fetch_all(conn, """ | |
| 52 | + with docs as (select d.id from documents d where d.entity_id = :id | |
| 53 | + union select s.document_id from claims c join snapshots s on s.id = c.snapshot_id where c.entity_id = :id | |
| 54 | + union select s.document_id from relations r join snapshots s on s.id = r.snapshot_id where r.subject_id = :id or r.object_id = :id) | |
| 55 | + select d.url, d.doc_type, d.title, greatest(d.last_fetched_at, d.last_changed_at) as last_observed_at, s.id as source_id, s.name as source_name, | |
| 56 | + s.domain, s.tier, (select count(*) from snapshots x where x.document_id = d.id) as snapshots | |
| 57 | + from docs join documents d on d.id = docs.id left join sources s on s.id = d.source_id | |
| 58 | + order by s.tier nulls last, last_observed_at desc nulls last limit :lim""", id=entity_id, lim=limit) | |
| 59 | + return [{"source_id": r["source_id"], "source_name": r["source_name"], "domain": r["domain"], "tier": r["tier"], "url": r["url"], "doc_type": r["doc_type"], | |
| 60 | + "title": r.get("title"), "last_observed_at": r["last_observed_at"], "snapshots": int(r["snapshots"] or 0)} for r in rows] | |
| 61 | + | |
| 62 | + | |
| 63 | +def _scope_sql(is_org: bool) -> str: | |
| 64 | + if is_org: | |
| 65 | + return ("(ev.entity_id = :id or ev.entity_id in (select id from entities where organization_id = :id union " | |
| 66 | + "select object_id from relations where subject_id = :id and predicate in ('develops','owns','operates','published') and valid_to is null))") | |
| 67 | + return "ev.entity_id = :id" | |
| 68 | + | |
| 69 | + | |
| 70 | +async def timeline_of(conn: AsyncConnection, entity_id: str, entity_type: str, *, limit: int = 30, before: Any = None, | |
| 71 | + include_documents: bool = False) -> list[dict[str, Any]]: | |
| 72 | + where = [_scope_sql(entity_type in COMPANY_TYPES)] | |
| 73 | + params: dict[str, Any] = {"id": entity_id, "lim": limit} | |
| 74 | + if before is not None: | |
| 75 | + where.append("ev.observed_at < :before") | |
| 76 | + params["before"] = before | |
| 77 | + if not include_documents: | |
| 78 | + where.append("ev.event_type <> 'DOCUMENT_CHANGED'") | |
| 79 | + rows = await fetch_all(conn, f"select {EVENT_COLS} from {EVENT_FROM} where {' and '.join(where)} order by ev.observed_at desc, ev.id desc limit :lim", **params) | |
| 80 | + return [change_event(r) for r in rows] | |
| 81 | + | |
| 82 | + | |
| 83 | +async def prices_of_model(conn: AsyncConnection, model_id: str, *, current_only: bool) -> list[dict[str, Any]]: | |
| 84 | + cond = "and p.valid_to is null" if current_only else "" | |
| 85 | + order = "p.input_per_mtok nulls last, pv.canonical_name" if current_only else "p.valid_from, p.id" | |
| 86 | + rows = await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where p.model_id = :id {cond} order by {order} limit 500", id=model_id) | |
| 87 | + return [price_row(r) for r in rows] | |
| 88 | + | |
| 89 | + | |
| 90 | +async def prices_of_provider(conn: AsyncConnection, provider_id: str) -> list[dict[str, Any]]: | |
| 91 | + rows = await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where p.provider_id = :id and p.valid_to is null order by m.canonical_name limit 500", id=provider_id) | |
| 92 | + return [price_row(r) for r in rows] | |
| 93 | + | |
| 94 | + | |
| 95 | +async def results_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]: | |
| 96 | + rows = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where r.model_id = :id and r.valid_to is null order by b.canonical_name, r.observed_at desc limit 300", id=model_id) | |
| 97 | + return [result_row(r) for r in rows] | |
| 98 | + | |
| 99 | + | |
| 100 | +async def leaderboard(conn: AsyncConnection, benchmark_id: str, *, limit: int = 100, offset: int = 0, config: str | None = None, history: bool = False) -> list[dict[str, Any]]: | |
| 101 | + where = ["r.benchmark_id = :id"] | |
| 102 | + params: dict[str, Any] = {"id": benchmark_id, "lim": limit, "off": offset} | |
| 103 | + if not history: | |
| 104 | + where.append("r.valid_to is null") | |
| 105 | + if config: | |
| 106 | + where.append("r.config::text ilike :cfg") | |
| 107 | + params["cfg"] = f"%{config}%" | |
| 108 | + rows = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where {' and '.join(where)} order by {RESULT_ORDER} limit :lim offset :off", **params) | |
| 109 | + return [result_row(r) for r in rows] | |
| 110 | + | |
| 111 | + | |
| 112 | +async def lineage_of(conn: AsyncConnection, model_id: str) -> dict[str, list[dict[str, Any]]]: | |
| 113 | + preds = list(LINEAGE_PREDICATES) | |
| 114 | + desc_preds = [p for p in preds if p != "quantized_from"] | |
| 115 | + ancestors = await fetch_all(conn, f""" | |
| 116 | + with recursive up as ( | |
| 117 | + select r.object_id as id, 1 as depth from relations r where r.subject_id = :id and r.valid_to is null and r.predicate = any(cast(:preds as text[])) | |
| 118 | + union | |
| 119 | + select r.object_id, up.depth + 1 from up join relations r on r.subject_id = up.id and r.valid_to is null and r.predicate = any(cast(:preds as text[])) where up.depth < 3) | |
| 120 | + select distinct on (e.id) up.depth, {ENTITY_COLS} from up join entities e on e.id = up.id left join entities eo on eo.id = e.organization_id | |
| 121 | + where e.id <> :id and e.merged_into is null order by e.id, up.depth""", id=model_id, preds=preds) | |
| 122 | + descendants = await fetch_all(conn, f""" | |
| 123 | + with recursive down as ( | |
| 124 | + select r.subject_id as id, 1 as depth from relations r where r.object_id = :id and r.valid_to is null and r.predicate = any(cast(:preds as text[])) | |
| 125 | + union | |
| 126 | + select r.subject_id, down.depth + 1 from down join relations r on r.object_id = down.id and r.valid_to is null and r.predicate = any(cast(:preds as text[])) where down.depth < 3) | |
| 127 | + select distinct on (e.id) down.depth, {ENTITY_COLS} from down join entities e on e.id = down.id left join entities eo on eo.id = e.organization_id | |
| 128 | + where e.id <> :id and e.merged_into is null order by e.id, down.depth""", id=model_id, preds=desc_preds) | |
| 129 | + quants = await fetch_all(conn, f"""select {ENTITY_COLS} from relations r join entities e on e.id = r.subject_id left join entities eo on eo.id = e.organization_id | |
| 130 | + where r.object_id = :id and r.predicate = 'quantized_from' and r.valid_to is null and e.merged_into is null | |
| 131 | + order by e.canonical_name limit 100""", id=model_id) | |
| 132 | + key = lambda r: (r.get("depth", 0), r["canonical_name"] or "") | |
| 133 | + return {"ancestors": [entity_summary(r) for r in sorted(ancestors, key=key)], "descendants": [entity_summary(r) for r in sorted(descendants, key=key)], | |
| 134 | + "quantizations": [entity_summary(r) for r in quants]} | |
| 135 | + | |
| 136 | + | |
| 137 | +async def providers_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]: | |
| 138 | + rows = await fetch_all(conn, f""" | |
| 139 | + with ids as (select r.object_id as id from relations r where r.subject_id = :id and r.predicate = 'available_through' and r.valid_to is null | |
| 140 | + union select p.provider_id from prices p where p.model_id = :id and p.valid_to is null) | |
| 141 | + select {ENTITY_COLS} from ids join entities e on e.id = ids.id left join entities eo on eo.id = e.organization_id | |
| 142 | + where e.merged_into is null order by e.canonical_name limit 100""", id=model_id) | |
| 143 | + return [entity_summary(r) for r in rows] | |
| 144 | + | |
| 145 | + | |
| 146 | +async def hardware_fit_of_model(conn: AsyncConnection, attrs: dict[str, Any]) -> list[dict[str, Any]] | None: | |
| 147 | + params = hf.parameter_count(attrs) | |
| 148 | + if params is None: | |
| 149 | + return None | |
| 150 | + context = 8192 | |
| 151 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'hardware' and e.merged_into is null and e.attributes ? 'memory_gb' order by e.canonical_name limit 200") | |
| 152 | + out: list[dict[str, Any]] = [] | |
| 153 | + for r in rows: | |
| 154 | + mem = hf.hardware_memory_gb(r["attributes"]) | |
| 155 | + if mem is None: | |
| 156 | + continue | |
| 157 | + f = hf.fit(params, mem, "4bit", context) | |
| 158 | + out.append({"hardware": entity_summary(r), "memory_gb": mem, **f}) | |
| 159 | + out.sort(key=lambda x: (not x["fits"], -x["headroom_gb"])) | |
| 160 | + return out | |
| 161 | + | |
| 162 | + | |
| 163 | +async def related_by_type(conn: AsyncConnection, entity_id: str, etype: str, *, limit: int = 50, offset: int = 0, include_org_children: bool = False) -> tuple[list[dict[str, Any]], int]: | |
| 164 | + """Entities of `etype` linked to `entity_id` by any live relation (either direction) — optionally also those whose organization is `entity_id`.""" | |
| 165 | + org_sql = "union select id from entities where organization_id = :id" if include_org_children else "" | |
| 166 | + rows = await fetch_all(conn, f""" | |
| 167 | + with ids as (select r.object_id as id from relations r where r.subject_id = :id and r.valid_to is null | |
| 168 | + union select r.subject_id from relations r where r.object_id = :id and r.valid_to is null {org_sql}) | |
| 169 | + select count(*) over () as total, {ENTITY_COLS} from ids join entities e on e.id = ids.id left join entities eo on eo.id = e.organization_id | |
| 170 | + where e.entity_type = :t and e.merged_into is null order by coalesce(e.attributes->>'release_date', e.attributes->>'published_at', '') desc, e.updated_at desc | |
| 171 | + limit :lim offset :off""", id=entity_id, t=etype, lim=limit, off=offset) | |
| 172 | + return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0 | |
| 173 | + | |
| 174 | + | |
| 175 | +async def entity_detail(row: dict[str, Any]) -> dict[str, Any]: | |
| 176 | + """Full detail. Blocks run on separate pooled connections concurrently (read-only, so no transaction needed).""" | |
| 177 | + eid, etype = row["id"], row["entity_type"] | |
| 178 | + | |
| 179 | + async def one(fn: Any, *args: Any, **kw: Any) -> Any: | |
| 180 | + async with connection() as conn: | |
| 181 | + return await fn(conn, *args, **kw) | |
| 182 | + | |
| 183 | + async def base(conn: AsyncConnection) -> dict[str, Any]: | |
| 184 | + aliases = await fetch_all(conn, "select alias from entity_aliases where entity_id = :id order by kind, alias limit 200", id=eid) | |
| 185 | + idents = await fetch_all(conn, "select scheme, value from entity_identifiers where entity_id = :id order by scheme, value limit 200", id=eid) | |
| 186 | + return {"aliases": [a["alias"] for a in aliases], "identifiers": [{"scheme": i["scheme"], "value": i["value"]} for i in idents]} | |
| 187 | + | |
| 188 | + tasks: dict[str, Any] = {"base": one(base), "relations": one(relations_grouped, eid), "sources": one(sources_of, eid), "timeline": one(timeline_of, eid, etype)} | |
| 189 | + if etype == "model": | |
| 190 | + tasks["prices"] = one(prices_of_model, eid, current_only=True) | |
| 191 | + tasks["price_history"] = one(prices_of_model, eid, current_only=False) | |
| 192 | + tasks["results"] = one(results_of_model, eid) | |
| 193 | + tasks["lineage"] = one(lineage_of, eid) | |
| 194 | + tasks["providers"] = one(providers_of_model, eid) | |
| 195 | + tasks["hardware_fit"] = one(hardware_fit_of_model, row["attributes"] or {}) | |
| 196 | + tasks["papers"] = one(related_by_type, eid, "paper", limit=24) | |
| 197 | + tasks["repositories"] = one(related_by_type, eid, "repository", limit=24) | |
| 198 | + elif etype in COMPANY_TYPES: | |
| 199 | + tasks["models"] = one(related_by_type, eid, "model", limit=50, include_org_children=True) | |
| 200 | + tasks["papers"] = one(related_by_type, eid, "paper", limit=24, include_org_children=True) | |
| 201 | + tasks["repositories"] = one(related_by_type, eid, "repository", limit=24, include_org_children=True) | |
| 202 | + elif etype == "provider": | |
| 203 | + tasks["prices"] = one(prices_of_provider, eid) | |
| 204 | + tasks["models"] = one(_provider_models, eid) | |
| 205 | + elif etype == "benchmark": | |
| 206 | + tasks["results"] = one(leaderboard, eid, limit=100) | |
| 207 | + elif etype == "hardware": | |
| 208 | + tasks["models"] = one(related_by_type, eid, "model", limit=50) | |
| 209 | + elif etype == "framework": | |
| 210 | + tasks["repositories"] = one(related_by_type, eid, "repository", limit=24) | |
| 211 | + | |
| 212 | + keys = list(tasks) | |
| 213 | + values = await asyncio.gather(*(tasks[k] for k in keys)) | |
| 214 | + blocks = dict(zip(keys, values, strict=True)) | |
| 215 | + | |
| 216 | + detail = entity_summary(row) or {} | |
| 217 | + detail["attributes"] = row.get("attributes") or {} | |
| 218 | + detail["provenance"] = row.get("provenance") or {} | |
| 219 | + detail.update(blocks.pop("base")) | |
| 220 | + detail["relations"] = blocks.pop("relations") | |
| 221 | + detail["sources"] = blocks.pop("sources") | |
| 222 | + detail["timeline"] = blocks.pop("timeline") | |
| 223 | + for k, v in blocks.items(): | |
| 224 | + if k in ("models",) and isinstance(v, tuple): | |
| 225 | + items, total = v | |
| 226 | + detail[k] = {"items": items, "total": total, "limit": 50, "offset": 0} | |
| 227 | + elif k in ("papers", "repositories") and isinstance(v, tuple): | |
| 228 | + detail[k] = v[0] | |
| 229 | + elif v is not None: | |
| 230 | + detail[k] = v | |
| 231 | + if etype == "model" and "hardware_fit" in detail: | |
| 232 | + detail["hardware_fit_assumptions"] = hf.ASSUMPTIONS | |
| 233 | + return detail | |
| 234 | + | |
| 235 | + | |
| 236 | +async def _provider_models(conn: AsyncConnection, provider_id: str) -> tuple[list[dict[str, Any]], int]: | |
| 237 | + rows = await fetch_all(conn, f""" | |
| 238 | + with ids as (select p.model_id as id from prices p where p.provider_id = :id and p.valid_to is null | |
| 239 | + union select r.subject_id from relations r where r.object_id = :id and r.predicate = 'available_through' and r.valid_to is null) | |
| 240 | + select count(*) over () as total, {ENTITY_COLS} from ids join entities e on e.id = ids.id left join entities eo on eo.id = e.organization_id | |
| 241 | + where e.merged_into is null order by e.canonical_name limit 50""", id=provider_id) | |
| 242 | + return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0 | |
| 243 | + | |
| 244 | + | |
| 245 | +__all__ = [ | |
| 246 | + "entity_detail", | |
| 247 | + "hardware_fit_of_model", | |
| 248 | + "leaderboard", | |
| 249 | + "lineage_of", | |
| 250 | + "prices_of_model", | |
| 251 | + "prices_of_provider", | |
| 252 | + "providers_of_model", | |
| 253 | + "related_by_type", | |
| 254 | + "relations_grouped", | |
| 255 | + "results_of_model", | |
| 256 | + "sources_of", | |
| 257 | + "timeline_of", | |
| 258 | +] | |
added
src/aiatlas/api/main.py
+137 −0
@@ -0,0 +1,137 @@ | ||
| 1 | +"""FastAPI application — `/api/v1`. Loopback-only in production (Next.js proxies `/api/v1/*`); docs at `/api/v1/docs`.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +import logging | |
| 6 | +from contextlib import asynccontextmanager | |
| 7 | +from datetime import UTC, datetime | |
| 8 | +from typing import Any | |
| 9 | + | |
| 10 | +from fastapi import FastAPI, Request | |
| 11 | +from fastapi.exceptions import RequestValidationError | |
| 12 | +from fastapi.middleware.cors import CORSMiddleware | |
| 13 | +from fastapi.middleware.gzip import GZipMiddleware | |
| 14 | +from starlette.exceptions import HTTPException as StarletteHTTPException | |
| 15 | + | |
| 16 | +import aiatlas | |
| 17 | +from aiatlas.api.common import AtlasJSONResponse | |
| 18 | +from aiatlas.api.routers import ( | |
| 19 | + admin, | |
| 20 | + benchmarks, | |
| 21 | + changes, | |
| 22 | + companies, | |
| 23 | + compare, | |
| 24 | + diff, | |
| 25 | + entities, | |
| 26 | + explore, | |
| 27 | + hardware, | |
| 28 | + misc, | |
| 29 | + models, | |
| 30 | + papers, | |
| 31 | + prices, | |
| 32 | + providers, | |
| 33 | + search, | |
| 34 | + sources, | |
| 35 | + stats, | |
| 36 | + timeline, | |
| 37 | +) | |
| 38 | +from aiatlas.config import settings | |
| 39 | +from aiatlas.db import connection, dispose, fetch_val | |
| 40 | +from aiatlas.logging import setup_logging | |
| 41 | +from aiatlas.services import cache | |
| 42 | +from aiatlas.services.llm import gateway | |
| 43 | + | |
| 44 | +log = logging.getLogger("aiatlas.api") | |
| 45 | + | |
| 46 | + | |
| 47 | +@asynccontextmanager | |
| 48 | +async def lifespan(app: FastAPI): # type: ignore[no-untyped-def] | |
| 49 | + setup_logging(service="aia-api") | |
| 50 | + settings.ensure_dirs() | |
| 51 | + log.info("api started", extra={"version": aiatlas.__version__, "env": settings.app_env, "port": settings.api_port}) | |
| 52 | + yield | |
| 53 | + await cache.close() | |
| 54 | + await dispose() | |
| 55 | + | |
| 56 | + | |
| 57 | +app = FastAPI(title="AI Atlas API", version=aiatlas.__version__, lifespan=lifespan, default_response_class=AtlasJSONResponse, | |
| 58 | + docs_url="/api/v1/docs", redoc_url=None, openapi_url="/api/v1/openapi.json", | |
| 59 | + description="The global intelligence layer for artificial intelligence: models, companies, papers, providers, prices, benchmarks, hardware — " | |
| 60 | + "with provenance and history on every fact. Contract: docs/API.md.") | |
| 61 | + | |
| 62 | +app.add_middleware(GZipMiddleware, minimum_size=1024) | |
| 63 | +_origins = {settings.site_url, "https://www.ai-atlas.co", "https://ai-atlas.co", "http://localhost:8320", "http://127.0.0.1:8320", "http://localhost:8330", "http://127.0.0.1:8330"} | |
| 64 | +app.add_middleware(CORSMiddleware, allow_origins=sorted(o for o in _origins if o), allow_methods=["GET", "POST", "PATCH", "OPTIONS"], | |
| 65 | + allow_headers=["*"], max_age=600) | |
| 66 | + | |
| 67 | + | |
| 68 | +@app.middleware("http") | |
| 69 | +async def security_headers(request: Request, call_next): # type: ignore[no-untyped-def] | |
| 70 | + try: | |
| 71 | + response = await call_next(request) | |
| 72 | + except Exception: | |
| 73 | + log.exception("unhandled error", extra={"route": request.url.path}) | |
| 74 | + return AtlasJSONResponse({"detail": "internal server error"}, status_code=500) | |
| 75 | + response.headers["x-content-type-options"] = "nosniff" | |
| 76 | + response.headers["referrer-policy"] = "strict-origin-when-cross-origin" | |
| 77 | + return response | |
| 78 | + | |
| 79 | + | |
| 80 | +@app.exception_handler(StarletteHTTPException) | |
| 81 | +async def http_error(request: Request, exc: StarletteHTTPException): # type: ignore[no-untyped-def] | |
| 82 | + detail = exc.detail if isinstance(exc.detail, str) else str(exc.detail) | |
| 83 | + return AtlasJSONResponse({"detail": detail}, status_code=exc.status_code, headers=dict(exc.headers or {})) | |
| 84 | + | |
| 85 | + | |
| 86 | +@app.exception_handler(RequestValidationError) | |
| 87 | +async def validation_error(request: Request, exc: RequestValidationError): # type: ignore[no-untyped-def] | |
| 88 | + errs = exc.errors()[:5] | |
| 89 | + msg = "; ".join(f"{'.'.join(str(x) for x in e.get('loc', []) if x not in ('query', 'body', 'path'))}: {e.get('msg')}" for e in errs) or "invalid request" | |
| 90 | + return AtlasJSONResponse({"detail": msg, "errors": [{"loc": e.get("loc"), "msg": e.get("msg"), "type": e.get("type")} for e in errs]}, status_code=422) | |
| 91 | + | |
| 92 | + | |
| 93 | +@app.exception_handler(Exception) | |
| 94 | +async def unhandled(request: Request, exc: Exception): # type: ignore[no-untyped-def] | |
| 95 | + log.exception("unhandled error", extra={"route": request.url.path}) | |
| 96 | + return AtlasJSONResponse({"detail": "internal server error"}, status_code=500) | |
| 97 | + | |
| 98 | + | |
| 99 | +async def _health() -> dict[str, Any]: | |
| 100 | + db_ok = redis_ok = False | |
| 101 | + try: | |
| 102 | + async with connection() as conn: | |
| 103 | + db_ok = (await asyncio.wait_for(fetch_val(conn, "select 1"), timeout=3)) == 1 | |
| 104 | + except Exception: # noqa: BLE001 | |
| 105 | + db_ok = False | |
| 106 | + try: | |
| 107 | + redis_ok = bool(await asyncio.wait_for(cache.redis().ping(), timeout=2)) | |
| 108 | + except Exception: # noqa: BLE001 | |
| 109 | + redis_ok = False | |
| 110 | + llm: dict[str, Any] = {"available": gateway.available} | |
| 111 | + if gateway.available: | |
| 112 | + cached = await cache.cache_get("health:llm") | |
| 113 | + if cached is None: | |
| 114 | + try: | |
| 115 | + cached = {"reachable": bool(await asyncio.wait_for(gateway.engine.health(), timeout=2.5))} | |
| 116 | + except Exception: # noqa: BLE001 | |
| 117 | + cached = {"reachable": False} | |
| 118 | + await cache.cache_set("health:llm", cached, 300) | |
| 119 | + llm.update(cached) | |
| 120 | + return {"status": "ok" if db_ok else "degraded", "version": aiatlas.__version__, "db": db_ok, "redis": redis_ok, "llm": llm, "time": datetime.now(UTC)} | |
| 121 | + | |
| 122 | + | |
| 123 | +@app.get("/health", tags=["health"]) | |
| 124 | +async def health_root() -> dict[str, Any]: | |
| 125 | + return await _health() | |
| 126 | + | |
| 127 | + | |
| 128 | +@app.get("/api/v1/health", tags=["health"]) | |
| 129 | +async def health_v1() -> dict[str, Any]: | |
| 130 | + return await _health() | |
| 131 | + | |
| 132 | + | |
| 133 | +# Order matters only where a literal path and a `{param}` path share a prefix — literal routes live in the same router and are declared first. | |
| 134 | +for r in (stats, search, models, companies, papers, providers, prices, benchmarks, hardware, explore, changes, timeline, compare, diff, sources, misc, entities, admin): | |
| 135 | + app.include_router(r.router) | |
| 136 | + | |
| 137 | +__all__ = ["app"] | |
added
src/aiatlas/api/routers/__init__.py
+1 −0
@@ -0,0 +1 @@ | ||
| 1 | +"""One router per docs/API.md group, all mounted under `/api/v1`.""" | |
added
src/aiatlas/api/routers/admin.py
+539 −0
@@ -0,0 +1,539 @@ | ||
| 1 | +"""Admin API (`x-aia-admin-token`): overview, connectors, runs/errors, documents & snapshots (cleaned text only), queues, LLM | |
| 2 | +accounting, review queue with entity merging, duplicates, curation, infrastructure, cache. Never returns raw archive paths.""" | |
| 3 | +from __future__ import annotations | |
| 4 | + | |
| 5 | +import platform | |
| 6 | +import socket | |
| 7 | +import time | |
| 8 | +from datetime import UTC, datetime | |
| 9 | +from typing import Any | |
| 10 | + | |
| 11 | +from fastapi import APIRouter, Depends, Query | |
| 12 | +from pydantic import BaseModel, Field | |
| 13 | + | |
| 14 | +from aiatlas.api.common import PAGINATION, ApiError, Pagination, page, rate_limit, require_admin | |
| 15 | +from aiatlas.config import settings | |
| 16 | +from aiatlas.connectors import registry as connector_registry | |
| 17 | +from aiatlas.db import connection, execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 18 | +from aiatlas.ids import new_id | |
| 19 | +from aiatlas.sdk import archive | |
| 20 | +from aiatlas.services import cache | |
| 21 | +from aiatlas.services.jobs import enqueue, queue_depth | |
| 22 | +from aiatlas.services.llm import gateway | |
| 23 | +from aiatlas.services.merge import merge_entities | |
| 24 | +from aiatlas.services.stats import live_counts | |
| 25 | + | |
| 26 | +router = APIRouter(prefix="/api/v1/admin", tags=["admin"], dependencies=[Depends(require_admin), Depends(rate_limit("admin"))]) | |
| 27 | +_STARTED = time.time() | |
| 28 | +SNAPSHOT_TEXT_LIMIT = 20 * 1024 | |
| 29 | +HIDDEN = ("raw_path", "text_path") | |
| 30 | + | |
| 31 | + | |
| 32 | +def _public(row: dict[str, Any]) -> dict[str, Any]: | |
| 33 | + return {k: v for k, v in row.items() if k not in HIDDEN} | |
| 34 | + | |
| 35 | + | |
| 36 | +# ------------------------------------------------------------------------------------------------------------------ overview | |
| 37 | + | |
| 38 | + | |
| 39 | +@router.get("/overview") | |
| 40 | +async def overview() -> dict[str, Any]: | |
| 41 | + async with connection() as conn: | |
| 42 | + stats = await live_counts(conn) | |
| 43 | + queue = await queue_depth(conn) | |
| 44 | + health = await fetch_all(conn, "select health, count(*) as n from connectors group by 1") | |
| 45 | + errors = await fetch_val(conn, "select count(*) from connector_errors where created_at > now() - interval '24 hours'") | |
| 46 | + llm = await fetch_one(conn, """select count(*) as jobs_24h, coalesce(sum(input_tokens), 0) + coalesce(sum(output_tokens), 0) as tokens_24h, | |
| 47 | + count(*) filter (where status <> 'ok') as failed_24h from llm_jobs where created_at > now() - interval '24 hours'""") | |
| 48 | + by_stage = await fetch_all(conn, "select stage, status, count(*) as n from llm_jobs where created_at > now() - interval '24 hours' group by 1, 2 order by 1, 2") | |
| 49 | + recent_runs = await fetch_all(conn, "select * from connector_runs order by started_at desc limit 12") | |
| 50 | + review_kinds = await fetch_all(conn, "select kind, count(*) as n from review_queue where status = 'pending' group by 1") | |
| 51 | + counts = {h["health"]: int(h["n"]) for h in health} | |
| 52 | + return {"stats": stats, "queue": queue, "heartbeats": await cache.heartbeats(), | |
| 53 | + "connectors": {k: counts.get(k, 0) for k in ("ok", "degraded", "failing", "disabled", "unknown")}, "review_pending": stats.get("review_pending", 0), | |
| 54 | + "review_by_kind": {r["kind"]: int(r["n"]) for r in review_kinds}, "recent_errors": int(errors or 0), | |
| 55 | + "llm": {"available": gateway.available, "jobs_24h": int(llm["jobs_24h"] or 0), "tokens_24h": int(llm["tokens_24h"] or 0), "failed_24h": int(llm["failed_24h"] or 0), | |
| 56 | + "by_stage": [{**r, "n": int(r["n"])} for r in by_stage]}, | |
| 57 | + "recent_runs": recent_runs, "archive": archive.archive_size(), "computed_at": datetime.now(UTC)} | |
| 58 | + | |
| 59 | + | |
| 60 | +# ------------------------------------------------------------------------------------------------------------------ connectors | |
| 61 | + | |
| 62 | + | |
| 63 | +@router.get("/connectors") | |
| 64 | +async def connectors() -> dict[str, Any]: | |
| 65 | + async with connection() as conn: | |
| 66 | + rows = await fetch_all(conn, """ | |
| 67 | + select c.*, s.key as source_key, s.name as source_name, s.tier as source_tier, s.domain as source_domain, | |
| 68 | + (select row_to_json(r) from connector_runs r where r.connector_name = c.name order by r.started_at desc limit 1) as last_run, | |
| 69 | + (select count(*) from documents d where d.connector_name = c.name) as documents, | |
| 70 | + (select count(*) from snapshots x join documents d on d.id = x.document_id where d.connector_name = c.name) as snapshots, | |
| 71 | + (select count(*) from connector_errors e where e.connector_name = c.name and e.created_at > now() - interval '7 days') as errors_7d | |
| 72 | + from connectors c left join sources s on s.id = c.source_id order by c.priority, c.name""") | |
| 73 | + known = connector_registry() | |
| 74 | + return {"items": [{**r, "documents": int(r["documents"]), "snapshots": int(r["snapshots"]), "errors_7d": int(r["errors_7d"]), "in_code": r["name"] in known, | |
| 75 | + "run_now_pending": await _run_now_pending(r["name"])} for r in rows], | |
| 76 | + "unregistered_in_db": sorted(set(known) - {r["name"] for r in rows})} | |
| 77 | + | |
| 78 | + | |
| 79 | +async def _run_now_pending(name: str) -> bool: | |
| 80 | + try: | |
| 81 | + return bool(await cache.redis().exists(f"aia:run-now:{name}")) | |
| 82 | + except Exception: # noqa: BLE001 | |
| 83 | + return False | |
| 84 | + | |
| 85 | + | |
| 86 | +class RunBody(BaseModel): | |
| 87 | + force: bool = False | |
| 88 | + | |
| 89 | + | |
| 90 | +@router.post("/connectors/{name}/run") | |
| 91 | +async def run_connector(name: str, body: RunBody | None = None) -> dict[str, Any]: | |
| 92 | + async with connection() as conn: | |
| 93 | + exists = await fetch_one(conn, "select name from connectors where name = :n", n=name) | |
| 94 | + if not exists and name not in connector_registry(): | |
| 95 | + raise ApiError(404, f"unknown connector {name!r}") | |
| 96 | + try: | |
| 97 | + await cache.redis().set(f"aia:run-now:{name}", b"force" if (body and body.force) else b"1", ex=6 * 3600) | |
| 98 | + except Exception as exc: | |
| 99 | + raise ApiError(503, f"redis unavailable: {exc.__class__.__name__}") from exc | |
| 100 | + return {"queued": True, "connector": name, "force": bool(body and body.force), "note": "consumed by the scheduler tick (aia:run-now:<name>)"} | |
| 101 | + | |
| 102 | + | |
| 103 | +class ConnectorPatch(BaseModel): | |
| 104 | + enabled: bool | None = None | |
| 105 | + interval_seconds: int | None = Field(None, ge=60, le=30 * 86400) | |
| 106 | + priority: int | None = Field(None, ge=0, le=9) | |
| 107 | + | |
| 108 | + | |
| 109 | +@router.patch("/connectors/{name}") | |
| 110 | +async def patch_connector(name: str, body: ConnectorPatch) -> dict[str, Any]: | |
| 111 | + sets, params = [], {"n": name} | |
| 112 | + if body.enabled is not None: | |
| 113 | + sets.append("enabled = :enabled") | |
| 114 | + sets.append("health = case when :enabled then (case when health = 'disabled' then 'unknown' else health end) else 'disabled' end") | |
| 115 | + sets.append("circuit_open_until = case when :enabled then null else circuit_open_until end") | |
| 116 | + params["enabled"] = body.enabled | |
| 117 | + if body.interval_seconds is not None: | |
| 118 | + sets.append("interval_seconds = :iv") | |
| 119 | + sets.append("next_run_at = least(coalesce(next_run_at, now()), coalesce(last_success_at, now()) + make_interval(secs => :iv))") | |
| 120 | + params["iv"] = body.interval_seconds | |
| 121 | + if body.priority is not None: | |
| 122 | + sets.append("priority = :pr") | |
| 123 | + params["pr"] = body.priority | |
| 124 | + if not sets: | |
| 125 | + raise ApiError(400, "nothing to update") | |
| 126 | + async with transaction() as conn: | |
| 127 | + row = await fetch_one(conn, f"update connectors set {', '.join(sets)}, updated_at = now() where name = :n returning *", **params) | |
| 128 | + if not row: | |
| 129 | + raise ApiError(404, f"unknown connector {name!r}") | |
| 130 | + return row | |
| 131 | + | |
| 132 | + | |
| 133 | +@router.get("/runs") | |
| 134 | +async def runs(connector: str | None = None, status: str | None = None, limit: int = Query(50, ge=1, le=500), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 135 | + where = ["true"] | |
| 136 | + params: dict[str, Any] = {"lim": limit, "off": offset} | |
| 137 | + if connector: | |
| 138 | + where.append("connector_name = :c") | |
| 139 | + params["c"] = connector | |
| 140 | + if status: | |
| 141 | + where.append("status = :s") | |
| 142 | + params["s"] = status | |
| 143 | + async with connection() as conn: | |
| 144 | + rows = await fetch_all(conn, f"select * from connector_runs where {' and '.join(where)} order by started_at desc limit :lim offset :off", **params) | |
| 145 | + total = await fetch_val(conn, f"select count(*) from connector_runs where {' and '.join(where)}", **{k: v for k, v in params.items() if k not in ('lim', 'off')}) | |
| 146 | + return {"items": rows, "total": int(total or 0), "limit": limit, "offset": offset} | |
| 147 | + | |
| 148 | + | |
| 149 | +@router.get("/errors") | |
| 150 | +async def errors(connector: str | None = None, limit: int = Query(100, ge=1, le=1000), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 151 | + where = "connector_name = :c" if connector else "true" | |
| 152 | + async with connection() as conn: | |
| 153 | + rows = await fetch_all(conn, f"select * from connector_errors where {where} order by created_at desc limit :lim offset :off", c=connector, lim=limit, off=offset) | |
| 154 | + total = await fetch_val(conn, f"select count(*) from connector_errors where {where}", c=connector) | |
| 155 | + return {"items": rows, "total": int(total or 0), "limit": limit, "offset": offset} | |
| 156 | + | |
| 157 | + | |
| 158 | +# ------------------------------------------------------------------------------------------------------------------ documents & snapshots | |
| 159 | + | |
| 160 | + | |
| 161 | +@router.get("/documents") | |
| 162 | +async def documents(connector: str | None = None, status: str | None = None, q: str | None = Query(None, max_length=300), entity: str | None = None, | |
| 163 | + needs_llm: int | None = Query(None, ge=0, le=1), p: Pagination = PAGINATION) -> dict[str, Any]: | |
| 164 | + where = ["true"] | |
| 165 | + params: dict[str, Any] = {} | |
| 166 | + if connector: | |
| 167 | + where.append("d.connector_name = :c") | |
| 168 | + params["c"] = connector | |
| 169 | + if status: | |
| 170 | + where.append("d.status = :s") | |
| 171 | + params["s"] = status | |
| 172 | + if q: | |
| 173 | + where.append("(d.url ilike :q or d.title ilike :q)") | |
| 174 | + params["q"] = f"%{q}%" | |
| 175 | + if entity: | |
| 176 | + where.append("(d.entity_id = :e or e.slug = :e)") | |
| 177 | + params["e"] = entity | |
| 178 | + if needs_llm is not None: | |
| 179 | + where.append("d.needs_llm = :nl") | |
| 180 | + params["nl"] = bool(needs_llm) | |
| 181 | + where_sql = " and ".join(where) | |
| 182 | + async with connection() as conn: | |
| 183 | + rows = await fetch_all(conn, f"""select d.*, e.slug as entity_slug, e.canonical_name as entity_name, e.entity_type, | |
| 184 | + (select count(*) from snapshots x where x.document_id = d.id) as snapshots, | |
| 185 | + (select x.processing_status from snapshots x where x.document_id = d.id order by x.observed_at desc limit 1) as last_processing_status | |
| 186 | + from documents d left join entities e on e.id = d.entity_id where {where_sql} | |
| 187 | + order by d.last_fetched_at desc nulls last, d.first_seen_at desc limit :lim offset :off""", lim=p.limit, off=p.offset, **params) | |
| 188 | + total = await fetch_val(conn, f"select count(*) from documents d left join entities e on e.id = d.entity_id where {where_sql}", **params) | |
| 189 | + return page([{**r, "snapshots": int(r["snapshots"])} for r in rows], int(total or 0), p) | |
| 190 | + | |
| 191 | + | |
| 192 | +@router.get("/documents/{doc_id}") | |
| 193 | +async def document(doc_id: str) -> dict[str, Any]: | |
| 194 | + async with connection() as conn: | |
| 195 | + doc = await fetch_one(conn, """select d.*, e.slug as entity_slug, e.canonical_name as entity_name, e.entity_type, s.name as source_name, s.tier as source_tier | |
| 196 | + from documents d left join entities e on e.id = d.entity_id left join sources s on s.id = d.source_id where d.id = :id or d.url = :id limit 1""", id=doc_id) | |
| 197 | + if not doc: | |
| 198 | + raise ApiError(404, "document not found") | |
| 199 | + snaps = await fetch_all(conn, """select id, run_id, url, final_url, observed_at, http_status, content_type, content_hash, byte_size, text_hash, parser_version, connector_version, | |
| 200 | + transport, changed, processing_status, created_at, structured is not null as has_structured, diff is not null as has_diff, text_path is not null as has_text | |
| 201 | + from snapshots where document_id = :id order by observed_at desc limit 200""", id=doc["id"]) | |
| 202 | + claims = await fetch_val(conn, "select count(*) from claims c join snapshots s on s.id = c.snapshot_id where s.document_id = :id", id=doc["id"]) | |
| 203 | + return {**doc, "snapshots": snaps, "claims_from_document": int(claims or 0)} | |
| 204 | + | |
| 205 | + | |
| 206 | +@router.get("/snapshots/{snap_id}") | |
| 207 | +async def snapshot(snap_id: str, text_limit: int = Query(SNAPSHOT_TEXT_LIMIT, ge=0, le=SNAPSHOT_TEXT_LIMIT)) -> dict[str, Any]: | |
| 208 | + async with connection() as conn: | |
| 209 | + snap = await fetch_one(conn, """select s.*, d.url as document_url, d.doc_type, d.connector_name, d.entity_id, e.slug as entity_slug, e.canonical_name as entity_name | |
| 210 | + from snapshots s join documents d on d.id = s.document_id left join entities e on e.id = d.entity_id where s.id = :id""", id=snap_id) | |
| 211 | + if not snap: | |
| 212 | + raise ApiError(404, "snapshot not found") | |
| 213 | + claims = await fetch_all(conn, """select c.id, c.entity_id, e.slug as entity_slug, c.property, c.value, c.unit, c.status, c.confidence, c.extractor, c.observed_at | |
| 214 | + from claims c left join entities e on e.id = c.entity_id where c.snapshot_id = :id order by e.slug, c.property limit 500""", id=snap_id) | |
| 215 | + events = await fetch_all(conn, "select id, entity_id, event_type, category, property, summary, importance, observed_at from change_events where snapshot_id = :id order by observed_at desc limit 200", id=snap_id) | |
| 216 | + llm = await fetch_all(conn, "select id, task_type, stage, model, status, input_tokens, output_tokens, duration_ms, error, created_at from llm_jobs where snapshot_id = :id order by created_at desc limit 50", id=snap_id) | |
| 217 | + text, text_error, text_len = None, None, 0 | |
| 218 | + if snap.get("text_path") and text_limit: | |
| 219 | + try: | |
| 220 | + full = archive.load_text(snap["text_path"]) | |
| 221 | + text_len = len(full) | |
| 222 | + text = full[:text_limit] | |
| 223 | + except OSError as exc: | |
| 224 | + text_error = f"cleaned text unavailable ({exc.__class__.__name__})" | |
| 225 | + out = _public(snap) | |
| 226 | + out.update({"text": text, "text_chars": text_len, "text_truncated": text_len > text_limit, "text_error": text_error, "has_raw": bool(snap.get("raw_path")), | |
| 227 | + "claims": claims, "events": events, "llm_jobs": llm}) | |
| 228 | + return out | |
| 229 | + | |
| 230 | + | |
| 231 | +# ------------------------------------------------------------------------------------------------------------------ queues | |
| 232 | + | |
| 233 | + | |
| 234 | +@router.get("/jobs") | |
| 235 | +async def jobs(status: str | None = None, kind: str | None = None, limit: int = Query(100, ge=1, le=1000), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 236 | + where = ["true"] | |
| 237 | + params: dict[str, Any] = {} | |
| 238 | + if status: | |
| 239 | + where.append("status = any(cast(:s as text[]))") | |
| 240 | + params["s"] = [x.strip() for x in status.split(",") if x.strip()] | |
| 241 | + if kind: | |
| 242 | + where.append("kind = :k") | |
| 243 | + params["k"] = kind | |
| 244 | + where_sql = " and ".join(where) | |
| 245 | + async with connection() as conn: | |
| 246 | + rows = await fetch_all(conn, f"select * from jobs where {where_sql} order by case status when 'running' then 0 when 'queued' then 1 when 'failed' then 2 when 'dead' then 3 else 4 end, " | |
| 247 | + f"priority, coalesce(finished_at, run_after) desc limit :lim offset :off", lim=limit, off=offset, **params) | |
| 248 | + total = await fetch_val(conn, f"select count(*) from jobs where {where_sql}", **params) | |
| 249 | + depth = await queue_depth(conn) | |
| 250 | + return {"items": rows, "total": int(total or 0), "limit": limit, "offset": offset, "depth": depth} | |
| 251 | + | |
| 252 | + | |
| 253 | +@router.post("/jobs/{job_id}/retry") | |
| 254 | +async def retry_job(job_id: str) -> dict[str, Any]: | |
| 255 | + async with transaction() as conn: | |
| 256 | + row = await fetch_one(conn, """update jobs set status = 'queued', attempts = 0, error = null, locked_by = null, locked_at = null, finished_at = null, run_after = now() | |
| 257 | + where id = :id and status in ('failed', 'dead', 'done') returning id, kind, status""", id=job_id) | |
| 258 | + if not row: | |
| 259 | + raise ApiError(404, "job not found or not retryable (must be failed, dead or done)") | |
| 260 | + return {"ok": True, **row} | |
| 261 | + | |
| 262 | + | |
| 263 | +@router.post("/jobs/requeue-dead") | |
| 264 | +async def requeue_dead(kind: str | None = None) -> dict[str, Any]: | |
| 265 | + where = "status = 'dead'" + (" and kind = :k" if kind else "") | |
| 266 | + async with transaction() as conn: | |
| 267 | + n = await fetch_val(conn, f"with u as (update jobs set status = 'queued', attempts = 0, error = null, locked_by = null, locked_at = null, finished_at = null, run_after = now() " | |
| 268 | + f"where {where} returning 1) select count(*) from u", k=kind) | |
| 269 | + return {"requeued": int(n or 0)} | |
| 270 | + | |
| 271 | + | |
| 272 | +@router.get("/llm-jobs") | |
| 273 | +async def llm_jobs(limit: int = Query(100, ge=1, le=1000), offset: int = Query(0, ge=0), status: str | None = None, task: str | None = None) -> dict[str, Any]: | |
| 274 | + where = ["true"] | |
| 275 | + params: dict[str, Any] = {} | |
| 276 | + if status: | |
| 277 | + where.append("status = :s") | |
| 278 | + params["s"] = status | |
| 279 | + if task: | |
| 280 | + where.append("task_type = :t") | |
| 281 | + params["t"] = task | |
| 282 | + where_sql = " and ".join(where) | |
| 283 | + async with connection() as conn: | |
| 284 | + rows = await fetch_all(conn, f"select id, job_id, task_type, stage, engine, model, node, schema_name, snapshot_id, entity_id, input_tokens, output_tokens, duration_ms, status, error, created_at " | |
| 285 | + f"from llm_jobs where {where_sql} order by created_at desc limit :lim offset :off", lim=limit, off=offset, **params) | |
| 286 | + total = await fetch_val(conn, f"select count(*) from llm_jobs where {where_sql}", **params) | |
| 287 | + totals = await fetch_all(conn, """select stage, model, status, count(*) as n, coalesce(sum(input_tokens), 0) as input_tokens, coalesce(sum(output_tokens), 0) as output_tokens, | |
| 288 | + coalesce(avg(duration_ms), 0)::int as avg_ms from llm_jobs group by 1, 2, 3 order by 1, 2, 3""") | |
| 289 | + return {"items": rows, "total": int(total or 0), "limit": limit, "offset": offset, "totals": [{**t, "n": int(t["n"]), "input_tokens": int(t["input_tokens"]), "output_tokens": int(t["output_tokens"])} for t in totals]} | |
| 290 | + | |
| 291 | + | |
| 292 | +@router.get("/llm/health") | |
| 293 | +async def llm_health() -> dict[str, Any]: | |
| 294 | + h = await gateway.health() | |
| 295 | + h["base_url_configured"] = bool(settings.llm_base_url) | |
| 296 | + return h | |
| 297 | + | |
| 298 | + | |
| 299 | +class LLMEnqueue(BaseModel): | |
| 300 | + limit: int = Field(200, ge=1, le=5000) | |
| 301 | + task: str | None = None | |
| 302 | + connector: str | None = None | |
| 303 | + | |
| 304 | + | |
| 305 | +@router.post("/llm/enqueue") | |
| 306 | +async def llm_enqueue(body: LLMEnqueue | None = None) -> dict[str, Any]: | |
| 307 | + body = body or LLMEnqueue() | |
| 308 | + where = "s.processing_status in ('stored', 'extracted', 'llm_pending') and s.text_path is not null and (d.needs_llm or s.processing_status = 'llm_pending')" | |
| 309 | + params: dict[str, Any] = {"lim": body.limit} | |
| 310 | + if body.connector: | |
| 311 | + where += " and d.connector_name = :c" | |
| 312 | + params["c"] = body.connector | |
| 313 | + queued = 0 | |
| 314 | + async with transaction() as conn: | |
| 315 | + snaps = await fetch_all(conn, f"select s.id from snapshots s join documents d on d.id = s.document_id where {where} order by d.priority, s.observed_at desc limit :lim", **params) | |
| 316 | + for s in snaps: | |
| 317 | + payload = {"snapshot_id": s["id"], **({"task": body.task} if body.task else {})} | |
| 318 | + if await enqueue(conn, "llm_extract", payload, priority=6, dedupe_key=f"llm:{s['id']}"): | |
| 319 | + queued += 1 | |
| 320 | + if snaps: | |
| 321 | + await execute(conn, "update snapshots set processing_status = 'llm_pending' where id = any(cast(:ids as text[])) and processing_status <> 'llm_pending'", ids=[s["id"] for s in snaps]) | |
| 322 | + return {"queued": queued, "candidates": len(snaps), "llm_available": gateway.available} | |
| 323 | + | |
| 324 | + | |
| 325 | +class ReprocessBody(BaseModel): | |
| 326 | + connector: str | |
| 327 | + url: str | None = None | |
| 328 | + | |
| 329 | + | |
| 330 | +@router.post("/reprocess") | |
| 331 | +async def reprocess(body: ReprocessBody) -> dict[str, Any]: | |
| 332 | + if body.connector not in connector_registry(): | |
| 333 | + raise ApiError(404, f"unknown connector {body.connector!r}") | |
| 334 | + async with transaction() as conn: | |
| 335 | + jid = await enqueue(conn, "reprocess_snapshot", {"connector": body.connector, **({"url": body.url} if body.url else {})}, priority=3, | |
| 336 | + dedupe_key=f"reprocess:{body.connector}:{body.url or '*'}") | |
| 337 | + return {"queued": jid is not None, "job_id": jid} | |
| 338 | + | |
| 339 | + | |
| 340 | +# ------------------------------------------------------------------------------------------------------------------ review queue & curation | |
| 341 | + | |
| 342 | + | |
| 343 | +@router.get("/review") | |
| 344 | +async def review(status: str = "pending", kind: str | None = None, limit: int = Query(100, ge=1, le=1000), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 345 | + where = ["r.status = :s"] if status != "all" else ["true"] | |
| 346 | + params: dict[str, Any] = {"s": status} | |
| 347 | + if kind: | |
| 348 | + where.append("r.kind = :k") | |
| 349 | + params["k"] = kind | |
| 350 | + where_sql = " and ".join(where) | |
| 351 | + async with connection() as conn: | |
| 352 | + rows = await fetch_all(conn, f"""select r.*, (select jsonb_agg(jsonb_build_object('id', e.id, 'slug', e.slug, 'name', e.canonical_name, 'entity_type', e.entity_type, 'status', e.status) order by e.slug) | |
| 353 | + from entities e where e.id = any(r.entity_ids)) as entities | |
| 354 | + from review_queue r where {where_sql} order by r.created_at desc limit :lim offset :off""", lim=limit, off=offset, **params) | |
| 355 | + total = await fetch_val(conn, f"select count(*) from review_queue r where {where_sql}", **params) | |
| 356 | + kinds = await fetch_all(conn, "select kind, status, count(*) as n from review_queue group by 1, 2 order by 1, 2") | |
| 357 | + return {"items": rows, "total": int(total or 0), "limit": limit, "offset": offset, "by_kind": [{**k, "n": int(k["n"])} for k in kinds]} | |
| 358 | + | |
| 359 | + | |
| 360 | +class ReviewAction(BaseModel): | |
| 361 | + action: str = Field(..., pattern="^(approve|reject|edit)$") | |
| 362 | + resolution: dict[str, Any] | None = None | |
| 363 | + | |
| 364 | + | |
| 365 | +@router.post("/review/{review_id}") | |
| 366 | +async def review_action(review_id: str, body: ReviewAction) -> dict[str, Any]: | |
| 367 | + async with transaction() as conn: | |
| 368 | + item = await fetch_one(conn, "select * from review_queue where id = :id for update", id=review_id) | |
| 369 | + if not item: | |
| 370 | + raise ApiError(404, "review item not found") | |
| 371 | + if item["status"] != "pending": | |
| 372 | + raise ApiError(409, f"review item already {item['status']}") | |
| 373 | + resolution: dict[str, Any] = dict(body.resolution or {}) | |
| 374 | + effect: dict[str, Any] | None = None | |
| 375 | + if body.action == "approve" and item["kind"] == "merge_candidate": | |
| 376 | + payload = item["payload"] or {} | |
| 377 | + ids = list(item["entity_ids"] or []) | |
| 378 | + source = resolution.get("source_id") or payload.get("source_id") or (ids[0] if len(ids) > 1 else None) | |
| 379 | + target = resolution.get("target_id") or payload.get("target_id") or (ids[1] if len(ids) > 1 else None) | |
| 380 | + if not source or not target: | |
| 381 | + raise ApiError(400, "merge_candidate needs source_id and target_id (payload or resolution)") | |
| 382 | + try: | |
| 383 | + effect = await merge_entities(conn, source, target) | |
| 384 | + except (ValueError, LookupError) as exc: | |
| 385 | + raise ApiError(400, str(exc)) from exc | |
| 386 | + resolution.update({"merged": effect}) | |
| 387 | + elif body.action == "approve" and item["kind"] == "conflict": | |
| 388 | + keep = resolution.get("keep_claim_id") | |
| 389 | + if keep: | |
| 390 | + effect = await _accept_claim(conn, keep) | |
| 391 | + resolution["accepted"] = effect | |
| 392 | + status = {"approve": "approved", "reject": "rejected", "edit": "edited"}[body.action] | |
| 393 | + await execute(conn, "update review_queue set status = :st, resolution = cast(:r as jsonb), resolved_at = now() where id = :id", st=status, r=jsonb(resolution), id=review_id) | |
| 394 | + await cache.cache_invalidate() | |
| 395 | + return {"ok": True, "id": review_id, "status": status, "effect": effect} | |
| 396 | + | |
| 397 | + | |
| 398 | +async def _accept_claim(conn: Any, claim_id: str) -> dict[str, Any]: | |
| 399 | + """Conflict resolution: promote one claim to current, supersede the others for the same property.""" | |
| 400 | + c = await fetch_one(conn, "select * from claims where id = :id", id=claim_id) | |
| 401 | + if not c: | |
| 402 | + raise ApiError(404, "claim not found") | |
| 403 | + await execute(conn, "update claims set status = 'superseded', valid_to = now() where entity_id = :e and property = :p and id <> :id and status in ('current', 'conflicting')", | |
| 404 | + e=c["entity_id"], p=c["property"], id=claim_id) | |
| 405 | + await execute(conn, "update claims set status = 'current', confidence = case when confidence = 'conflicted' then 'high' else confidence end, valid_to = null where id = :id", id=claim_id) | |
| 406 | + prov = {"source_id": c["source_id"], "snapshot_id": c["snapshot_id"], "url": c["source_url"], "tier": c["tier"], "confidence": "high", "extractor": c["extractor"], | |
| 407 | + "observed_at": c["observed_at"].isoformat() if c["observed_at"] else None, "resolved_by": "review"} | |
| 408 | + await execute(conn, "update entities set attributes = attributes || jsonb_build_object(:p, cast(:v as jsonb)), provenance = provenance || jsonb_build_object(:p, cast(:pv as jsonb)), updated_at = now() where id = :e", | |
| 409 | + p=c["property"], v=jsonb(c["value"]), pv=jsonb(prov), e=c["entity_id"]) | |
| 410 | + return {"claim_id": claim_id, "entity_id": c["entity_id"], "property": c["property"]} | |
| 411 | + | |
| 412 | + | |
| 413 | +@router.get("/entities/duplicates") | |
| 414 | +async def duplicates(type: str | None = Query(None, alias="type"), limit: int = Query(100, ge=1, le=500), threshold: float = Query(0.8, ge=0.3, le=1.0)) -> dict[str, Any]: | |
| 415 | + where = "a.entity_type = :t" if type else "true" | |
| 416 | + async with connection() as conn: | |
| 417 | + await execute(conn, "select set_limit(:th)", th=threshold) | |
| 418 | + rows = await fetch_all(conn, f""" | |
| 419 | + select a.id as a_id, a.slug as a_slug, a.canonical_name as a_name, a.entity_type, a.first_seen_at as a_first_seen_at, oa.canonical_name as a_org, | |
| 420 | + b.id as b_id, b.slug as b_slug, b.canonical_name as b_name, b.first_seen_at as b_first_seen_at, ob.canonical_name as b_org, | |
| 421 | + similarity(a.canonical_name, b.canonical_name) as similarity, | |
| 422 | + (select count(*) from claims c where c.entity_id = a.id and c.status = 'current') as a_claims, | |
| 423 | + (select count(*) from claims c where c.entity_id = b.id and c.status = 'current') as b_claims | |
| 424 | + from entities a join entities b on b.entity_type = a.entity_type and b.id > a.id and a.canonical_name % b.canonical_name | |
| 425 | + left join entities oa on oa.id = a.organization_id left join entities ob on ob.id = b.organization_id | |
| 426 | + where {where} and a.merged_into is null and b.merged_into is null and similarity(a.canonical_name, b.canonical_name) > :th | |
| 427 | + order by similarity desc, a.canonical_name limit :lim""", t=type, th=threshold, lim=limit) | |
| 428 | + return {"threshold": threshold, "items": [{"entity_type": r["entity_type"], "similarity": round(float(r["similarity"]), 3), | |
| 429 | + "a": {"id": r["a_id"], "slug": r["a_slug"], "name": r["a_name"], "organization": r["a_org"], "first_seen_at": r["a_first_seen_at"], "claims": int(r["a_claims"])}, | |
| 430 | + "b": {"id": r["b_id"], "slug": r["b_slug"], "name": r["b_name"], "organization": r["b_org"], "first_seen_at": r["b_first_seen_at"], "claims": int(r["b_claims"])}} | |
| 431 | + for r in rows]} | |
| 432 | + | |
| 433 | + | |
| 434 | +class MergeBody(BaseModel): | |
| 435 | + source_id: str | |
| 436 | + target_id: str | |
| 437 | + | |
| 438 | + | |
| 439 | +@router.post("/entities/merge") | |
| 440 | +async def merge(body: MergeBody) -> dict[str, Any]: | |
| 441 | + async with transaction() as conn: | |
| 442 | + src = await fetch_one(conn, "select id from entities where id = :k or slug = :k limit 1", k=body.source_id) | |
| 443 | + dst = await fetch_one(conn, "select id from entities where id = :k or slug = :k limit 1", k=body.target_id) | |
| 444 | + if not src or not dst: | |
| 445 | + raise ApiError(404, "source or target entity not found") | |
| 446 | + try: | |
| 447 | + result = await merge_entities(conn, src["id"], dst["id"]) | |
| 448 | + except (ValueError, LookupError) as exc: | |
| 449 | + raise ApiError(400, str(exc)) from exc | |
| 450 | + await cache.cache_invalidate() | |
| 451 | + return result | |
| 452 | + | |
| 453 | + | |
| 454 | +@router.post("/entities/{entity_id}/claims/{claim_id}/retract") | |
| 455 | +async def retract_claim(entity_id: str, claim_id: str) -> dict[str, Any]: | |
| 456 | + async with transaction() as conn: | |
| 457 | + ent = await fetch_one(conn, "select id from entities where id = :k or slug = :k limit 1", k=entity_id) | |
| 458 | + if not ent: | |
| 459 | + raise ApiError(404, "entity not found") | |
| 460 | + c = await fetch_one(conn, "select * from claims where id = :c and entity_id = :e", c=claim_id, e=ent["id"]) | |
| 461 | + if not c: | |
| 462 | + raise ApiError(404, "claim not found for this entity") | |
| 463 | + was_current = c["status"] == "current" | |
| 464 | + await execute(conn, "update claims set status = 'retracted', valid_to = coalesce(valid_to, now()) where id = :c", c=claim_id) | |
| 465 | + restored = None | |
| 466 | + if was_current: | |
| 467 | + prev = await fetch_one(conn, """select * from claims where entity_id = :e and property = :p and id <> :c and status in ('superseded', 'conflicting') | |
| 468 | + order by tier, valid_from desc limit 1""", e=ent["id"], p=c["property"], c=claim_id) | |
| 469 | + if prev: | |
| 470 | + await execute(conn, "update claims set status = 'current', valid_to = null where id = :id", id=prev["id"]) | |
| 471 | + prov = {"source_id": prev["source_id"], "snapshot_id": prev["snapshot_id"], "url": prev["source_url"], "tier": prev["tier"], "confidence": prev["confidence"], | |
| 472 | + "extractor": prev["extractor"], "observed_at": prev["observed_at"].isoformat() if prev["observed_at"] else None, "restored_by": "retraction"} | |
| 473 | + await execute(conn, "update entities set attributes = attributes || jsonb_build_object(:p, cast(:v as jsonb)), provenance = provenance || jsonb_build_object(:p, cast(:pv as jsonb)), updated_at = now() where id = :e", | |
| 474 | + p=c["property"], v=jsonb(prev["value"]), pv=jsonb(prov), e=ent["id"]) | |
| 475 | + restored = prev["id"] | |
| 476 | + else: | |
| 477 | + await execute(conn, "update entities set attributes = attributes - :p, provenance = provenance - :p, updated_at = now() where id = :e", p=c["property"], e=ent["id"]) | |
| 478 | + await execute(conn, """insert into change_events (id, entity_id, event_type, category, property, old_value, summary, importance, connector_name, dedupe_key) | |
| 479 | + values (:id, :e, 'CLAIM_RETRACTED', 'source', :p, cast(:v as jsonb), :s, 1, 'curation', :dk) on conflict (dedupe_key) do nothing""", | |
| 480 | + id=new_id("change_event"), e=ent["id"], p=c["property"], v=jsonb(c["value"]), | |
| 481 | + s=f"Retracted claim {c['property']}", dk=f"retract:{claim_id}") | |
| 482 | + await cache.cache_invalidate() | |
| 483 | + return {"ok": True, "claim_id": claim_id, "property": c["property"], "was_current": was_current, "restored_claim_id": restored} | |
| 484 | + | |
| 485 | + | |
| 486 | +# ------------------------------------------------------------------------------------------------------------------ infrastructure & maintenance | |
| 487 | + | |
| 488 | + | |
| 489 | +@router.get("/infrastructure") | |
| 490 | +async def infrastructure() -> dict[str, Any]: | |
| 491 | + async with connection() as conn: | |
| 492 | + db = await fetch_one(conn, """select pg_database_size(current_database()) as db_bytes, current_database() as database, version() as pg_version, | |
| 493 | + (select count(*) from pg_stat_activity where datname = current_database()) as connections, | |
| 494 | + (select setting from pg_settings where name = 'max_connections') as max_connections""") | |
| 495 | + tables = await fetch_all(conn, """select c.relname as table, pg_total_relation_size(c.oid) as total_bytes, pg_relation_size(c.oid) as data_bytes, | |
| 496 | + coalesce(s.n_live_tup, 0) as rows_estimate from pg_class c join pg_namespace n on n.oid = c.relnamespace | |
| 497 | + left join pg_stat_user_tables s on s.relid = c.oid where n.nspname = 'public' and c.relkind = 'r' | |
| 498 | + order by pg_total_relation_size(c.oid) desc limit 20""") | |
| 499 | + extensions = await fetch_all(conn, "select extname, extversion from pg_extension order by 1") | |
| 500 | + redis_info: dict[str, Any] = {"ok": False} | |
| 501 | + try: | |
| 502 | + info = await cache.redis().info("memory") | |
| 503 | + redis_info = {"ok": True, "used_memory": info.get("used_memory"), "used_memory_human": info.get("used_memory_human"), | |
| 504 | + "api_cache_keys": sum([1 async for _ in cache.redis().scan_iter(match="aia:api:*", count=1000)])} | |
| 505 | + except Exception as exc: # noqa: BLE001 | |
| 506 | + redis_info = {"ok": False, "error": exc.__class__.__name__} | |
| 507 | + return {"hostname": socket.gethostname(), "python": platform.python_version(), "platform": platform.platform(), "api_uptime_s": int(time.time() - _STARTED), | |
| 508 | + "env": settings.app_env, "heartbeats": await cache.heartbeats(), "archive": archive.archive_size(), "data_dir_exists": settings.data_dir.exists(), | |
| 509 | + "database": {**(db or {}), "db_bytes": int((db or {}).get("db_bytes") or 0), "tables": tables, "extensions": extensions}, "redis": redis_info, | |
| 510 | + "llm": {"available": gateway.available, "engine": gateway.engine.name}, "scheduler_tick_s": settings.scheduler_tick_s} | |
| 511 | + | |
| 512 | + | |
| 513 | +@router.post("/cache/flush") | |
| 514 | +async def cache_flush(prefix: str = "") -> dict[str, Any]: | |
| 515 | + return {"flushed": await cache.cache_invalidate(prefix)} | |
| 516 | + | |
| 517 | + | |
| 518 | +@router.post("/stats/recompute") | |
| 519 | +async def stats_recompute() -> dict[str, Any]: | |
| 520 | + from aiatlas.services.stats import compute_stats | |
| 521 | + | |
| 522 | + counts = await compute_stats() | |
| 523 | + await cache.cache_invalidate("/api/v1/stats") | |
| 524 | + return {"ok": True, "entities_total": counts.get("entities_total")} | |
| 525 | + | |
| 526 | + | |
| 527 | +class QualityBody(BaseModel): | |
| 528 | + entity_ids: list[str] | None = None | |
| 529 | + limit: int = Field(20000, ge=1, le=200000) | |
| 530 | + | |
| 531 | + | |
| 532 | +@router.post("/quality/recompute") | |
| 533 | +async def quality_recompute(body: QualityBody | None = None) -> dict[str, Any]: | |
| 534 | + from aiatlas.services.quality import recompute | |
| 535 | + | |
| 536 | + body = body or QualityBody() | |
| 537 | + res = await recompute(entity_ids=body.entity_ids, limit=body.limit) | |
| 538 | + await cache.cache_invalidate() | |
| 539 | + return {"ok": True, **res} | |
added
src/aiatlas/api/routers/benchmarks.py
+83 −0
@@ -0,0 +1,83 @@ | ||
| 1 | +"""/benchmarks listing, /benchmarks/{slug}, /benchmarks/{slug}/results (leaderboard), /benchmarks/{slug}/history.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + ENTITY_COLS, | |
| 10 | + ENTITY_FROM, | |
| 11 | + PAGINATION, | |
| 12 | + RESULT_COLS, | |
| 13 | + RESULT_FROM, | |
| 14 | + Pagination, | |
| 15 | + cached, | |
| 16 | + entity_cols, | |
| 17 | + entity_summary, | |
| 18 | + page, | |
| 19 | + resolve_entity, | |
| 20 | + resolve_id, | |
| 21 | + result_row, | |
| 22 | +) | |
| 23 | +from aiatlas.api.detail import leaderboard | |
| 24 | +from aiatlas.api.routers.entities import detail_for_type | |
| 25 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 26 | + | |
| 27 | +router = APIRouter(prefix="/api/v1/benchmarks", tags=["benchmarks"]) | |
| 28 | + | |
| 29 | + | |
| 30 | +@router.get("") | |
| 31 | +@cached(300) | |
| 32 | +async def list_benchmarks(request: Request, category: str | None = None) -> dict[str, Any]: | |
| 33 | + where = "e.entity_type = 'benchmark' and e.merged_into is null" + (" and e.attributes->>'category' ilike :cat" if category else "") | |
| 34 | + async with connection() as conn: | |
| 35 | + rows = await fetch_all(conn, f""" | |
| 36 | + select {ENTITY_COLS}, | |
| 37 | + (select count(*) from benchmark_results r where r.benchmark_id = e.id and r.valid_to is null) as result_count, | |
| 38 | + (select count(distinct r.model_id) from benchmark_results r where r.benchmark_id = e.id and r.valid_to is null) as model_count, | |
| 39 | + top.score as top_score, {entity_cols("tm", "t_")} | |
| 40 | + from {ENTITY_FROM} | |
| 41 | + left join lateral (select r.model_id, r.score from benchmark_results r where r.benchmark_id = e.id and r.valid_to is null | |
| 42 | + order by case when r.higher_is_better then -r.score else r.score end limit 1) top on true | |
| 43 | + left join entities tm on tm.id = top.model_id left join entities tmo on tmo.id = tm.organization_id | |
| 44 | + where {where} order by result_count desc, e.canonical_name""", cat=category) | |
| 45 | + items = [] | |
| 46 | + for r in rows: | |
| 47 | + top_model = entity_summary(r, "t_") | |
| 48 | + items.append({**(entity_summary(r) or {}), "result_count": int(r["result_count"] or 0), "model_count": int(r["model_count"] or 0), | |
| 49 | + "top": {"model": top_model, "score": r["top_score"]} if top_model else None}) | |
| 50 | + return {"items": items} | |
| 51 | + | |
| 52 | + | |
| 53 | +@router.get("/{slug}") | |
| 54 | +@cached(300) | |
| 55 | +async def get_benchmark(request: Request, slug: str) -> dict[str, Any]: | |
| 56 | + return await detail_for_type(slug, ("benchmark",)) | |
| 57 | + | |
| 58 | + | |
| 59 | +@router.get("/{slug}/results") | |
| 60 | +@cached(300) | |
| 61 | +async def benchmark_results(request: Request, slug: str, p: Pagination = PAGINATION, config: str | None = Query(None, max_length=200), history: int = Query(0, ge=0, le=1)) -> dict[str, Any]: | |
| 62 | + async with connection() as conn: | |
| 63 | + bench = await resolve_entity(conn, slug, ("benchmark",)) | |
| 64 | + items = await leaderboard(conn, bench["id"], limit=p.limit, offset=p.offset, config=config, history=bool(history)) | |
| 65 | + where = "r.benchmark_id = :id" + ("" if history else " and r.valid_to is null") + (" and r.config::text ilike :cfg" if config else "") | |
| 66 | + total = await fetch_val(conn, f"select count(*) from benchmark_results r where {where}", id=bench["id"], cfg=f"%{config}%" if config else None) | |
| 67 | + out = page(items, int(total or 0), p) | |
| 68 | + out["benchmark"] = entity_summary(bench) | |
| 69 | + return out | |
| 70 | + | |
| 71 | + | |
| 72 | +@router.get("/{slug}/history") | |
| 73 | +@cached(300) | |
| 74 | +async def benchmark_history(request: Request, slug: str, model: str | None = None, limit: int = Query(1000, ge=1, le=5000)) -> dict[str, Any]: | |
| 75 | + async with connection() as conn: | |
| 76 | + bench = await resolve_entity(conn, slug, ("benchmark",)) | |
| 77 | + where = ["r.benchmark_id = :id"] | |
| 78 | + params: dict[str, Any] = {"id": bench["id"], "lim": limit} | |
| 79 | + if model: | |
| 80 | + where.append("r.model_id = :model") | |
| 81 | + params["model"] = await resolve_id(conn, model) | |
| 82 | + rows = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where {' and '.join(where)} order by r.observed_at asc, r.id limit :lim", **params) | |
| 83 | + return {"benchmark": entity_summary(bench), "items": [result_row(r) for r in rows]} | |
added
src/aiatlas/api/routers/changes.py
+125 −0
@@ -0,0 +1,125 @@ | ||
| 1 | +"""/changes (cursor feed) · /changes/daily ("What changed in AI today") · /changes/categories.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from fastapi import APIRouter, Query, Request | |
| 8 | + | |
| 9 | +from aiatlas.api.common import ( | |
| 10 | + ENTITY_COLS, | |
| 11 | + ENTITY_FROM, | |
| 12 | + EVENT_COLS, | |
| 13 | + EVENT_FROM, | |
| 14 | + ApiError, | |
| 15 | + cached, | |
| 16 | + change_event, | |
| 17 | + csv, | |
| 18 | + day_bounds, | |
| 19 | + entity_summary, | |
| 20 | + parse_date, | |
| 21 | + parse_ts, | |
| 22 | + resolve_id, | |
| 23 | +) | |
| 24 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 25 | + | |
| 26 | +router = APIRouter(prefix="/api/v1/changes", tags=["changes"]) | |
| 27 | + | |
| 28 | +CATEGORY_LABELS = {"model": "Models", "price": "Pricing", "benchmark": "Benchmarks", "paper": "Research", "release": "Releases", "company": "Companies & labs", | |
| 29 | + "provider": "Providers", "hardware": "Hardware", "framework": "Frameworks", "dataset": "Datasets", "regulation": "Regulation", | |
| 30 | + "incident": "Incidents", "repository": "Repositories", "tool": "Tools", "source": "Sources & curation", "update": "Updates"} | |
| 31 | +CATEGORY_ORDER = list(CATEGORY_LABELS) | |
| 32 | +TOTAL_CAP = 10_000 | |
| 33 | + | |
| 34 | + | |
| 35 | +def _filters(*, category: str | None, types: list[str], entity_type: str | None, importance_min: int | None, since: datetime | None, until: datetime | None, | |
| 36 | + q: str | None, include_documents: bool, entity_id: str | None = None, before: datetime | None = None) -> tuple[list[str], dict[str, Any]]: | |
| 37 | + where: list[str] = [] | |
| 38 | + p: dict[str, Any] = {} | |
| 39 | + if category: | |
| 40 | + where.append("ev.category = any(cast(:cats as text[]))") | |
| 41 | + p["cats"] = csv(category) | |
| 42 | + if types: | |
| 43 | + where.append("ev.event_type = any(cast(:types as text[]))") | |
| 44 | + p["types"] = [t.upper() for t in types] | |
| 45 | + elif not include_documents: | |
| 46 | + where.append("ev.event_type <> 'DOCUMENT_CHANGED'") | |
| 47 | + if entity_type: | |
| 48 | + where.append("e.entity_type = any(cast(:etypes as text[]))") | |
| 49 | + p["etypes"] = csv(entity_type) | |
| 50 | + if importance_min is not None: | |
| 51 | + where.append("ev.importance >= :imp") | |
| 52 | + p["imp"] = importance_min | |
| 53 | + if since is not None: | |
| 54 | + where.append("ev.observed_at >= :since") | |
| 55 | + p["since"] = since | |
| 56 | + if until is not None: | |
| 57 | + where.append("ev.observed_at <= :until") | |
| 58 | + p["until"] = until | |
| 59 | + if before is not None: | |
| 60 | + where.append("ev.observed_at < :before") | |
| 61 | + p["before"] = before | |
| 62 | + if q: | |
| 63 | + where.append("(ev.summary ilike :qlike or e.canonical_name ilike :qlike)") | |
| 64 | + p["qlike"] = f"%{q}%" | |
| 65 | + if entity_id: | |
| 66 | + where.append("ev.entity_id = :eid") | |
| 67 | + p["eid"] = entity_id | |
| 68 | + return where or ["true"], p | |
| 69 | + | |
| 70 | + | |
| 71 | +@router.get("") | |
| 72 | +@cached(60) | |
| 73 | +async def list_changes(request: Request, category: str | None = None, type: str | None = Query(None, alias="type"), entity_type: str | None = None, | |
| 74 | + importance_min: int | None = Query(None, ge=0, le=3), since: str | None = None, until: str | None = None, q: str | None = Query(None, max_length=200), | |
| 75 | + entity: str | None = None, limit: int = Query(50, ge=1, le=200), before: str | None = None, offset: int = Query(0, ge=0, le=10000), | |
| 76 | + include_documents: int = Query(0, ge=0, le=1)) -> dict[str, Any]: | |
| 77 | + before_ts, since_ts, until_ts = parse_ts(before, "before"), parse_ts(since, "since"), parse_ts(until, "until") | |
| 78 | + async with connection() as conn: | |
| 79 | + eid = await resolve_id(conn, entity) if entity else None | |
| 80 | + where, params = _filters(category=category, types=csv(type), entity_type=entity_type, importance_min=importance_min, since=since_ts, until=until_ts, q=q, | |
| 81 | + include_documents=bool(include_documents), entity_id=eid, before=before_ts) | |
| 82 | + where_sql = " and ".join(where) | |
| 83 | + rows = await fetch_all(conn, f"select {EVENT_COLS} from {EVENT_FROM} where {where_sql} order by ev.observed_at desc, ev.id desc limit :lim offset :off", | |
| 84 | + lim=limit, off=offset, **params) | |
| 85 | + total = await fetch_val(conn, f"select count(*) from (select 1 from {EVENT_FROM} where {where_sql} limit {TOTAL_CAP}) t", **params) | |
| 86 | + items = [change_event(r) for r in rows] | |
| 87 | + return {"items": items, "total": int(total or 0), "limit": limit, "offset": offset, "next_before": items[-1]["observed_at"] if len(items) == limit else None} | |
| 88 | + | |
| 89 | + | |
| 90 | +@router.get("/daily") | |
| 91 | +@cached(120) | |
| 92 | +async def changes_daily(request: Request, date: str | None = None, per_section: int = Query(30, ge=1, le=100)) -> dict[str, Any]: | |
| 93 | + d = parse_date(date, "date") or datetime.now(UTC).date() | |
| 94 | + start, end = day_bounds(d) | |
| 95 | + async with connection() as conn: | |
| 96 | + rows = await fetch_all(conn, f"""select {EVENT_COLS} from ( | |
| 97 | + select ev.*, row_number() over (partition by ev.category order by ev.importance desc, ev.observed_at desc) as rn | |
| 98 | + from change_events ev where ev.observed_at >= :s and ev.observed_at <= :e and ev.event_type <> 'DOCUMENT_CHANGED') ev | |
| 99 | + left join entities e on e.id = ev.entity_id left join entities eo on eo.id = e.organization_id | |
| 100 | + where ev.rn <= :n order by ev.category, ev.rn""", s=start, e=end, n=per_section) | |
| 101 | + counts = await fetch_all(conn, "select category, count(*) as n from change_events where observed_at >= :s and observed_at <= :e and event_type <> 'DOCUMENT_CHANGED' group by 1", s=start, e=end) | |
| 102 | + new_models = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.entity_type = 'model' and e.merged_into is null and e.first_seen_at >= :s and e.first_seen_at <= :e " | |
| 103 | + f"order by e.first_seen_at desc limit 100", s=start, e=end) | |
| 104 | + prev = await fetch_val(conn, "select max(observed_at)::date from change_events where observed_at < :s and event_type <> 'DOCUMENT_CHANGED'", s=start) | |
| 105 | + nxt = await fetch_val(conn, "select min(observed_at)::date from change_events where observed_at > :e and event_type <> 'DOCUMENT_CHANGED'", e=end) | |
| 106 | + by_cat: dict[str, list[dict[str, Any]]] = {} | |
| 107 | + for r in rows: | |
| 108 | + by_cat.setdefault(r["category"], []).append(change_event(r)) | |
| 109 | + order = {c: i for i, c in enumerate(CATEGORY_ORDER)} | |
| 110 | + sections = [{"category": c, "label": CATEGORY_LABELS.get(c, c.title()), "items": items} for c, items in sorted(by_cat.items(), key=lambda kv: (order.get(kv[0], 99), kv[0]))] | |
| 111 | + return {"date": d.isoformat(), "counts": {r["category"]: int(r["n"]) for r in counts}, "total": sum(int(r["n"]) for r in counts), "sections": sections, | |
| 112 | + "new_models": [entity_summary(r) for r in new_models], "labels": CATEGORY_LABELS, | |
| 113 | + "previous_day": prev.isoformat() if prev else None, "next_day": nxt.isoformat() if nxt else None} | |
| 114 | + | |
| 115 | + | |
| 116 | +@router.get("/categories") | |
| 117 | +@cached(300) | |
| 118 | +async def changes_categories(request: Request, days: int = Query(7, ge=1, le=365)) -> dict[str, Any]: | |
| 119 | + async with connection() as conn: | |
| 120 | + rows = await fetch_all(conn, """select category, event_type, count(*) as count from change_events where observed_at > now() - make_interval(days => :d) | |
| 121 | + and event_type <> 'DOCUMENT_CHANGED' group by 1, 2 order by 3 desc, 1, 2""", d=days) | |
| 122 | + return {"days": days, "items": [{"category": r["category"], "label": CATEGORY_LABELS.get(r["category"], r["category"].title()), "event_type": r["event_type"], "count": int(r["count"])} for r in rows]} | |
| 123 | + | |
| 124 | + | |
| 125 | +__all__ = ["CATEGORY_LABELS", "ApiError", "router"] | |
added
src/aiatlas/api/routers/companies.py
+64 −0
@@ -0,0 +1,64 @@ | ||
| 1 | +"""/companies listing (company, organization, lab, university) and /companies/{slug} alias.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + COMPANY_TYPES, | |
| 10 | + ENTITY_COLS, | |
| 11 | + ENTITY_FROM, | |
| 12 | + PAGINATION, | |
| 13 | + ApiError, | |
| 14 | + Pagination, | |
| 15 | + cached, | |
| 16 | + entity_summary, | |
| 17 | + page, | |
| 18 | +) | |
| 19 | +from aiatlas.api.routers.entities import detail_for_type | |
| 20 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 21 | + | |
| 22 | +router = APIRouter(prefix="/api/v1/companies", tags=["companies"]) | |
| 23 | + | |
| 24 | +MODEL_COUNT = "(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null)" | |
| 25 | +PAPER_COUNT = ("(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) + " | |
| 26 | + "(select count(*) from relations r join entities pp on pp.id = r.subject_id where r.object_id = e.id and r.valid_to is null and pp.entity_type = 'paper' and pp.organization_id is distinct from e.id)") | |
| 27 | +SORTS = {"models": "model_count desc, e.canonical_name", "name": "e.canonical_name asc", "updated": "e.updated_at desc", | |
| 28 | + "quality": "coalesce((e.quality->>'score')::float, 0) desc, e.canonical_name", "papers": "paper_count desc, e.canonical_name"} | |
| 29 | + | |
| 30 | + | |
| 31 | +@router.get("") | |
| 32 | +@cached(300) | |
| 33 | +async def list_companies(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), country: str | None = None, kind: str | None = None, | |
| 34 | + sort: str = "models", facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]: | |
| 35 | + if sort not in SORTS: | |
| 36 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 37 | + where = ["e.entity_type = any(cast(:types as text[]))", "e.merged_into is null"] | |
| 38 | + params: dict[str, Any] = {"types": list(COMPANY_TYPES)} | |
| 39 | + if q: | |
| 40 | + where.append("(e.canonical_name ilike :qlike or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))") | |
| 41 | + params["qlike"] = f"%{q}%" | |
| 42 | + if country: | |
| 43 | + where.append("e.attributes->>'country' ilike :country") | |
| 44 | + params["country"] = country | |
| 45 | + if kind: | |
| 46 | + where.append("(e.attributes->>'org_kind' ilike :kind or e.entity_type = :kind)") | |
| 47 | + params["kind"] = kind | |
| 48 | + where_sql = " and ".join(where) | |
| 49 | + async with connection() as conn: | |
| 50 | + rows = await fetch_all(conn, f"select {ENTITY_COLS}, {MODEL_COUNT} as model_count, {PAPER_COUNT} as paper_count from {ENTITY_FROM} where {where_sql} " | |
| 51 | + f"order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) | |
| 52 | + total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) | |
| 53 | + out = page([{**(entity_summary(r) or {}), "model_count": int(r["model_count"]), "paper_count": int(r["paper_count"])} for r in rows], int(total or 0), p) | |
| 54 | + if facets: | |
| 55 | + countries = await fetch_all(conn, f"select e.attributes->>'country' as value, count(*) as count from {ENTITY_FROM} where {where_sql} and e.attributes ? 'country' group by 1 order by 2 desc limit 60", **params) | |
| 56 | + kinds = await fetch_all(conn, f"select coalesce(e.attributes->>'org_kind', e.entity_type) as value, count(*) as count from {ENTITY_FROM} where {where_sql} group by 1 order by 2 desc", **params) | |
| 57 | + out["facets"] = {"countries": [{"value": r["value"], "count": int(r["count"])} for r in countries], "kinds": [{"value": r["value"], "count": int(r["count"])} for r in kinds]} | |
| 58 | + return out | |
| 59 | + | |
| 60 | + | |
| 61 | +@router.get("/{slug}") | |
| 62 | +@cached(300) | |
| 63 | +async def get_company(request: Request, slug: str) -> dict[str, Any]: | |
| 64 | + return await detail_for_type(slug, COMPANY_TYPES) | |
added
src/aiatlas/api/routers/compare.py
+116 −0
@@ -0,0 +1,116 @@ | ||
| 1 | +"""/compare?ids=a,b,… — side-by-side of 2–6 entities of the same type with per-type dimension sets and provenance.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + COMPANY_TYPES, | |
| 10 | + PRICE_COLS, | |
| 11 | + PRICE_FROM, | |
| 12 | + RESULT_COLS, | |
| 13 | + RESULT_FROM, | |
| 14 | + ApiError, | |
| 15 | + cached, | |
| 16 | + csv, | |
| 17 | + entity_summary, | |
| 18 | + price_row, | |
| 19 | + resolve_entity, | |
| 20 | + result_row, | |
| 21 | +) | |
| 22 | +from aiatlas.db import connection, fetch_all | |
| 23 | + | |
| 24 | +router = APIRouter(prefix="/api/v1/compare", tags=["compare"]) | |
| 25 | + | |
| 26 | +D = lambda key, label, kind="text", unit=None, source="attr": {"key": key, "label": label, "kind": kind, **({"unit": unit} if unit else {}), "source": source} | |
| 27 | + | |
| 28 | +DIMENSIONS: dict[str, list[dict[str, Any]]] = { | |
| 29 | + "model": [D("parameter_count", "Parameters", "number", "params"), D("active_parameter_count", "Active parameters", "number", "params"), | |
| 30 | + D("context_length", "Context window", "number", "tokens"), D("max_output_tokens", "Max output", "number", "tokens"), D("openness", "Openness"), | |
| 31 | + D("license", "License"), D("modalities", "Modalities", "list"), D("release_date", "Release date", "date"), D("knowledge_cutoff", "Knowledge cutoff", "date"), | |
| 32 | + D("status", "Status"), D("family", "Family"), D("architecture", "Architecture"), | |
| 33 | + D("best_input_per_mtok", "Best input price", "number", "USD / 1M tokens", "prices"), D("best_output_per_mtok", "Best output price", "number", "USD / 1M tokens", "prices"), | |
| 34 | + D("provider_count", "Providers", "number", None, "prices")], | |
| 35 | + "provider": [D("website", "Website"), D("pricing_url", "Pricing page"), D("model_count", "Models priced", "number", None, "prices"), | |
| 36 | + D("min_input_per_mtok", "Min input price", "number", "USD / 1M tokens", "prices"), D("min_output_per_mtok", "Min output price", "number", "USD / 1M tokens", "prices"), | |
| 37 | + D("features", "Features", "list", None, "prices")], | |
| 38 | + "hardware": [D("kind", "Kind"), D("manufacturer", "Manufacturer"), D("memory_gb", "Memory", "number", "GB"), D("memory_bandwidth_gbs", "Memory bandwidth", "number", "GB/s"), | |
| 39 | + D("tdp_watts", "TDP", "number", "W"), D("architecture", "Architecture"), D("memory_type", "Memory type"), D("runtimes", "Runtimes", "list"), D("release_date", "Release", "date")], | |
| 40 | + "framework": [D("latest_version", "Latest version"), D("latest_release_at", "Latest release", "date"), D("license", "License"), D("language", "Language"), | |
| 41 | + D("metric.stars", "GitHub stars", "number"), D("metric.forks", "Forks", "number"), D("repository_url", "Repository")], | |
| 42 | + "company": [D("country", "Country"), D("founded", "Founded", "date"), D("headquarters", "Headquarters"), D("org_kind", "Kind"), D("website", "Website"), | |
| 43 | + D("model_count", "Models", "number", None, "graph"), D("paper_count", "Papers", "number", None, "graph")], | |
| 44 | + "benchmark": [D("category", "Category"), D("metric", "Metric"), D("unit", "Unit"), D("task", "Task"), D("result_count", "Results", "number", None, "results")], | |
| 45 | +} | |
| 46 | + | |
| 47 | + | |
| 48 | +@router.get("") | |
| 49 | +@cached(300) | |
| 50 | +async def compare(request: Request, ids: str = Query(..., description="2–6 slugs or ids, comma-separated")) -> dict[str, Any]: | |
| 51 | + keys = csv(ids) | |
| 52 | + if not 2 <= len(keys) <= 6: | |
| 53 | + raise ApiError(400, "ids must list between 2 and 6 entities") | |
| 54 | + async with connection() as conn: | |
| 55 | + rows = [await resolve_entity(conn, k) for k in keys] | |
| 56 | + types = {r["entity_type"] for r in rows} | |
| 57 | + etype = rows[0]["entity_type"] | |
| 58 | + if etype in COMPANY_TYPES: | |
| 59 | + etype = "company" | |
| 60 | + types = {"company"} | |
| 61 | + if len(types) != 1: | |
| 62 | + raise ApiError(400, f"all entities must share one type, got {', '.join(sorted(types))}") | |
| 63 | + dims = [dict(d) for d in DIMENSIONS.get(etype, [])] | |
| 64 | + eids = [r["id"] for r in rows] | |
| 65 | + items: list[dict[str, Any]] = [] | |
| 66 | + extra: dict[str, dict[str, Any]] = {eid: {} for eid in eids} | |
| 67 | + prices_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids} | |
| 68 | + results_by: dict[str, list[dict[str, Any]]] = {eid: [] for eid in eids} | |
| 69 | + if etype == "model": | |
| 70 | + for p in await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where p.model_id = any(cast(:ids as text[])) and p.valid_to is null order by p.input_per_mtok nulls last", ids=eids): | |
| 71 | + prices_by[p["m_id"]].append(price_row(p)) | |
| 72 | + for eid, plist in prices_by.items(): | |
| 73 | + ins = [x["input_per_mtok"] for x in plist if x["input_per_mtok"] is not None] | |
| 74 | + outs = [x["output_per_mtok"] for x in plist if x["output_per_mtok"] is not None] | |
| 75 | + extra[eid] = {"best_input_per_mtok": min(ins) if ins else None, "best_output_per_mtok": min(outs) if outs else None, | |
| 76 | + "provider_count": len({x["provider"]["id"] for x in plist if x["provider"]})} | |
| 77 | + res = await fetch_all(conn, f"select {RESULT_COLS} from {RESULT_FROM} where r.model_id = any(cast(:ids as text[])) and r.valid_to is null order by b.canonical_name, r.observed_at desc", ids=eids) | |
| 78 | + per_bench: dict[str, dict[str, float]] = {} | |
| 79 | + bench_meta: dict[str, dict[str, Any]] = {} | |
| 80 | + for r in res: | |
| 81 | + results_by[r["m_id"]].append(result_row(r)) | |
| 82 | + per_bench.setdefault(r["b_id"], {}).setdefault(r["m_id"], r["score"]) | |
| 83 | + bench_meta.setdefault(r["b_id"], {"slug": r["b_slug"], "name": r["b_canonical_name"], "unit": r.get("unit") or r.get("b_attributes", {}).get("unit"), "higher_is_better": r["higher_is_better"]}) | |
| 84 | + for bid, scores in per_bench.items(): | |
| 85 | + if all(eid in scores for eid in eids): | |
| 86 | + meta = bench_meta[bid] | |
| 87 | + dims.append({"key": f"bench:{meta['slug']}", "label": meta["name"], "kind": "number", "unit": meta["unit"], "source": "results", "higher_is_better": meta["higher_is_better"]}) | |
| 88 | + for eid in eids: | |
| 89 | + extra[eid][f"bench:{meta['slug']}"] = scores[eid] | |
| 90 | + elif etype == "provider": | |
| 91 | + agg = await fetch_all(conn, """select provider_id, count(distinct model_id) as model_count, min(nullif(input_per_mtok, 0)) as min_input_per_mtok, | |
| 92 | + min(nullif(output_per_mtok, 0)) as min_output_per_mtok, jsonb_agg(distinct k.key) filter (where k.key is not null) as features | |
| 93 | + from prices p left join lateral jsonb_object_keys(p.features) k(key) on true | |
| 94 | + where provider_id = any(cast(:ids as text[])) and valid_to is null group by 1""", ids=eids) | |
| 95 | + for a in agg: | |
| 96 | + extra[a["provider_id"]] = {"model_count": int(a["model_count"]), "min_input_per_mtok": a["min_input_per_mtok"], "min_output_per_mtok": a["min_output_per_mtok"], "features": a["features"] or []} | |
| 97 | + elif etype == "company": | |
| 98 | + agg = await fetch_all(conn, """select e.id, (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null) as model_count, | |
| 99 | + (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) as paper_count | |
| 100 | + from entities e where e.id = any(cast(:ids as text[]))""", ids=eids) | |
| 101 | + for a in agg: | |
| 102 | + extra[a["id"]] = {"model_count": int(a["model_count"]), "paper_count": int(a["paper_count"])} | |
| 103 | + elif etype == "benchmark": | |
| 104 | + agg = await fetch_all(conn, "select benchmark_id, count(*) as n from benchmark_results where benchmark_id = any(cast(:ids as text[])) and valid_to is null group by 1", ids=eids) | |
| 105 | + for a in agg: | |
| 106 | + extra[a["benchmark_id"]] = {"result_count": int(a["n"])} | |
| 107 | + for r in rows: | |
| 108 | + attrs = r.get("attributes") or {} | |
| 109 | + prov = r.get("provenance") or {} | |
| 110 | + values = {d["key"]: (extra[r["id"]].get(d["key"]) if d.get("source") != "attr" else attrs.get(d["key"])) for d in dims} | |
| 111 | + item: dict[str, Any] = {"entity": entity_summary(r), "values": values, "provenance": {k: prov[k] for k in values if k in prov}} | |
| 112 | + if etype == "model": | |
| 113 | + item["prices"] = prices_by[r["id"]] | |
| 114 | + item["results"] = results_by[r["id"]] | |
| 115 | + items.append(item) | |
| 116 | + return {"entity_type": etype, "dimensions": dims, "items": items} | |
added
src/aiatlas/api/routers/diff.py
+85 −0
@@ -0,0 +1,85 @@ | ||
| 1 | +"""/diff?a=&b=&scope= — what changed between two dates: new / gone entities, property, price and benchmark changes.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime | |
| 5 | +from datetime import time as dtime | |
| 6 | +from typing import Any | |
| 7 | + | |
| 8 | +from fastapi import APIRouter, Query, Request | |
| 9 | + | |
| 10 | +from aiatlas.api.common import ( | |
| 11 | + ENTITY_COLS, | |
| 12 | + ENTITY_FROM, | |
| 13 | + EVENT_COLS, | |
| 14 | + EVENT_FROM, | |
| 15 | + ApiError, | |
| 16 | + cached, | |
| 17 | + change_event, | |
| 18 | + entity_summary, | |
| 19 | + parse_date, | |
| 20 | +) | |
| 21 | +from aiatlas.db import connection, fetch_all, fetch_one, fetch_val | |
| 22 | + | |
| 23 | +router = APIRouter(prefix="/api/v1/diff", tags=["diff"]) | |
| 24 | +GONE_TYPES = ("RETIREMENT_ANNOUNCED", "DEPRECATION_ANNOUNCED", "STATUS_CHANGED", "ENTITY_MERGED", "MODEL_RETIRED", "MODEL_DEPRECATED") | |
| 25 | + | |
| 26 | + | |
| 27 | +async def _scope(conn: Any, scope: str) -> tuple[str, str, dict[str, Any], dict[str, Any]]: | |
| 28 | + """Returns (entity where-fragment on alias e, event where-fragment on alias e, params, scope description).""" | |
| 29 | + s = (scope or "all").strip() | |
| 30 | + if s == "all": | |
| 31 | + return "true", "true", {}, {"kind": "all"} | |
| 32 | + if s == "models": | |
| 33 | + return "e.entity_type = 'model'", "e.entity_type = 'model'", {}, {"kind": "models"} | |
| 34 | + if s.startswith("org:"): | |
| 35 | + key = s[4:] | |
| 36 | + org = await fetch_one(conn, "select id, slug, canonical_name from entities where slug = :k or id = :k limit 1", k=key) | |
| 37 | + if not org: | |
| 38 | + raise ApiError(404, f"organization {key!r} not found") | |
| 39 | + frag = "(e.organization_id = :org or e.id = :org)" | |
| 40 | + return frag, frag, {"org": org["id"]}, {"kind": "org", "organization": {"id": org["id"], "slug": org["slug"], "name": org["canonical_name"]}} | |
| 41 | + if s.startswith("family:"): | |
| 42 | + frag = "e.attributes->>'family' ilike :family" | |
| 43 | + return frag, frag, {"family": s[7:]}, {"kind": "family", "family": s[7:]} | |
| 44 | + raise ApiError(400, "scope must be all | models | org:<slug> | family:<name>") | |
| 45 | + | |
| 46 | + | |
| 47 | +@router.get("") | |
| 48 | +@cached(300) | |
| 49 | +async def diff(request: Request, a: str = Query(..., description="YYYY-MM-DD"), b: str = Query(..., description="YYYY-MM-DD"), scope: str = "all", | |
| 50 | + limit: int = Query(200, ge=1, le=1000)) -> dict[str, Any]: | |
| 51 | + da, db = parse_date(a, "a"), parse_date(b, "b") | |
| 52 | + assert da is not None and db is not None | |
| 53 | + if da > db: | |
| 54 | + da, db = db, da | |
| 55 | + start, end = datetime.combine(da, dtime.max, UTC), datetime.combine(db, dtime.max, UTC) | |
| 56 | + async with connection() as conn: | |
| 57 | + ent_frag, ev_frag, params, scope_desc = await _scope(conn, scope) | |
| 58 | + new_entities = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.merged_into is null and {ent_frag} and e.first_seen_at > :s and e.first_seen_at <= :e " | |
| 59 | + f"order by e.first_seen_at desc limit :lim", s=start, e=end, lim=limit, **params) | |
| 60 | + gone = await fetch_all(conn, f"""select distinct on (e.id) {ENTITY_COLS} from change_events ev join entities e on e.id = ev.entity_id left join entities eo on eo.id = e.organization_id | |
| 61 | + where {ev_frag} and ev.observed_at > :s and ev.observed_at <= :e and (ev.event_type = any(cast(:gt as text[])) or e.merged_into is not null) | |
| 62 | + and (e.status in ('retired','deprecated','merged') or ev.event_type in ('RETIREMENT_ANNOUNCED','DEPRECATION_ANNOUNCED')) | |
| 63 | + order by e.id limit :lim""", s=start, e=end, gt=list(GONE_TYPES), lim=limit, **params) | |
| 64 | + | |
| 65 | + async def events(cond: str) -> list[dict[str, Any]]: | |
| 66 | + rows = await fetch_all(conn, f"select {EVENT_COLS} from {EVENT_FROM} where {ev_frag} and ev.observed_at > :s and ev.observed_at <= :e and {cond} " | |
| 67 | + f"order by ev.importance desc, ev.observed_at desc limit :lim", s=start, e=end, lim=limit, **params) | |
| 68 | + return [change_event(r) for r in rows] | |
| 69 | + | |
| 70 | + prop_changes = await events("ev.property is not null and ev.category not in ('price','benchmark') and ev.event_type <> 'DOCUMENT_CHANGED'") | |
| 71 | + price_changes = await events("ev.category = 'price'") | |
| 72 | + bench_changes = await events("ev.category = 'benchmark'") | |
| 73 | + counts = await fetch_one(conn, f"""select (select count(*) from entities e where e.merged_into is null and {ent_frag} and e.first_seen_at > :s and e.first_seen_at <= :e) as new_entities, | |
| 74 | + (select count(*) from change_events ev left join entities e on e.id = ev.entity_id where {ev_frag} and ev.observed_at > :s and ev.observed_at <= :e and ev.event_type <> 'DOCUMENT_CHANGED') as events, | |
| 75 | + (select count(*) from change_events ev left join entities e on e.id = ev.entity_id where {ev_frag} and ev.observed_at > :s and ev.observed_at <= :e and ev.category = 'price') as price_changes, | |
| 76 | + (select count(*) from change_events ev left join entities e on e.id = ev.entity_id where {ev_frag} and ev.observed_at > :s and ev.observed_at <= :e and ev.category = 'benchmark') as benchmark_changes, | |
| 77 | + (select count(*) from prices p join entities e on e.id = p.model_id where {ent_frag} and p.valid_from > :s and p.valid_from <= :e) as price_rows_opened, | |
| 78 | + (select count(*) from prices p join entities e on e.id = p.model_id where {ent_frag} and p.valid_to > :s and p.valid_to <= :e) as price_rows_closed, | |
| 79 | + (select count(*) from claims c join entities e on e.id = c.entity_id where {ent_frag} and c.valid_to > :s and c.valid_to <= :e) as claims_superseded""", | |
| 80 | + s=start, e=end, **params) | |
| 81 | + entities_at_a = await fetch_val(conn, f"select count(*) from entities e where e.merged_into is null and {ent_frag} and e.first_seen_at <= :s", s=start, **params) | |
| 82 | + entities_at_b = await fetch_val(conn, f"select count(*) from entities e where e.merged_into is null and {ent_frag} and e.first_seen_at <= :e", e=end, **params) | |
| 83 | + return {"a": da.isoformat(), "b": db.isoformat(), "scope": scope_desc, "new_entities": [entity_summary(r) for r in new_entities], "gone_entities": [entity_summary(r) for r in gone], | |
| 84 | + "property_changes": prop_changes, "price_changes": price_changes, "benchmark_changes": bench_changes, | |
| 85 | + "counts": {**{k: int(v or 0) for k, v in (counts or {}).items()}, "gone_entities": len(gone), "entities_at_a": int(entities_at_a or 0), "entities_at_b": int(entities_at_b or 0)}} | |
added
src/aiatlas/api/routers/entities.py
+152 −0
@@ -0,0 +1,152 @@ | ||
| 1 | +"""/entities/{slug_or_id} and its sub-resources (timeline, history, asof, graph, sources, related).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime | |
| 5 | +from datetime import time as dtime | |
| 6 | +from typing import Any | |
| 7 | + | |
| 8 | +from fastapi import APIRouter, Query, Request | |
| 9 | + | |
| 10 | +from aiatlas.api.common import ( | |
| 11 | + CLAIM_COLS, | |
| 12 | + CLAIM_FROM, | |
| 13 | + ENTITY_COLS, | |
| 14 | + ENTITY_FROM, | |
| 15 | + ApiError, | |
| 16 | + cached, | |
| 17 | + claim_row, | |
| 18 | + entity_summary, | |
| 19 | + parse_date, | |
| 20 | + parse_ts, | |
| 21 | + resolve_entity, | |
| 22 | +) | |
| 23 | +from aiatlas.api.detail import entity_detail, sources_of, timeline_of | |
| 24 | +from aiatlas.db import connection, fetch_all, fetch_one | |
| 25 | + | |
| 26 | +router = APIRouter(prefix="/api/v1/entities", tags=["entities"]) | |
| 27 | + | |
| 28 | + | |
| 29 | +@router.get("/{slug_or_id}") | |
| 30 | +@cached(300) | |
| 31 | +async def get_entity(request: Request, slug_or_id: str) -> dict[str, Any]: | |
| 32 | + async with connection() as conn: | |
| 33 | + row = await resolve_entity(conn, slug_or_id) | |
| 34 | + return await entity_detail(row) | |
| 35 | + | |
| 36 | + | |
| 37 | +@router.get("/{slug_or_id}/timeline") | |
| 38 | +@cached(60) | |
| 39 | +async def entity_timeline(request: Request, slug_or_id: str, limit: int = Query(50, ge=1, le=200), before: str | None = None, | |
| 40 | + include_documents: int = Query(0, ge=0, le=1)) -> dict[str, Any]: | |
| 41 | + before_ts = parse_ts(before, "before") | |
| 42 | + async with connection() as conn: | |
| 43 | + row = await resolve_entity(conn, slug_or_id) | |
| 44 | + items = await timeline_of(conn, row["id"], row["entity_type"], limit=limit, before=before_ts, include_documents=bool(include_documents)) | |
| 45 | + return {"items": items, "next_before": items[-1]["observed_at"] if len(items) == limit else None} | |
| 46 | + | |
| 47 | + | |
| 48 | +@router.get("/{slug_or_id}/history") | |
| 49 | +@cached(120) | |
| 50 | +async def entity_history(request: Request, slug_or_id: str, property: str | None = Query(None, max_length=120), limit: int = Query(500, ge=1, le=2000)) -> dict[str, Any]: | |
| 51 | + async with connection() as conn: | |
| 52 | + row = await resolve_entity(conn, slug_or_id) | |
| 53 | + where = "c.entity_id = :id" + (" and c.property = :p" if property else "") | |
| 54 | + rows = await fetch_all(conn, f"select {CLAIM_COLS} from {CLAIM_FROM} where {where} order by c.valid_from desc, c.observed_at desc limit :lim", | |
| 55 | + id=row["id"], p=property, lim=limit) | |
| 56 | + return {"items": [claim_row(r) for r in rows]} | |
| 57 | + | |
| 58 | + | |
| 59 | +@router.get("/{slug_or_id}/asof") | |
| 60 | +@cached(300) | |
| 61 | +async def entity_asof(request: Request, slug_or_id: str, date: str = Query(..., description="YYYY-MM-DD")) -> dict[str, Any]: | |
| 62 | + d = parse_date(date, "date") | |
| 63 | + assert d is not None | |
| 64 | + at = datetime.combine(d, dtime.max, UTC) # end of that day, UTC | |
| 65 | + async with connection() as conn: | |
| 66 | + row = await resolve_entity(conn, slug_or_id) | |
| 67 | + existed = row["first_seen_at"] is not None and row["first_seen_at"] <= at | |
| 68 | + rows = await fetch_all(conn, f"""select distinct on (c.property) {CLAIM_COLS} from {CLAIM_FROM} | |
| 69 | + where c.entity_id = :id and c.status <> 'retracted' and c.valid_from <= :at and (c.valid_to is null or c.valid_to > :at) | |
| 70 | + order by c.property, c.tier, c.valid_from desc""", id=row["id"], at=at) if existed else [] | |
| 71 | + claims = [claim_row(r) for r in rows] | |
| 72 | + return {"id": row["id"], "slug": row["slug"], "name": row["canonical_name"], "entity_type": row["entity_type"], "existed": bool(existed), | |
| 73 | + "first_seen_at": row["first_seen_at"], "date": d.isoformat(), "attributes": {c["property"]: c["value"] for c in claims}, "claims": claims} | |
| 74 | + | |
| 75 | + | |
| 76 | +@router.get("/{slug_or_id}/graph") | |
| 77 | +@cached(300) | |
| 78 | +async def entity_graph(request: Request, slug_or_id: str, depth: int = Query(1, ge=1, le=2), limit: int = Query(80, ge=2, le=300)) -> dict[str, Any]: | |
| 79 | + async with connection() as conn: | |
| 80 | + row = await resolve_entity(conn, slug_or_id) | |
| 81 | + root = row["id"] | |
| 82 | + seen = {root} | |
| 83 | + frontier = [root] | |
| 84 | + edges: list[dict[str, str]] = [] | |
| 85 | + for _ in range(depth): | |
| 86 | + if not frontier or len(seen) >= limit: | |
| 87 | + break | |
| 88 | + rows = await fetch_all(conn, """select r.subject_id, r.predicate, r.object_id from relations r | |
| 89 | + where r.valid_to is null and (r.subject_id = any(cast(:ids as text[])) or r.object_id = any(cast(:ids as text[]))) | |
| 90 | + order by r.observed_at desc limit :lim""", ids=frontier, lim=limit * 3) | |
| 91 | + nxt: list[str] = [] | |
| 92 | + for r in rows: | |
| 93 | + other = r["object_id"] if r["subject_id"] in seen else r["subject_id"] | |
| 94 | + if other not in seen: | |
| 95 | + if len(seen) >= limit: | |
| 96 | + continue | |
| 97 | + seen.add(other) | |
| 98 | + nxt.append(other) | |
| 99 | + edges.append({"source": r["subject_id"], "target": r["object_id"], "predicate": r["predicate"]}) | |
| 100 | + frontier = nxt | |
| 101 | + nodes = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.id = any(cast(:ids as text[]))", ids=list(seen)) | |
| 102 | + node_ids = {n["id"] for n in nodes} | |
| 103 | + uniq = {(e["source"], e["target"], e["predicate"]): e for e in edges if e["source"] in node_ids and e["target"] in node_ids} | |
| 104 | + return {"root": root, "nodes": [{"id": n["id"], "slug": n["slug"], "name": n["canonical_name"], "entity_type": n["entity_type"], | |
| 105 | + "organization_name": n["organization_name"]} for n in nodes], "edges": list(uniq.values())} | |
| 106 | + | |
| 107 | + | |
| 108 | +@router.get("/{slug_or_id}/sources") | |
| 109 | +@cached(300) | |
| 110 | +async def entity_sources(request: Request, slug_or_id: str) -> dict[str, Any]: | |
| 111 | + async with connection() as conn: | |
| 112 | + row = await resolve_entity(conn, slug_or_id) | |
| 113 | + items = await sources_of(conn, row["id"], limit=200) | |
| 114 | + return {"items": items} | |
| 115 | + | |
| 116 | + | |
| 117 | +@router.get("/{slug_or_id}/related") | |
| 118 | +@cached(300) | |
| 119 | +async def entity_related(request: Request, slug_or_id: str, limit: int = Query(12, ge=1, le=60)) -> dict[str, Any]: | |
| 120 | + async with connection() as conn: | |
| 121 | + row = await resolve_entity(conn, slug_or_id) | |
| 122 | + fam = (row.get("attributes") or {}).get("family") | |
| 123 | + rows = await fetch_all(conn, f""" | |
| 124 | + with cand as ( | |
| 125 | + select e.id, 3 as w from entities e where e.entity_type = :t and e.organization_id is not null and e.organization_id = :org and e.id <> :id | |
| 126 | + union all | |
| 127 | + select e.id, 4 from entities e where :fam <> '' and e.entity_type = :t and e.attributes->>'family' = :fam and e.id <> :id | |
| 128 | + union all | |
| 129 | + select r2.subject_id, 2 from relations r1 join relations r2 on r2.object_id = r1.object_id and r2.predicate = r1.predicate | |
| 130 | + where r1.subject_id = :id and r1.valid_to is null and r2.valid_to is null and r2.subject_id <> :id | |
| 131 | + union all | |
| 132 | + select r2.object_id, 1 from relations r1 join relations r2 on r2.subject_id = r1.subject_id and r2.predicate = r1.predicate | |
| 133 | + where r1.object_id = :id and r1.valid_to is null and r2.valid_to is null and r2.object_id <> :id), | |
| 134 | + scored as (select id, sum(w) as w from cand group by id) | |
| 135 | + select {ENTITY_COLS} from scored join entities e on e.id = scored.id left join entities eo on eo.id = e.organization_id | |
| 136 | + where e.merged_into is null order by scored.w desc, coalesce((e.quality->>'score')::float, 0) desc, e.updated_at desc limit :lim""", | |
| 137 | + id=row["id"], t=row["entity_type"], org=row.get("organization_id") or "", fam=str(fam or ""), lim=limit) | |
| 138 | + return {"items": [entity_summary(r) for r in rows]} | |
| 139 | + | |
| 140 | + | |
| 141 | +async def detail_for_type(slug: str, types: tuple[str, ...]) -> dict[str, Any]: | |
| 142 | + async with connection() as conn: | |
| 143 | + row = await resolve_entity(conn, slug, types) | |
| 144 | + return await entity_detail(row) | |
| 145 | + | |
| 146 | + | |
| 147 | +async def entity_exists(slug: str) -> dict[str, Any] | None: | |
| 148 | + async with connection() as conn: | |
| 149 | + return await fetch_one(conn, "select id, entity_type from entities where slug = :s or id = :s limit 1", s=slug) | |
| 150 | + | |
| 151 | + | |
| 152 | +__all__ = ["ApiError", "detail_for_type", "router"] | |
added
src/aiatlas/api/routers/explore.py
+84 −0
@@ -0,0 +1,84 @@ | ||
| 1 | +"""/explore/types · /explore/{type} — generic listing for every entity type, plus /frameworks, /datasets, /tools aliases.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + ENTITY_COLS, | |
| 10 | + ENTITY_FROM, | |
| 11 | + PAGINATION, | |
| 12 | + TYPE_LABELS, | |
| 13 | + ApiError, | |
| 14 | + Pagination, | |
| 15 | + cached, | |
| 16 | + entity_summary, | |
| 17 | + page, | |
| 18 | +) | |
| 19 | +from aiatlas.api.routers.entities import detail_for_type | |
| 20 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 21 | +from aiatlas.ids import ENTITY_TYPES | |
| 22 | + | |
| 23 | +router = APIRouter(prefix="/api/v1", tags=["explore"]) | |
| 24 | +SORTS = {"updated": "e.updated_at desc", "name": "e.canonical_name asc", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "first_seen": "e.first_seen_at desc", | |
| 25 | + "release": "coalesce(e.attributes->>'release_date', e.attributes->>'published_at', e.attributes->>'latest_release_at') desc nulls last", | |
| 26 | + "stars": "(case when e.attributes->>'metric.stars' ~ '^[0-9]+$' then (e.attributes->>'metric.stars')::bigint end) desc nulls last"} | |
| 27 | + | |
| 28 | + | |
| 29 | +@router.get("/explore/types") | |
| 30 | +@cached(300) | |
| 31 | +async def explore_types(request: Request) -> dict[str, Any]: | |
| 32 | + async with connection() as conn: | |
| 33 | + rows = await fetch_all(conn, "select entity_type, count(*) as count from entities where merged_into is null group by 1 order by 2 desc, 1") | |
| 34 | + return {"items": [{"entity_type": r["entity_type"], "count": int(r["count"]), "label": TYPE_LABELS.get(r["entity_type"], r["entity_type"].replace("_", " ").title())} for r in rows]} | |
| 35 | + | |
| 36 | + | |
| 37 | +@router.get("/explore/{type}") | |
| 38 | +@cached(300) | |
| 39 | +async def explore_type(request: Request, type: str, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, status: str | None = None, | |
| 40 | + sort: str = "updated") -> dict[str, Any]: | |
| 41 | + etype = type.strip().lower().rstrip("s") if type not in ENTITY_TYPES else type | |
| 42 | + aliases = {"companie": "company", "librarie": "library", "universitie": "university", "regulations": "regulation", "hardware": "hardware"} | |
| 43 | + etype = aliases.get(etype, etype) | |
| 44 | + if etype not in ENTITY_TYPES: | |
| 45 | + raise ApiError(404, f"unknown entity type {type!r}") | |
| 46 | + if sort not in SORTS: | |
| 47 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 48 | + where = ["e.entity_type = :t", "e.merged_into is null"] | |
| 49 | + params: dict[str, Any] = {"t": etype} | |
| 50 | + if q: | |
| 51 | + where.append("(e.canonical_name ilike :qlike or e.search @@ plainto_tsquery('simple', :q) or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))") | |
| 52 | + params["q"], params["qlike"] = q, f"%{q}%" | |
| 53 | + if org: | |
| 54 | + where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)") | |
| 55 | + params["org"] = org | |
| 56 | + if status: | |
| 57 | + where.append("e.status = any(cast(:status as text[]))") | |
| 58 | + params["status"] = [s.strip() for s in status.split(",") if s.strip()] | |
| 59 | + where_sql = " and ".join(where) | |
| 60 | + async with connection() as conn: | |
| 61 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) | |
| 62 | + total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) | |
| 63 | + out = page([entity_summary(r) for r in rows], int(total or 0), p) | |
| 64 | + out["entity_type"] = etype | |
| 65 | + out["label"] = TYPE_LABELS.get(etype, etype.title()) | |
| 66 | + return out | |
| 67 | + | |
| 68 | + | |
| 69 | +@router.get("/frameworks/{slug}") | |
| 70 | +@cached(300) | |
| 71 | +async def get_framework(request: Request, slug: str) -> dict[str, Any]: | |
| 72 | + return await detail_for_type(slug, ("framework", "library", "runtime")) | |
| 73 | + | |
| 74 | + | |
| 75 | +@router.get("/datasets/{slug}") | |
| 76 | +@cached(300) | |
| 77 | +async def get_dataset(request: Request, slug: str) -> dict[str, Any]: | |
| 78 | + return await detail_for_type(slug, ("dataset",)) | |
| 79 | + | |
| 80 | + | |
| 81 | +@router.get("/tools/{slug}") | |
| 82 | +@cached(300) | |
| 83 | +async def get_tool(request: Request, slug: str) -> dict[str, Any]: | |
| 84 | + return await detail_for_type(slug, ("tool", "agent", "application", "product", "mcp_server")) | |
added
src/aiatlas/api/routers/hardware.py
+92 −0
@@ -0,0 +1,92 @@ | ||
| 1 | +"""/hardware listing, /hardware/fit (ESTIMATED memory fit) and /hardware/{slug} alias.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + ENTITY_COLS, | |
| 10 | + ENTITY_FROM, | |
| 11 | + PAGINATION, | |
| 12 | + ApiError, | |
| 13 | + Pagination, | |
| 14 | + attr_num, | |
| 15 | + cached, | |
| 16 | + entity_summary, | |
| 17 | + page, | |
| 18 | +) | |
| 19 | +from aiatlas.api.routers.entities import detail_for_type | |
| 20 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 21 | +from aiatlas.services import hardware_fit as hf | |
| 22 | + | |
| 23 | +router = APIRouter(prefix="/api/v1/hardware", tags=["hardware"]) | |
| 24 | + | |
| 25 | +MEMORY = ("(case when jsonb_typeof(e.attributes->'memory_gb') = 'array' then (select max(x::text::double precision) from jsonb_array_elements(e.attributes->'memory_gb') x where jsonb_typeof(x) = 'number') " | |
| 26 | + "when jsonb_typeof(e.attributes->'memory_gb') = 'number' then (e.attributes->>'memory_gb')::double precision else " + attr_num("memory_gb") + " end)") | |
| 27 | +SORTS = {"memory": f"{MEMORY} desc nulls last", "name": "e.canonical_name asc", "release": "e.attributes->>'release_date' desc nulls last", "updated": "e.updated_at desc", | |
| 28 | + "bandwidth": attr_num("memory_bandwidth_gbs") + " desc nulls last"} | |
| 29 | + | |
| 30 | + | |
| 31 | +@router.get("") | |
| 32 | +@cached(300) | |
| 33 | +async def list_hardware(request: Request, p: Pagination = PAGINATION, kind: str | None = None, manufacturer: str | None = None, min_memory: float | None = Query(None, ge=0), | |
| 34 | + q: str | None = Query(None, max_length=200), sort: str = "memory") -> dict[str, Any]: | |
| 35 | + if sort not in SORTS: | |
| 36 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 37 | + where = ["e.entity_type = 'hardware'", "e.merged_into is null"] | |
| 38 | + params: dict[str, Any] = {} | |
| 39 | + if kind: | |
| 40 | + where.append("e.attributes->>'kind' ilike :kind") | |
| 41 | + params["kind"] = kind | |
| 42 | + if manufacturer: | |
| 43 | + where.append("(e.attributes->>'manufacturer' ilike :man or eo.slug = :man or eo.canonical_name ilike :man)") | |
| 44 | + params["man"] = manufacturer | |
| 45 | + if min_memory is not None: | |
| 46 | + where.append(f"{MEMORY} >= :min_memory") | |
| 47 | + params["min_memory"] = float(min_memory) | |
| 48 | + if q: | |
| 49 | + where.append("e.canonical_name ilike :qlike") | |
| 50 | + params["qlike"] = f"%{q}%" | |
| 51 | + where_sql = " and ".join(where) | |
| 52 | + async with connection() as conn: | |
| 53 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) | |
| 54 | + total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) | |
| 55 | + kinds = await fetch_all(conn, "select e.attributes->>'kind' as value, count(*) as count from entities e where e.entity_type = 'hardware' and e.merged_into is null and e.attributes ? 'kind' group by 1 order by 2 desc") | |
| 56 | + mans = await fetch_all(conn, "select e.attributes->>'manufacturer' as value, count(*) as count from entities e where e.entity_type = 'hardware' and e.merged_into is null and e.attributes ? 'manufacturer' group by 1 order by 2 desc") | |
| 57 | + out = page([entity_summary(r) for r in rows], int(total or 0), p) | |
| 58 | + out["facets"] = {"kinds": [{"value": r["value"], "count": int(r["count"])} for r in kinds], "manufacturers": [{"value": r["value"], "count": int(r["count"])} for r in mans]} | |
| 59 | + return out | |
| 60 | + | |
| 61 | + | |
| 62 | +@router.get("/fit") | |
| 63 | +@cached(300) | |
| 64 | +async def hardware_fit(request: Request, memory_gb: float = Query(..., gt=0, le=100000), quant: str = Query("4bit"), context: int = Query(8192, ge=0, le=10_000_000), | |
| 65 | + limit: int = Query(100, ge=1, le=500), openness: str | None = None) -> dict[str, Any]: | |
| 66 | + if quant not in hf.BYTES_PER_PARAM: | |
| 67 | + raise ApiError(400, f"quant must be one of {', '.join(hf.BYTES_PER_PARAM)}") | |
| 68 | + where = "e.entity_type = 'model' and e.merged_into is null and e.attributes ? 'parameter_count'" | |
| 69 | + params: dict[str, Any] = {} | |
| 70 | + if openness: | |
| 71 | + vals = [v.strip() for v in openness.split(",") if v.strip()] | |
| 72 | + if "open" in vals: | |
| 73 | + vals += ["open-weights", "open-source"] | |
| 74 | + where += " and e.attributes->>'openness' = any(cast(:openness as text[]))" | |
| 75 | + params["openness"] = vals | |
| 76 | + async with connection() as conn: | |
| 77 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where} order by {attr_num('parameter_count')} desc nulls last limit 2000", **params) | |
| 78 | + items = [] | |
| 79 | + for r in rows: | |
| 80 | + pc = hf.parameter_count(r["attributes"]) | |
| 81 | + if pc is None: | |
| 82 | + continue | |
| 83 | + items.append({"model": entity_summary(r), "parameter_count": pc, **hf.fit(pc, memory_gb, quant, context)}) | |
| 84 | + items.sort(key=lambda x: (not x["fits"], -x["parameter_count"] if x["fits"] else x["parameter_count"])) | |
| 85 | + return {"inputs": {"memory_gb": memory_gb, "quant": quant, "context": context}, "estimated": True, "assumptions": hf.ASSUMPTIONS, | |
| 86 | + "counts": {"fits": sum(1 for i in items if i["fits"]), "evaluated": len(items)}, "items": items[:limit]} | |
| 87 | + | |
| 88 | + | |
| 89 | +@router.get("/{slug}") | |
| 90 | +@cached(300) | |
| 91 | +async def get_hardware(request: Request, slug: str) -> dict[str, Any]: | |
| 92 | + return await detail_for_type(slug, ("hardware",)) | |
added
src/aiatlas/api/routers/misc.py
+101 −0
@@ -0,0 +1,101 @@ | ||
| 1 | +"""/methodology · /trending · POST /views · /sitemap · /api-keys/me.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import hashlib | |
| 5 | +from datetime import UTC, datetime | |
| 6 | +from typing import Any | |
| 7 | + | |
| 8 | +from fastapi import APIRouter, Depends, Query, Request | |
| 9 | +from pydantic import BaseModel, Field | |
| 10 | + | |
| 11 | +from aiatlas.api.common import ENTITY_COLS, ApiError, cached, entity_summary, rate_limit | |
| 12 | +from aiatlas.db import connection, execute, fetch_all, fetch_one, transaction | |
| 13 | +from aiatlas.ids import ENTITY_TYPES | |
| 14 | +from aiatlas.services.quality import EXPECTED_FIELDS, QUALITY_VERSION | |
| 15 | + | |
| 16 | +router = APIRouter(prefix="/api/v1", tags=["misc"]) | |
| 17 | + | |
| 18 | +CONFIDENCE_LEVELS = [ | |
| 19 | + {"key": "verified", "label": "Verified", "description": "Confirmed by at least two independent sources, one of them tier 1."}, | |
| 20 | + {"key": "high", "label": "High", "description": "Stated explicitly by an official (tier 1) source and extracted deterministically."}, | |
| 21 | + {"key": "medium", "label": "Medium", "description": "Quality secondary source, curated registry, or LLM extraction from an official document."}, | |
| 22 | + {"key": "low", "label": "Low", "description": "Community or unverified source, or LLM extraction from a secondary document."}, | |
| 23 | + {"key": "conflicted", "label": "Conflicted", "description": "Another source states a different value; both claims are kept and flagged, never averaged."}, | |
| 24 | +] | |
| 25 | +TIERS = [ | |
| 26 | + {"tier": 1, "label": "Official / primary", "description": "The organization's own site, docs, pricing pages, model cards, filings."}, | |
| 27 | + {"tier": 2, "label": "Quality secondary", "description": "Peer-reviewed venues, arXiv, curated registries, major leaderboards."}, | |
| 28 | + {"tier": 3, "label": "Community", "description": "Community-maintained hubs, forums, wikis."}, | |
| 29 | + {"tier": 4, "label": "Unverified", "description": "Anything else; never overrides a better tier."}, | |
| 30 | +] | |
| 31 | +EXTRACTORS = [ | |
| 32 | + {"key": "deterministic", "description": "Rule-based parsers (tables, JSON-LD, meta tags, embedded JSON, regex) — always runs first."}, | |
| 33 | + {"key": "curated", "description": "Hand-maintained registries shipped with the code (organizations, providers, benchmarks, hardware)."}, | |
| 34 | + {"key": "llm", "description": "Local LLM extraction validated against a strict JSON schema; medium/low confidence, never overrides deterministic tier-1 claims."}, | |
| 35 | +] | |
| 36 | + | |
| 37 | + | |
| 38 | +@router.get("/methodology") | |
| 39 | +@cached(600) | |
| 40 | +async def methodology(request: Request) -> dict[str, Any]: | |
| 41 | + async with connection() as conn: | |
| 42 | + metrics = await fetch_all(conn, "select key, label, version, description, formula from metric_definitions order by key") | |
| 43 | + event_types = await fetch_all(conn, "select event_type, category, count(*) as count, max(observed_at) as last_seen_at from change_events group by 1, 2 order by 3 desc") | |
| 44 | + return {"metrics": metrics, "quality_version": QUALITY_VERSION, "expected_fields": EXPECTED_FIELDS, "confidence_levels": CONFIDENCE_LEVELS, "tiers": TIERS, | |
| 45 | + "event_types": [{**r, "count": int(r["count"])} for r in event_types], "extractors": EXTRACTORS, | |
| 46 | + "principles": ["Never fabricate: missing data is reported as unavailable.", "Every fact carries provenance (source, snapshot, URL, tier, confidence, extractor).", | |
| 47 | + "History is append-only: claims, prices and benchmark results are never overwritten.", | |
| 48 | + "Conflicts between sources are stored side by side and flagged for review.", "Live counters and feeds are computed from the database."]} | |
| 49 | + | |
| 50 | + | |
| 51 | +@router.get("/trending") | |
| 52 | +@cached(300) | |
| 53 | +async def trending(request: Request, days: int = Query(7, ge=1, le=90), limit: int = Query(12, ge=1, le=60), type: str | None = Query(None, alias="type")) -> dict[str, Any]: | |
| 54 | + async with connection() as conn: | |
| 55 | + rows = await fetch_all(conn, f""" | |
| 56 | + with v as (select split_part(regexp_replace(path, '[?#].*$', ''), '/', 3) as slug, sum(views) as views from page_views | |
| 57 | + where day >= ((now() at time zone 'UTC') - make_interval(days => :d))::date and path ~ '^/[a-z-]+/[^/?#]+' group by 1) | |
| 58 | + select v.views, {ENTITY_COLS} from v join entities e on e.slug = v.slug left join entities eo on eo.id = e.organization_id | |
| 59 | + where e.merged_into is null {"and e.entity_type = :t" if type else ""} order by v.views desc, e.updated_at desc limit :lim""", d=days, lim=limit, t=type) | |
| 60 | + return {"days": days, "items": [{**(entity_summary(r) or {}), "views": int(r["views"] or 0)} for r in rows]} | |
| 61 | + | |
| 62 | + | |
| 63 | +class ViewBeacon(BaseModel): | |
| 64 | + path: str = Field(..., min_length=1, max_length=300) | |
| 65 | + | |
| 66 | + | |
| 67 | +@router.post("/views", dependencies=[Depends(rate_limit("views"))]) | |
| 68 | +async def record_view(body: ViewBeacon) -> dict[str, Any]: | |
| 69 | + path = body.path.strip() | |
| 70 | + if not path.startswith("/") or "\n" in path or "//" in path: | |
| 71 | + raise ApiError(400, "path must be a site-relative path") | |
| 72 | + path = path.split("?", 1)[0].split("#", 1)[0][:300] | |
| 73 | + async with transaction() as conn: | |
| 74 | + await execute(conn, "insert into page_views (path, day, views) values (:p, :d, 1) on conflict (path, day) do update set views = page_views.views + 1", | |
| 75 | + p=path, d=datetime.now(UTC).date()) | |
| 76 | + return {"ok": True} | |
| 77 | + | |
| 78 | + | |
| 79 | +@router.get("/sitemap") | |
| 80 | +@cached(600) | |
| 81 | +async def sitemap(request: Request, type: str | None = Query(None, alias="type"), limit: int = Query(5000, ge=1, le=5000), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 82 | + if type and type not in ENTITY_TYPES: | |
| 83 | + raise ApiError(400, f"unknown entity type {type!r}") | |
| 84 | + where = "e.merged_into is null" + (" and e.entity_type = :t" if type else "") | |
| 85 | + async with connection() as conn: | |
| 86 | + rows = await fetch_all(conn, f"select e.slug, e.entity_type, e.updated_at from entities e where {where} order by e.updated_at desc, e.id limit :lim offset :off", t=type, lim=limit, off=offset) | |
| 87 | + total = await fetch_one(conn, f"select count(*) as n from entities e where {where}", t=type) | |
| 88 | + return {"items": rows, "total": int(total["n"]) if total else 0, "limit": limit, "offset": offset} | |
| 89 | + | |
| 90 | + | |
| 91 | +@router.get("/api-keys/me") | |
| 92 | +async def api_key_me(request: Request) -> dict[str, Any]: | |
| 93 | + key = request.headers.get("x-api-key") or "" | |
| 94 | + if not key: | |
| 95 | + raise ApiError(401, "x-api-key header required") | |
| 96 | + digest = hashlib.sha256(key.encode()).hexdigest() | |
| 97 | + async with transaction() as conn: | |
| 98 | + row = await fetch_one(conn, "update api_keys set last_used_at = now(), usage_count = usage_count + 1 where key_hash = :h and enabled returning label, plan, rate_per_min, usage_count, created_at, last_used_at", h=digest) | |
| 99 | + if not row: | |
| 100 | + raise ApiError(401, "unknown or disabled API key") | |
| 101 | + return row | |
added
src/aiatlas/api/routers/models.py
+143 −0
@@ -0,0 +1,143 @@ | ||
| 1 | +"""/models listing (filters on `entities.attributes`, facets) and /models/{slug} alias.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from fastapi import APIRouter, Query, Request | |
| 8 | + | |
| 9 | +from aiatlas.api.common import ( | |
| 10 | + ENTITY_COLS, | |
| 11 | + ENTITY_FROM, | |
| 12 | + PAGINATION, | |
| 13 | + ApiError, | |
| 14 | + Pagination, | |
| 15 | + cached, | |
| 16 | + entity_summary, | |
| 17 | + flip_order, | |
| 18 | + normalize, | |
| 19 | + num_expr, | |
| 20 | + page, | |
| 21 | +) | |
| 22 | +from aiatlas.api.routers.entities import detail_for_type | |
| 23 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 24 | +from aiatlas.services import cache | |
| 25 | + | |
| 26 | +router = APIRouter(prefix="/api/v1/models", tags=["models"]) | |
| 27 | + | |
| 28 | +PARAMS = num_expr("e.attributes->>'parameter_count'") | |
| 29 | +CONTEXT = num_expr("e.attributes->>'context_length'") | |
| 30 | +DOWNLOADS = num_expr("e.attributes->>'metric.downloads'") | |
| 31 | +SORTS = { | |
| 32 | + "updated": "e.updated_at desc", "name": "e.canonical_name asc", "params": f"{PARAMS} desc nulls last", "context": f"{CONTEXT} desc nulls last", | |
| 33 | + "release": "e.attributes->>'release_date' desc nulls last", "quality": "coalesce((e.quality->>'score')::float, 0) desc", "downloads": f"{DOWNLOADS} desc nulls last", | |
| 34 | + "first_seen": "e.first_seen_at desc", | |
| 35 | +} | |
| 36 | + | |
| 37 | + | |
| 38 | +def model_filters(*, q: str | None, org: str | None, family: str | None, openness: str | None, modality: str | None, status: str | None, | |
| 39 | + min_params: float | None, max_params: float | None, min_context: int | None, year_from: int | None, year_to: int | None, | |
| 40 | + license: str | None) -> tuple[list[str], dict[str, Any]]: | |
| 41 | + where = ["e.entity_type = 'model'", "e.merged_into is null"] | |
| 42 | + p: dict[str, Any] = {} | |
| 43 | + if q: | |
| 44 | + where.append("(e.canonical_name ilike :qlike or e.search @@ plainto_tsquery('simple', :q) or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))") | |
| 45 | + p["q"], p["qlike"] = q, f"%{q}%" | |
| 46 | + if org: | |
| 47 | + where.append("(eo.slug = :org or eo.id = :org or eo.canonical_name ilike :org)") | |
| 48 | + p["org"] = org | |
| 49 | + if family: | |
| 50 | + where.append("e.attributes->>'family' ilike :family") | |
| 51 | + p["family"] = family | |
| 52 | + if openness: | |
| 53 | + vals = [v.strip() for v in openness.split(",") if v.strip()] | |
| 54 | + if "open" in vals: | |
| 55 | + vals += ["open-weights", "open-source"] | |
| 56 | + where.append("e.attributes->>'openness' = any(cast(:openness as text[]))") | |
| 57 | + p["openness"] = vals | |
| 58 | + if modality: | |
| 59 | + where.append("(e.attributes->'modalities' ? :modality or e.attributes->'modalities_input' ? :modality or e.attributes->'modalities_output' ? :modality)") | |
| 60 | + p["modality"] = modality | |
| 61 | + if status: | |
| 62 | + where.append("e.status = any(cast(:status as text[]))") | |
| 63 | + p["status"] = [v.strip() for v in status.split(",") if v.strip()] | |
| 64 | + if min_params is not None: | |
| 65 | + where.append(f"{PARAMS} >= :min_params") | |
| 66 | + p["min_params"] = float(min_params) | |
| 67 | + if max_params is not None: | |
| 68 | + where.append(f"{PARAMS} <= :max_params") | |
| 69 | + p["max_params"] = float(max_params) | |
| 70 | + if min_context is not None: | |
| 71 | + where.append(f"{CONTEXT} >= :min_context") | |
| 72 | + p["min_context"] = float(min_context) | |
| 73 | + if year_from is not None: | |
| 74 | + where.append("left(e.attributes->>'release_date', 4) >= :yf") | |
| 75 | + p["yf"] = str(year_from) | |
| 76 | + if year_to is not None: | |
| 77 | + where.append("left(e.attributes->>'release_date', 4) <= :yt") | |
| 78 | + p["yt"] = str(year_to) | |
| 79 | + if license: | |
| 80 | + where.append("e.attributes->>'license' ilike :license") | |
| 81 | + p["license"] = license | |
| 82 | + return where, p | |
| 83 | + | |
| 84 | + | |
| 85 | +async def model_facets(where_sql: str, params: dict[str, Any]) -> dict[str, Any]: | |
| 86 | + async def run(sql: str) -> list[dict[str, Any]]: | |
| 87 | + async with connection() as conn: | |
| 88 | + return await fetch_all(conn, sql, **params) | |
| 89 | + | |
| 90 | + base = f"from {ENTITY_FROM} where {where_sql}" | |
| 91 | + mods_from = (f"from {ENTITY_FROM} cross join lateral jsonb_array_elements_text(case when jsonb_typeof(e.attributes->'modalities') = 'array' " | |
| 92 | + f"then e.attributes->'modalities' else '[]'::jsonb end) m where {where_sql}") | |
| 93 | + orgs, openness, mods, fams, years, lics, status = await asyncio.gather( | |
| 94 | + run(f"select eo.slug, eo.canonical_name as name, count(*) as count {base} and eo.id is not null group by 1, 2 order by 3 desc, 2 limit 60"), | |
| 95 | + run(f"select e.attributes->>'openness' as value, count(*) as count {base} and e.attributes ? 'openness' group by 1 order by 2 desc"), | |
| 96 | + run(f"select m.value, count(*) as count {mods_from} group by 1 order by 2 desc limit 30"), | |
| 97 | + run(f"select e.attributes->>'family' as value, count(*) as count {base} and e.attributes ? 'family' group by 1 order by 2 desc, 1 limit 60"), | |
| 98 | + run(f"select left(e.attributes->>'release_date', 4) as value, count(*) as count {base} and e.attributes ? 'release_date' group by 1 order by 1 desc limit 30"), | |
| 99 | + run(f"select e.attributes->>'license' as value, count(*) as count {base} and e.attributes ? 'license' group by 1 order by 2 desc, 1 limit 40"), | |
| 100 | + run(f"select e.status as value, count(*) as count {base} group by 1 order by 2 desc"), | |
| 101 | + ) | |
| 102 | + conv = lambda rows: [{"value": r["value"], "count": int(r["count"])} for r in rows if r["value"] not in (None, "")] | |
| 103 | + return {"organizations": [{"slug": r["slug"], "name": r["name"], "count": int(r["count"])} for r in orgs], "openness": conv(openness), "modalities": conv(mods), | |
| 104 | + "families": conv(fams), "years": conv(years), "licenses": conv(lics), "status": conv(status)} | |
| 105 | + | |
| 106 | + | |
| 107 | +@router.get("") | |
| 108 | +@cached(300) | |
| 109 | +async def list_models(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), org: str | None = None, family: str | None = None, | |
| 110 | + openness: str | None = None, modality: str | None = None, status: str | None = None, min_params: float | None = Query(None, ge=0), | |
| 111 | + max_params: float | None = Query(None, ge=0), min_context: int | None = Query(None, ge=0), year_from: int | None = Query(None, ge=1950, le=2100), | |
| 112 | + year_to: int | None = Query(None, ge=1950, le=2100), license: str | None = None, sort: str = "updated", order: str = "", | |
| 113 | + facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]: | |
| 114 | + if sort not in SORTS: | |
| 115 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 116 | + where, params = model_filters(q=q, org=org, family=family, openness=openness, modality=modality, status=status, min_params=min_params, max_params=max_params, | |
| 117 | + min_context=min_context, year_from=year_from, year_to=year_to, license=license) | |
| 118 | + order_sql = flip_order(SORTS[sort], order) | |
| 119 | + where_sql = " and ".join(where) | |
| 120 | + async with connection() as conn: | |
| 121 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {order_sql}, e.id limit :lim offset :off", | |
| 122 | + lim=p.limit, off=p.offset, **params) | |
| 123 | + total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) | |
| 124 | + out = page([entity_summary(r) for r in rows], int(total or 0), p) | |
| 125 | + if facets: | |
| 126 | + out["facets"] = await _facets_cached(where_sql, params) | |
| 127 | + return out | |
| 128 | + | |
| 129 | + | |
| 130 | +async def _facets_cached(where_sql: str, params: dict[str, Any]) -> dict[str, Any]: | |
| 131 | + key = "facets:models:" + where_sql + ":" + repr(sorted(params.items())) | |
| 132 | + hit = await cache.cache_get(key) | |
| 133 | + if hit is not None: | |
| 134 | + return hit | |
| 135 | + value = normalize(await model_facets(where_sql, params)) | |
| 136 | + await cache.cache_set(key, value, 600) | |
| 137 | + return value | |
| 138 | + | |
| 139 | + | |
| 140 | +@router.get("/{slug}") | |
| 141 | +@cached(300) | |
| 142 | +async def get_model(request: Request, slug: str) -> dict[str, Any]: | |
| 143 | + return await detail_for_type(slug, ("model",)) | |
added
src/aiatlas/api/routers/papers.py
+60 −0
@@ -0,0 +1,60 @@ | ||
| 1 | +"""/papers listing and /papers/{slug} alias.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + ENTITY_COLS, | |
| 10 | + ENTITY_FROM, | |
| 11 | + PAGINATION, | |
| 12 | + ApiError, | |
| 13 | + Pagination, | |
| 14 | + cached, | |
| 15 | + entity_summary, | |
| 16 | + page, | |
| 17 | + parse_date, | |
| 18 | +) | |
| 19 | +from aiatlas.api.routers.entities import detail_for_type | |
| 20 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 21 | + | |
| 22 | +router = APIRouter(prefix="/api/v1/papers", tags=["papers"]) | |
| 23 | +SORTS = {"published": "e.attributes->>'published_at' desc nulls last", "updated": "e.updated_at desc", "name": "e.canonical_name asc", | |
| 24 | + "citations": "(case when e.attributes->>'metric.citations' ~ '^[0-9]+$' then (e.attributes->>'metric.citations')::bigint end) desc nulls last"} | |
| 25 | + | |
| 26 | + | |
| 27 | +@router.get("") | |
| 28 | +@cached(300) | |
| 29 | +async def list_papers(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), category: str | None = None, org: str | None = None, | |
| 30 | + since: str | None = None, until: str | None = None, sort: str = "published") -> dict[str, Any]: | |
| 31 | + if sort not in SORTS: | |
| 32 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 33 | + where = ["e.entity_type = 'paper'", "e.merged_into is null"] | |
| 34 | + params: dict[str, Any] = {} | |
| 35 | + if q: | |
| 36 | + where.append("(e.canonical_name ilike :qlike or e.search @@ plainto_tsquery('english', :q))") | |
| 37 | + params["q"], params["qlike"] = q, f"%{q}%" | |
| 38 | + if category: | |
| 39 | + where.append("(e.attributes->>'primary_category' = :cat or e.attributes->'categories' ? :cat)") | |
| 40 | + params["cat"] = category | |
| 41 | + if org: | |
| 42 | + where.append("(eo.slug = :org or eo.id = :org or exists (select 1 from relations r join entities x on x.id = r.object_id where r.subject_id = e.id and r.valid_to is null and (x.slug = :org or x.id = :org)))") | |
| 43 | + params["org"] = org | |
| 44 | + if since: | |
| 45 | + where.append("e.attributes->>'published_at' >= :since") | |
| 46 | + params["since"] = parse_date(since, "since").isoformat() # type: ignore[union-attr] | |
| 47 | + if until: | |
| 48 | + where.append("e.attributes->>'published_at' <= :until") | |
| 49 | + params["until"] = parse_date(until, "until").isoformat() # type: ignore[union-attr] | |
| 50 | + where_sql = " and ".join(where) | |
| 51 | + async with connection() as conn: | |
| 52 | + rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) | |
| 53 | + total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params) | |
| 54 | + return page([entity_summary(r) for r in rows], int(total or 0), p) | |
| 55 | + | |
| 56 | + | |
| 57 | +@router.get("/{slug}") | |
| 58 | +@cached(300) | |
| 59 | +async def get_paper(request: Request, slug: str) -> dict[str, Any]: | |
| 60 | + return await detail_for_type(slug, ("paper",)) | |
added
src/aiatlas/api/routers/prices.py
+90 −0
@@ -0,0 +1,90 @@ | ||
| 1 | +"""/prices · /prices/history · /prices/index — append-only pricing table (`valid_to` closes a row).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ( | |
| 9 | + EVENT_COLS, | |
| 10 | + EVENT_FROM, | |
| 11 | + PAGINATION, | |
| 12 | + PRICE_COLS, | |
| 13 | + PRICE_FROM, | |
| 14 | + ApiError, | |
| 15 | + Pagination, | |
| 16 | + cached, | |
| 17 | + change_event, | |
| 18 | + page, | |
| 19 | + price_row, | |
| 20 | + resolve_id, | |
| 21 | +) | |
| 22 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 23 | + | |
| 24 | +router = APIRouter(prefix="/api/v1/prices", tags=["prices"]) | |
| 25 | +SORTS = {"input": "p.input_per_mtok asc nulls last", "output": "p.output_per_mtok asc nulls last", "model": "m.canonical_name asc", "provider": "pv.canonical_name asc", | |
| 26 | + "observed": "p.observed_at desc", "valid_from": "p.valid_from desc"} | |
| 27 | + | |
| 28 | + | |
| 29 | +@router.get("") | |
| 30 | +@cached(300) | |
| 31 | +async def list_prices(request: Request, p: Pagination = PAGINATION, model: str | None = None, provider: str | None = None, sort: str = "model", order: str = "", | |
| 32 | + current: int = Query(1, ge=0, le=1)) -> dict[str, Any]: | |
| 33 | + if sort not in SORTS: | |
| 34 | + raise ApiError(400, f"sort must be one of {', '.join(SORTS)}") | |
| 35 | + order_sql = SORTS[sort] | |
| 36 | + if order == "desc" and " asc" in order_sql: | |
| 37 | + order_sql = order_sql.replace(" asc", " desc") | |
| 38 | + elif order == "asc" and " desc" in order_sql: | |
| 39 | + order_sql = order_sql.replace(" desc", " asc") | |
| 40 | + where = ["p.valid_to is null"] if current else ["true"] | |
| 41 | + params: dict[str, Any] = {} | |
| 42 | + async with connection() as conn: | |
| 43 | + if model: | |
| 44 | + where.append("p.model_id = :model") | |
| 45 | + params["model"] = await resolve_id(conn, model) | |
| 46 | + if provider: | |
| 47 | + where.append("p.provider_id = :provider") | |
| 48 | + params["provider"] = await resolve_id(conn, provider) | |
| 49 | + where_sql = " and ".join(where) | |
| 50 | + rows = await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where {where_sql} order by {order_sql}, p.id limit :lim offset :off", lim=p.limit, off=p.offset, **params) | |
| 51 | + total = await fetch_val(conn, f"select count(*) from prices p where {where_sql}", **params) | |
| 52 | + return page([price_row(r) for r in rows], int(total or 0), p) | |
| 53 | + | |
| 54 | + | |
| 55 | +@router.get("/history") | |
| 56 | +@cached(300) | |
| 57 | +async def price_history(request: Request, model: str | None = None, provider: str | None = None, limit: int = Query(2000, ge=1, le=5000)) -> dict[str, Any]: | |
| 58 | + if not model and not provider: | |
| 59 | + raise ApiError(400, "model or provider is required") | |
| 60 | + where = ["true"] | |
| 61 | + params: dict[str, Any] = {"lim": limit} | |
| 62 | + async with connection() as conn: | |
| 63 | + if model: | |
| 64 | + where.append("p.model_id = :model") | |
| 65 | + params["model"] = await resolve_id(conn, model) | |
| 66 | + if provider: | |
| 67 | + where.append("p.provider_id = :provider") | |
| 68 | + params["provider"] = await resolve_id(conn, provider) | |
| 69 | + rows = await fetch_all(conn, f"select {PRICE_COLS} from {PRICE_FROM} where {' and '.join(where)} order by p.valid_from asc, p.id limit :lim", **params) | |
| 70 | + return {"items": [price_row(r) for r in rows]} | |
| 71 | + | |
| 72 | + | |
| 73 | +@router.get("/index") | |
| 74 | +@cached(600) | |
| 75 | +async def price_index(request: Request, days: int = Query(180, ge=7, le=1825)) -> dict[str, Any]: | |
| 76 | + async with connection() as conn: | |
| 77 | + series = await fetch_all(conn, """ | |
| 78 | + with days as (select generate_series(((now() at time zone 'UTC') - make_interval(days => :d))::date::timestamp, (now() at time zone 'UTC')::date::timestamp, interval '1 day')::date as day), | |
| 79 | + live as (select p.model_id, p.input_per_mtok, p.output_per_mtok, p.valid_from, p.valid_to from prices p | |
| 80 | + where p.input_per_mtok is not null and p.input_per_mtok > 0 and p.valid_from < now()) | |
| 81 | + select d.day, percentile_cont(0.5) within group (order by l.input_per_mtok) as median_input, | |
| 82 | + percentile_cont(0.5) within group (order by l.output_per_mtok) as median_output, | |
| 83 | + min(l.input_per_mtok) as min_input, max(l.input_per_mtok) as max_input, count(distinct l.model_id) as models, count(l.model_id) as offers | |
| 84 | + from days d left join live l on l.valid_from < (d.day + 1)::timestamp at time zone 'UTC' and (l.valid_to is null or l.valid_to >= (d.day + 1)::timestamp at time zone 'UTC') | |
| 85 | + group by d.day order by d.day""", d=days) | |
| 86 | + movers = await fetch_all(conn, f"select {EVENT_COLS} from {EVENT_FROM} where ev.category = 'price' and ev.observed_at > now() - make_interval(days => :d) " | |
| 87 | + f"order by ev.importance desc, ev.observed_at desc limit 30", d=days) | |
| 88 | + return {"days": days, "series": [{"day": r["day"].isoformat(), "median_input": r["median_input"], "median_output": r["median_output"], "min_input": r["min_input"], | |
| 89 | + "max_input": r["max_input"], "models": int(r["models"] or 0), "offers": int(r["offers"] or 0)} for r in series], | |
| 90 | + "movers": [change_event(r) for r in movers], "note": "Daily medians of live USD-per-1M-token input/output prices across all provider offers valid at the end of each UTC day."} | |
added
src/aiatlas/api/routers/providers.py
+35 −0
@@ -0,0 +1,35 @@ | ||
| 1 | +"""/providers listing (with pricing aggregates) and /providers/{slug} alias.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, cached, entity_summary | |
| 9 | +from aiatlas.api.routers.entities import detail_for_type | |
| 10 | +from aiatlas.db import connection, fetch_all | |
| 11 | + | |
| 12 | +router = APIRouter(prefix="/api/v1/providers", tags=["providers"]) | |
| 13 | + | |
| 14 | + | |
| 15 | +@router.get("") | |
| 16 | +@cached(300) | |
| 17 | +async def list_providers(request: Request) -> dict[str, Any]: | |
| 18 | + async with connection() as conn: | |
| 19 | + rows = await fetch_all(conn, f""" | |
| 20 | + select {ENTITY_COLS}, | |
| 21 | + (select count(distinct p.model_id) from prices p where p.provider_id = e.id and p.valid_to is null) as priced_models, | |
| 22 | + (select count(*) from relations r where r.object_id = e.id and r.predicate = 'available_through' and r.valid_to is null) as listed_models, | |
| 23 | + (select count(*) from prices p where p.provider_id = e.id and p.valid_to is null) as price_count, | |
| 24 | + (select min(p.input_per_mtok) from prices p where p.provider_id = e.id and p.valid_to is null and p.input_per_mtok > 0) as min_input_per_mtok, | |
| 25 | + (select min(p.output_per_mtok) from prices p where p.provider_id = e.id and p.valid_to is null and p.output_per_mtok > 0) as min_output_per_mtok | |
| 26 | + from {ENTITY_FROM} where e.entity_type = 'provider' and e.merged_into is null | |
| 27 | + order by price_count desc, listed_models desc, e.canonical_name""") | |
| 28 | + return {"items": [{**(entity_summary(r) or {}), "model_count": int(max(r["priced_models"] or 0, r["listed_models"] or 0)), "price_count": int(r["price_count"] or 0), | |
| 29 | + "min_input_per_mtok": r["min_input_per_mtok"], "min_output_per_mtok": r["min_output_per_mtok"]} for r in rows]} | |
| 30 | + | |
| 31 | + | |
| 32 | +@router.get("/{slug}") | |
| 33 | +@cached(300) | |
| 34 | +async def get_provider(request: Request, slug: str) -> dict[str, Any]: | |
| 35 | + return await detail_for_type(slug, ("provider",)) | |
added
src/aiatlas/api/routers/search.py
+111 −0
@@ -0,0 +1,111 @@ | ||
| 1 | +"""/search · /search/suggest — natural-language compiler + FTS/trigram (+ embeddings when the gateway is reachable).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +import re | |
| 6 | +import time | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +from fastapi import APIRouter, Depends, Query, Request | |
| 10 | + | |
| 11 | +from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, ApiError, cached, entity_summary, rate_limit | |
| 12 | +from aiatlas.db import connection, fetch_all, fetch_val | |
| 13 | +from aiatlas.services import cache | |
| 14 | +from aiatlas.services.embeddings import embed_query | |
| 15 | +from aiatlas.services.llm import gateway | |
| 16 | +from aiatlas.services.search import Query as SearchQuery | |
| 17 | +from aiatlas.services.search import compile_query, search_entities, suggest | |
| 18 | + | |
| 19 | +router = APIRouter(prefix="/api/v1/search", tags=["search"]) | |
| 20 | +TOTAL_CAP = 10_000 | |
| 21 | +EMBED_TIMEOUT_S = 2.0 | |
| 22 | +EMBED_BACKOFF_S = 300 | |
| 23 | +EMBED_DEGRADED_KEY = "search:embed-degraded" | |
| 24 | + | |
| 25 | + | |
| 26 | +def _count_sql(q: SearchQuery) -> tuple[str, dict[str, Any]]: | |
| 27 | + """Same filters as `search_entities`, without ranking — capped estimate.""" | |
| 28 | + where = ["e.merged_into is null"] | |
| 29 | + params: dict[str, Any] = {} | |
| 30 | + if q.entity_type: | |
| 31 | + where.append("e.entity_type = :etype") | |
| 32 | + params["etype"] = q.entity_type | |
| 33 | + if q.openness == "open": | |
| 34 | + where.append("(e.attributes->>'openness' in ('open-weights','open-source','open') or e.attributes->>'weights_availability' = 'open')") | |
| 35 | + elif q.openness == "proprietary": | |
| 36 | + where.append("e.attributes->>'openness' in ('proprietary','closed')") | |
| 37 | + if q.year_from: | |
| 38 | + where.append("left(e.attributes->>'release_date', 4) >= :yf") | |
| 39 | + params["yf"] = str(q.year_from) | |
| 40 | + if q.year_to: | |
| 41 | + where.append("left(e.attributes->>'release_date', 4) <= :yt") | |
| 42 | + params["yt"] = str(q.year_to) | |
| 43 | + if q.params_min: | |
| 44 | + where.append("(e.attributes->>'parameter_count')::double precision >= :pmin") | |
| 45 | + params["pmin"] = float(q.params_min) | |
| 46 | + if q.params_max: | |
| 47 | + where.append("(e.attributes->>'parameter_count')::double precision <= :pmax") | |
| 48 | + params["pmax"] = float(q.params_max) | |
| 49 | + if q.context_min: | |
| 50 | + where.append("(e.attributes->>'context_length')::double precision >= :cmin") | |
| 51 | + params["cmin"] = float(q.context_min) | |
| 52 | + for i, mod in enumerate(q.modalities): | |
| 53 | + where.append(f"e.attributes->'modalities' ? :mod{i}") | |
| 54 | + params[f"mod{i}"] = mod | |
| 55 | + if q.organization: | |
| 56 | + where.append("exists (select 1 from entities o where o.id = e.organization_id and o.canonical_name ilike :org)") | |
| 57 | + params["org"] = f"%{q.organization}%" | |
| 58 | + text = (q.filters.get("residual") or "").strip() or ("" if any([q.entity_type, q.openness, q.year_from, q.params_min, q.context_min, q.modalities]) else q.text) | |
| 59 | + if text: | |
| 60 | + params["q"] = text | |
| 61 | + params["qlike"] = f"%{text}%" | |
| 62 | + params["qprefix"] = " & ".join(f"{w}:*" for w in re.findall(r"\w+", text)[:8]) or text | |
| 63 | + where.append("(e.search @@ to_tsquery('simple', :qprefix) or e.canonical_name ilike :qlike or e.canonical_name % :q " | |
| 64 | + "or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))") | |
| 65 | + return f"select count(*) from (select 1 from entities e where {' and '.join(where)} limit {TOTAL_CAP}) t", params | |
| 66 | + | |
| 67 | + | |
| 68 | +@router.get("", dependencies=[Depends(rate_limit("search"))]) | |
| 69 | +@cached(120) | |
| 70 | +async def search(request: Request, q: str = Query("", max_length=300), type: str | None = Query(None, alias="type"), | |
| 71 | + limit: int = Query(30, ge=1, le=200), offset: int = Query(0, ge=0)) -> dict[str, Any]: | |
| 72 | + compiled = compile_query(q) | |
| 73 | + if type: | |
| 74 | + compiled.entity_type = type | |
| 75 | + if not q.strip() and not type: | |
| 76 | + raise ApiError(400, "q is required") | |
| 77 | + embedding = None | |
| 78 | + residual = (compiled.filters.get("residual") or "").strip() | |
| 79 | + if gateway.available and residual and not await cache.cache_get(EMBED_DEGRADED_KEY): | |
| 80 | + try: | |
| 81 | + embedding = await asyncio.wait_for(embed_query(residual), timeout=EMBED_TIMEOUT_S) | |
| 82 | + except Exception: # noqa: BLE001 — never wait on the LLM | |
| 83 | + embedding = None | |
| 84 | + if embedding is None: # back off: FTS-only for a while instead of paying the timeout on every query | |
| 85 | + await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time()}, EMBED_BACKOFF_S) | |
| 86 | + async with connection() as conn: | |
| 87 | + hits = await search_entities(conn, compiled, limit=limit, offset=offset, embedding=embedding) | |
| 88 | + ids = [h["id"] for h in hits] | |
| 89 | + full = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.id = any(cast(:ids as text[]))", ids=ids) if ids else [] | |
| 90 | + count_sql, count_params = _count_sql(compiled) | |
| 91 | + total = await fetch_val(conn, count_sql, **count_params) | |
| 92 | + by_id = {r["id"]: r for r in full} | |
| 93 | + items = [] | |
| 94 | + for h in hits: | |
| 95 | + row = by_id.get(h["id"]) | |
| 96 | + s = entity_summary(row) if row else None | |
| 97 | + if s: | |
| 98 | + s["rank"] = float(h["rank"] or 0) | |
| 99 | + items.append(s) | |
| 100 | + return {"query": {**compiled.as_dict(), "semantic": embedding is not None}, "items": items, "total": int(total or 0), "limit": limit, "offset": offset} | |
| 101 | + | |
| 102 | + | |
| 103 | +@router.get("/suggest", dependencies=[Depends(rate_limit("search"))]) | |
| 104 | +@cached(120) | |
| 105 | +async def search_suggest(request: Request, q: str = Query("", max_length=120), limit: int = Query(8, ge=1, le=20)) -> dict[str, Any]: | |
| 106 | + prefix = q.strip() | |
| 107 | + if len(prefix) < 1: | |
| 108 | + return {"items": []} | |
| 109 | + async with connection() as conn: | |
| 110 | + rows = await suggest(conn, prefix, limit=limit) | |
| 111 | + return {"items": [{"id": r["id"], "entity_type": r["entity_type"], "slug": r["slug"], "name": r["canonical_name"], "organization_name": r["organization_name"]} for r in rows]} | |
added
src/aiatlas/api/routers/sources.py
+34 −0
@@ -0,0 +1,34 @@ | ||
| 1 | +"""/sources — public transparency page: every source with tier, connectors and health (no credentials, no internal hosts).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import cached | |
| 9 | +from aiatlas.db import connection, fetch_all | |
| 10 | + | |
| 11 | +router = APIRouter(prefix="/api/v1/sources", tags=["sources"]) | |
| 12 | + | |
| 13 | + | |
| 14 | +@router.get("") | |
| 15 | +@cached(300) | |
| 16 | +async def list_sources(request: Request) -> dict[str, Any]: | |
| 17 | + async with connection() as conn: | |
| 18 | + rows = await fetch_all(conn, """ | |
| 19 | + select s.id, s.key, s.name, s.domain, s.tier, s.kind, s.category, s.base_url, s.robots_policy, s.enabled, s.priority, s.notes, | |
| 20 | + o.id as org_id, o.slug as org_slug, o.canonical_name as org_name, | |
| 21 | + (select count(*) from documents d where d.source_id = s.id) as documents, | |
| 22 | + (select count(*) from snapshots x join documents d on d.id = x.document_id where d.source_id = s.id) as snapshots, | |
| 23 | + (select max(d.last_fetched_at) from documents d where d.source_id = s.id) as last_crawled_at, | |
| 24 | + (select count(*) from claims c where c.source_id = s.id and c.status = 'current') as claims, | |
| 25 | + coalesce((select jsonb_agg(jsonb_build_object('name', c.name, 'label', c.label, 'health', c.health, 'enabled', c.enabled, | |
| 26 | + 'last_success_at', to_char(c.last_success_at at time zone 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.US"Z"'), | |
| 27 | + 'interval_seconds', c.interval_seconds, 'priority', c.priority, 'parser_version', c.parser_version) order by c.name) | |
| 28 | + from connectors c where c.source_id = s.id), '[]'::jsonb) as connectors | |
| 29 | + from sources s left join entities o on o.id = s.organization_id order by s.tier, s.priority, s.name""") | |
| 30 | + items = [{"key": r["key"], "name": r["name"], "domain": r["domain"], "tier": r["tier"], "kind": r["kind"], "category": r["category"], "base_url": r["base_url"], | |
| 31 | + "robots_policy": r["robots_policy"], "organization": {"id": r["org_id"], "slug": r["org_slug"], "name": r["org_name"]} if r["org_id"] else None, | |
| 32 | + "enabled": r["enabled"], "priority": r["priority"], "notes": r["notes"], "documents": int(r["documents"] or 0), "snapshots": int(r["snapshots"] or 0), | |
| 33 | + "claims": int(r["claims"] or 0), "last_crawled_at": r["last_crawled_at"], "connectors": r["connectors"] or []} for r in rows] | |
| 34 | + return {"items": items, "total": len(items), "tiers": {1: "official / primary", 2: "quality secondary", 3: "community", 4: "unverified"}} | |
added
src/aiatlas/api/routers/stats.py
+33 −0
@@ -0,0 +1,33 @@ | ||
| 1 | +"""/stats · /stats/history — always live from the database.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +from datetime import UTC, datetime | |
| 6 | +from typing import Any | |
| 7 | + | |
| 8 | +from fastapi import APIRouter, Query, Request | |
| 9 | + | |
| 10 | +from aiatlas.api.common import cached | |
| 11 | +from aiatlas.db import connection | |
| 12 | +from aiatlas.sdk.archive import archive_size | |
| 13 | +from aiatlas.services.stats import history, live_counts | |
| 14 | + | |
| 15 | +router = APIRouter(prefix="/api/v1/stats", tags=["stats"]) | |
| 16 | + | |
| 17 | + | |
| 18 | +@router.get("") | |
| 19 | +@cached(60) | |
| 20 | +async def stats(request: Request) -> dict[str, Any]: | |
| 21 | + async with connection() as conn: | |
| 22 | + counts = await live_counts(conn) | |
| 23 | + counts["archive"] = await asyncio.to_thread(archive_size) | |
| 24 | + counts["computed_at"] = datetime.now(UTC) | |
| 25 | + return counts | |
| 26 | + | |
| 27 | + | |
| 28 | +@router.get("/history") | |
| 29 | +@cached(300) | |
| 30 | +async def stats_history(request: Request, days: int = Query(90, ge=1, le=730)) -> dict[str, Any]: | |
| 31 | + async with connection() as conn: | |
| 32 | + rows = await history(conn, days) | |
| 33 | + return {"items": [{"day": r["day"].date().isoformat() if hasattr(r["day"], "date") else str(r["day"]), "counts": r["counts"]} for r in rows]} | |
added
src/aiatlas/api/routers/timeline.py
+43 −0
@@ -0,0 +1,43 @@ | ||
| 1 | +"""/timeline — events grouped by month (global, or one entity + what it develops).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from aiatlas.api.common import COMPANY_TYPES, EVENT_COLS, EVENT_FROM, cached, change_event, csv, resolve_entity | |
| 9 | +from aiatlas.db import connection, fetch_all | |
| 10 | + | |
| 11 | +router = APIRouter(prefix="/api/v1/timeline", tags=["timeline"]) | |
| 12 | + | |
| 13 | + | |
| 14 | +@router.get("") | |
| 15 | +@cached(120) | |
| 16 | +async def timeline(request: Request, entity: str | None = None, year: int | None = Query(None, ge=1950, le=2100), category: str | None = None, | |
| 17 | + importance_min: int | None = Query(None, ge=0, le=3), limit: int = Query(200, ge=1, le=1000)) -> dict[str, Any]: | |
| 18 | + where = ["ev.event_type <> 'DOCUMENT_CHANGED'"] | |
| 19 | + params: dict[str, Any] = {"lim": limit} | |
| 20 | + async with connection() as conn: | |
| 21 | + if entity: | |
| 22 | + row = await resolve_entity(conn, entity) | |
| 23 | + if row["entity_type"] in COMPANY_TYPES: | |
| 24 | + where.append("(ev.entity_id = :eid or ev.entity_id in (select id from entities where organization_id = :eid union " | |
| 25 | + "select object_id from relations where subject_id = :eid and predicate in ('develops','owns','operates','published') and valid_to is null))") | |
| 26 | + else: | |
| 27 | + where.append("ev.entity_id = :eid") | |
| 28 | + params["eid"] = row["id"] | |
| 29 | + if year: | |
| 30 | + where.append("extract(year from coalesce(ev.effective_at, ev.observed_at)) = :year") | |
| 31 | + params["year"] = year | |
| 32 | + if category: | |
| 33 | + where.append("ev.category = any(cast(:cats as text[]))") | |
| 34 | + params["cats"] = csv(category) | |
| 35 | + if importance_min is not None: | |
| 36 | + where.append("ev.importance >= :imp") | |
| 37 | + params["imp"] = importance_min | |
| 38 | + rows = await fetch_all(conn, f"select to_char(coalesce(ev.effective_at, ev.observed_at) at time zone 'UTC', 'YYYY-MM') as month, {EVENT_COLS} from {EVENT_FROM} " | |
| 39 | + f"where {' and '.join(where)} order by coalesce(ev.effective_at, ev.observed_at) desc, ev.id desc limit :lim", **params) | |
| 40 | + groups: dict[str, list[dict[str, Any]]] = {} | |
| 41 | + for r in rows: | |
| 42 | + groups.setdefault(r["month"], []).append(change_event(r)) | |
| 43 | + return {"items": [{"month": m, "count": len(evs), "events": evs} for m, evs in groups.items()], "total": len(rows)} | |
added
src/aiatlas/services/hardware_fit.py
+68 −0
@@ -0,0 +1,68 @@ | ||
| 1 | +"""ESTIMATED memory footprint of a model on a piece of hardware. Transparent formula, always labelled as an estimate: | |
| 2 | + | |
| 3 | + weights = parameter_count × bytes_per_param × 1.15 (runtime overhead: activations, buffers, fragmentation) | |
| 4 | + kv_cache = 0.5 GB per 8 192 tokens of context (coarse, architecture-agnostic) | |
| 5 | + fits = estimated_memory_gb ≤ hardware memory − 2 GB (OS / framework headroom) | |
| 6 | +""" | |
| 7 | +from __future__ import annotations | |
| 8 | + | |
| 9 | +from typing import Any | |
| 10 | + | |
| 11 | +BYTES_PER_PARAM = {"4bit": 0.5, "8bit": 1.0, "fp16": 2.0} | |
| 12 | +OVERHEAD = 1.15 | |
| 13 | +KV_GB_PER_8K = 0.5 | |
| 14 | +RESERVED_GB = 2.0 | |
| 15 | +ASSUMPTIONS = [ | |
| 16 | + "Estimated, not measured: weights = parameters × bytes/param × 1.15 runtime overhead.", | |
| 17 | + "bytes/param: 4bit = 0.5, 8bit = 1.0, fp16 = 2.0 (uniform quantization, no per-layer exceptions).", | |
| 18 | + "KV cache approximated at 0.5 GB per 8 192 tokens of context, independent of architecture (GQA/MLA models need less).", | |
| 19 | + "A model 'fits' when the estimate is at most the device memory minus 2 GB reserved for the OS and framework.", | |
| 20 | + "Mixture-of-experts models are estimated on total parameters (all experts must be resident); active parameters are ignored.", | |
| 21 | + "Device memory uses the largest configuration when several are listed (e.g. Apple silicon tiers).", | |
| 22 | +] | |
| 23 | + | |
| 24 | + | |
| 25 | +def estimate_memory_gb(parameter_count: float, quant: str = "4bit", context: int = 8192) -> float: | |
| 26 | + bpp = BYTES_PER_PARAM.get(quant, BYTES_PER_PARAM["4bit"]) | |
| 27 | + weights = parameter_count * bpp * OVERHEAD / 1e9 | |
| 28 | + kv = KV_GB_PER_8K * max(0.0, float(context)) / 8192.0 | |
| 29 | + return round(weights + kv, 2) | |
| 30 | + | |
| 31 | + | |
| 32 | +def hardware_memory_gb(attrs: dict[str, Any] | None) -> float | None: | |
| 33 | + """`memory_gb` may be a number or a list of configurations — use the largest.""" | |
| 34 | + v = (attrs or {}).get("memory_gb") | |
| 35 | + if isinstance(v, list): | |
| 36 | + nums = [float(x) for x in v if isinstance(x, (int, float)) and not isinstance(x, bool)] | |
| 37 | + return max(nums) if nums else None | |
| 38 | + if isinstance(v, (int, float)) and not isinstance(v, bool): | |
| 39 | + return float(v) | |
| 40 | + if isinstance(v, str): | |
| 41 | + try: | |
| 42 | + return float(v) | |
| 43 | + except ValueError: | |
| 44 | + return None | |
| 45 | + return None | |
| 46 | + | |
| 47 | + | |
| 48 | +def parameter_count(attrs: dict[str, Any] | None) -> float | None: | |
| 49 | + v = (attrs or {}).get("parameter_count") | |
| 50 | + if isinstance(v, (int, float)) and not isinstance(v, bool) and v > 0: | |
| 51 | + return float(v) | |
| 52 | + if isinstance(v, str): | |
| 53 | + try: | |
| 54 | + f = float(v) | |
| 55 | + return f if f > 0 else None | |
| 56 | + except ValueError: | |
| 57 | + return None | |
| 58 | + return None | |
| 59 | + | |
| 60 | + | |
| 61 | +def fit(parameter_count_: float, memory_gb: float, quant: str = "4bit", context: int = 8192) -> dict[str, Any]: | |
| 62 | + est = estimate_memory_gb(parameter_count_, quant, context) | |
| 63 | + headroom = round(memory_gb - RESERVED_GB - est, 2) | |
| 64 | + return {"quantization": quant, "estimated_memory_gb": est, "fits": headroom >= 0, "headroom_gb": headroom, | |
| 65 | + "note": f"estimated: {parameter_count_ / 1e9:.1f}B params × {BYTES_PER_PARAM.get(quant, 0.5)} B × 1.15 + KV cache for {context} tokens"} | |
| 66 | + | |
| 67 | + | |
| 68 | +__all__ = ["ASSUMPTIONS", "BYTES_PER_PARAM", "estimate_memory_gb", "fit", "hardware_memory_gb", "parameter_count"] | |
added
src/aiatlas/services/merge.py
+77 −0
@@ -0,0 +1,77 @@ | ||
| 1 | +"""Entity merging (curation): fold a duplicate `source` into `target`. Nothing is deleted — the source row stays with | |
| 2 | +`status='merged'` and `merged_into=target` so old slugs and ids keep resolving; every dependent row is re-pointed.""" | |
| 3 | +from __future__ import annotations | |
| 4 | + | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 8 | + | |
| 9 | +from aiatlas.db import execute, fetch_one, fetch_val | |
| 10 | +from aiatlas.ids import new_id, normalize_alias | |
| 11 | + | |
| 12 | + | |
| 13 | +async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str) -> dict[str, Any]: | |
| 14 | + if source_id == target_id: | |
| 15 | + raise ValueError("source and target are the same entity") | |
| 16 | + src = await fetch_one(conn, "select id, entity_type, canonical_name, slug, merged_into, attributes, provenance from entities where id = :id", id=source_id) | |
| 17 | + dst = await fetch_one(conn, "select id, entity_type, merged_into from entities where id = :id", id=target_id) | |
| 18 | + if not src or not dst: | |
| 19 | + raise LookupError("source or target entity not found") | |
| 20 | + if dst["merged_into"]: | |
| 21 | + raise ValueError("target is itself merged; merge into its survivor instead") | |
| 22 | + if src["merged_into"]: | |
| 23 | + raise ValueError("source is already merged") | |
| 24 | + if src["entity_type"] != dst["entity_type"]: | |
| 25 | + raise ValueError(f"cannot merge a {src['entity_type']} into a {dst['entity_type']}") | |
| 26 | + | |
| 27 | + moved: dict[str, int] = {} | |
| 28 | + | |
| 29 | + async def count_update(label: str, sql: str) -> None: | |
| 30 | + n = await fetch_val(conn, f"with u as ({sql} returning 1) select count(*) from u", s=source_id, t=target_id) | |
| 31 | + moved[label] = int(n or 0) | |
| 32 | + | |
| 33 | + await count_update("aliases", "insert into entity_aliases (entity_id, alias, alias_norm, kind, snapshot_id) " | |
| 34 | + "select :t, alias, alias_norm, kind, snapshot_id from entity_aliases where entity_id = :s on conflict (entity_id, alias_norm) do nothing") | |
| 35 | + await execute(conn, "insert into entity_aliases (entity_id, alias, alias_norm, kind) values (:t, :a, :n, 'former_name') " | |
| 36 | + "on conflict (entity_id, alias_norm) do update set kind = 'former_name'", | |
| 37 | + t=target_id, a=src["canonical_name"], n=normalize_alias(src["canonical_name"])) | |
| 38 | + await execute(conn, "delete from entity_aliases where entity_id = :s", s=source_id) | |
| 39 | + await count_update("identifiers", "update entity_identifiers x set entity_id = :t where entity_id = :s " | |
| 40 | + "and not exists (select 1 from entity_identifiers y where y.scheme = x.scheme and y.value = x.value and y.entity_id = :t)") | |
| 41 | + await execute(conn, "delete from entity_identifiers where entity_id = :s", s=source_id) | |
| 42 | + await count_update("claims", "update claims set entity_id = :t where entity_id = :s") | |
| 43 | + # relations: drop those that would duplicate a live edge on the target, then re-point | |
| 44 | + await execute(conn, """update relations r set valid_to = now() where r.valid_to is null and (r.subject_id = :s or r.object_id = :s) and exists ( | |
| 45 | + select 1 from relations x where x.valid_to is null and x.predicate = r.predicate | |
| 46 | + and x.subject_id = case when r.subject_id = :s then :t else r.subject_id end | |
| 47 | + and x.object_id = case when r.object_id = :s then :t else r.object_id end)""", s=source_id, t=target_id) | |
| 48 | + await count_update("relations", "update relations set subject_id = case when subject_id = :s then :t else subject_id end, " | |
| 49 | + "object_id = case when object_id = :s then :t else object_id end where subject_id = :s or object_id = :s") | |
| 50 | + await execute(conn, "update relations set valid_to = now() where subject_id = object_id and valid_to is null and subject_id = :t", t=target_id) | |
| 51 | + await count_update("events", "update change_events set entity_id = :t where entity_id = :s") | |
| 52 | + await execute(conn, "update prices set valid_to = now() where valid_to is null and model_id = :s and exists (select 1 from prices q where q.valid_to is null " | |
| 53 | + "and q.model_id = :t and q.provider_id = prices.provider_id and coalesce(q.provider_model_id,'') = coalesce(prices.provider_model_id,''))", | |
| 54 | + s=source_id, t=target_id) | |
| 55 | + await count_update("prices", "update prices set model_id = case when model_id = :s then :t else model_id end, " | |
| 56 | + "provider_id = case when provider_id = :s then :t else provider_id end where model_id = :s or provider_id = :s") | |
| 57 | + await count_update("results", "update benchmark_results set model_id = case when model_id = :s then :t else model_id end, " | |
| 58 | + "benchmark_id = case when benchmark_id = :s then :t else benchmark_id end where model_id = :s or benchmark_id = :s") | |
| 59 | + await count_update("documents", "update documents set entity_id = :t where entity_id = :s") | |
| 60 | + await count_update("children", "update entities set organization_id = :t where organization_id = :s") | |
| 61 | + await count_update("llm_jobs", "update llm_jobs set entity_id = :t where entity_id = :s") | |
| 62 | + # attributes the target lacks are inherited (with their provenance); target values always win | |
| 63 | + await execute(conn, """update entities t set attributes = coalesce(s.attributes, '{}'::jsonb) || t.attributes, | |
| 64 | + provenance = coalesce(s.provenance, '{}'::jsonb) || t.provenance, last_seen_at = greatest(t.last_seen_at, s.last_seen_at), | |
| 65 | + first_seen_at = least(t.first_seen_at, s.first_seen_at), updated_at = now() | |
| 66 | + from entities s where t.id = :t and s.id = :s""", s=source_id, t=target_id) | |
| 67 | + await execute(conn, "update entities set merged_into = :t, status = 'merged', updated_at = now() where id = :s", s=source_id, t=target_id) | |
| 68 | + await execute(conn, "update entities set merged_into = :t where merged_into = :s", s=source_id, t=target_id) | |
| 69 | + await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, connector_name, dedupe_key, meta) | |
| 70 | + values (:id, :t, 'ENTITY_MERGED', 'source', :sum, 1, 'curation', :dk, cast(:m as jsonb)) | |
| 71 | + on conflict (dedupe_key) do nothing""", | |
| 72 | + id=new_id("change_event"), t=target_id, sum=f"Merged duplicate '{src['canonical_name']}' ({src['slug']})", dk=f"merge:{source_id}:{target_id}", | |
| 73 | + m=f'{{"source_id": "{source_id}", "source_slug": "{src["slug"]}"}}') | |
| 74 | + return {"source_id": source_id, "target_id": target_id, "moved": moved} | |
| 75 | + | |
| 76 | + | |
| 77 | +__all__ = ["merge_entities"] | |
modified
src/aiatlas/services/scheduler.py
+21 −4
@@ -39,16 +39,33 @@ async def run_connector(name: str, *, force: bool = False) -> None: | ||
| 39 | 39 | _running.discard(name) |
| 40 | 40 | |
| 41 | 41 | |
| 42 | +async def _pop_run_now() -> list[str]: | |
| 43 | + """Admin `POST /admin/connectors/{name}/run` sets `aia:run-now:<name>`; consume (delete) every such key.""" | |
| 44 | + names: list[str] = [] | |
| 45 | + try: | |
| 46 | + r = cache.redis() | |
| 47 | + async for key in r.scan_iter(match="aia:run-now:*", count=100): | |
| 48 | + if await r.delete(key): | |
| 49 | + names.append(key.decode().rsplit(":", 1)[-1]) | |
| 50 | + except Exception as exc: # noqa: BLE001 | |
| 51 | + log.debug("run-now scan skipped", extra={"error": str(exc)}) | |
| 52 | + return names | |
| 53 | + | |
| 54 | + | |
| 42 | 55 | async def tick() -> None: |
| 56 | + known = registry() | |
| 57 | + forced = [n for n in await _pop_run_now() if n in known and n not in _running] | |
| 58 | + if forced: | |
| 59 | + log.info("run-now connectors", extra={"connectors": forced}) | |
| 60 | + await asyncio.gather(*(run_connector(n, force=True) for n in forced[:3])) | |
| 43 | 61 | async with transaction() as conn: |
| 44 | 62 | due = await fetch_all(conn, """select name from connectors where enabled and (next_run_at is null or next_run_at <= now()) |
| 45 | 63 | and (circuit_open_until is null or circuit_open_until <= now()) order by priority, coalesce(next_run_at, 'epoch') limit 6""") |
| 46 | − known = registry() | |
| 47 | − names = [r["name"] for r in due if r["name"] in known and r["name"] not in _running] | |
| 64 | + names = [r["name"] for r in due if r["name"] in known and r["name"] not in _running and r["name"] not in forced] | |
| 48 | 65 | if names: |
| 49 | 66 | log.info("due connectors", extra={"connectors": names}) |
| 50 | 67 | await asyncio.gather(*(run_connector(n) for n in names[:3])) |
| 51 | − await cache.heartbeat("scheduler", {"at": datetime.now(UTC).isoformat(), "running": sorted(_running), "due": names}) | |
| 68 | + await cache.heartbeat("scheduler", {"at": datetime.now(UTC).isoformat(), "running": sorted(_running), "due": names, "forced": forced}) | |
| 52 | 69 | |
| 53 | 70 | |
| 54 | 71 | async def hourly() -> None: |
@@ -116,4 +133,4 @@ async def main(*, with_worker: bool = True) -> None: | ||
| 116 | 133 | await cache.close() |
| 117 | 134 | |
| 118 | 135 | |
| 119 | −__all__ = ["main", "tick", "run_connector", "hourly"] | |
| 136 | +__all__ = ["hourly", "main", "run_connector", "tick"] | |
added
tests/test_api.py
+258 −0
@@ -0,0 +1,258 @@ | ||
| 1 | +"""API contract smoke tests against the live local database (docs/API.md). Read-only except for the `/views` beacon. | |
| 2 | +Counts are never asserted as fixed numbers: connectors add data while the suite runs.""" | |
| 3 | +from __future__ import annotations | |
| 4 | + | |
| 5 | +from collections.abc import AsyncIterator | |
| 6 | + | |
| 7 | +import pytest | |
| 8 | +from httpx import ASGITransport, AsyncClient | |
| 9 | + | |
| 10 | +from aiatlas import db | |
| 11 | +from aiatlas.api.main import app | |
| 12 | +from aiatlas.config import settings | |
| 13 | +from aiatlas.services import cache | |
| 14 | + | |
| 15 | +ADMIN = {"x-aia-admin-token": settings.admin_token or "dev-admin-token"} | |
| 16 | +MODEL = "claude-opus-5" | |
| 17 | +MODEL_B = "claude-sonnet-5" | |
| 18 | + | |
| 19 | + | |
| 20 | +@pytest.fixture | |
| 21 | +async def client() -> AsyncIterator[AsyncClient]: | |
| 22 | + """pytest-asyncio runs each test in its own loop: pools must not outlive the loop that created them.""" | |
| 23 | + await cache.cache_invalidate() | |
| 24 | + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as c: | |
| 25 | + yield c | |
| 26 | + await cache.close() | |
| 27 | + await db.dispose() | |
| 28 | + | |
| 29 | + | |
| 30 | +async def test_health(client: AsyncClient) -> None: | |
| 31 | + for path in ("/health", "/api/v1/health"): | |
| 32 | + r = await client.get(path) | |
| 33 | + assert r.status_code == 200 | |
| 34 | + body = r.json() | |
| 35 | + assert body["status"] in ("ok", "degraded") and body["db"] is True | |
| 36 | + assert {"version", "redis", "llm", "time"} <= set(body) | |
| 37 | + | |
| 38 | + | |
| 39 | +async def test_stats_keys_and_live(client: AsyncClient) -> None: | |
| 40 | + r = await client.get("/api/v1/stats") | |
| 41 | + assert r.status_code == 200 | |
| 42 | + s = r.json() | |
| 43 | + for key in ("entities", "entities_total", "sources", "connectors", "documents", "snapshots", "claims", "claims_current", "relations", "change_events", | |
| 44 | + "change_events_24h", "benchmark_results", "prices_current", "prices_total", "review_pending", "llm_jobs", "archive", "computed_at"): | |
| 45 | + assert key in s, key | |
| 46 | + assert s["entities_total"] == sum(s["entities"].values()) > 0 | |
| 47 | + assert s["entities"].get("model", 0) > 0 | |
| 48 | + assert {"raw_bytes", "raw_files", "text_bytes", "text_files"} <= set(s["archive"]) | |
| 49 | + hist = await client.get("/api/v1/stats/history?days=30") | |
| 50 | + assert hist.status_code == 200 and "items" in hist.json() | |
| 51 | + | |
| 52 | + | |
| 53 | +async def test_search_returns_claude_models(client: AsyncClient) -> None: | |
| 54 | + r = await client.get("/api/v1/search", params={"q": "claude", "limit": 10}) | |
| 55 | + assert r.status_code == 200 | |
| 56 | + body = r.json() | |
| 57 | + assert body["total"] > 0 and body["items"] | |
| 58 | + assert body["query"]["text"] == "claude" | |
| 59 | + assert all("rank" in it for it in body["items"]) | |
| 60 | + assert any(it["entity_type"] == "model" and it["slug"].startswith("claude") for it in body["items"]) | |
| 61 | + sug = await client.get("/api/v1/search/suggest", params={"q": "cla"}) | |
| 62 | + assert sug.status_code == 200 and sug.json()["items"] | |
| 63 | + assert {"id", "entity_type", "slug", "name", "organization_name"} <= set(sug.json()["items"][0]) | |
| 64 | + assert (await client.get("/api/v1/search", params={"q": ""})).status_code == 400 | |
| 65 | + | |
| 66 | + | |
| 67 | +async def test_models_filters(client: AsyncClient) -> None: | |
| 68 | + r = await client.get("/api/v1/models", params={"openness": "proprietary", "limit": 200}) | |
| 69 | + assert r.status_code == 200 | |
| 70 | + body = r.json() | |
| 71 | + assert {"items", "total", "limit", "offset"} <= set(body) and body["items"] | |
| 72 | + assert all(it["attributes"].get("openness") == "proprietary" for it in body["items"]) | |
| 73 | + r = await client.get("/api/v1/models", params={"min_context": 500000, "limit": 200}) | |
| 74 | + assert r.status_code == 200 | |
| 75 | + items = r.json()["items"] | |
| 76 | + assert items and all(int(it["attributes"]["context_length"]) >= 500000 for it in items) | |
| 77 | + r = await client.get("/api/v1/models", params={"facets": 1, "limit": 1}) | |
| 78 | + assert r.status_code == 200 | |
| 79 | + facets = r.json()["facets"] | |
| 80 | + assert {"organizations", "openness", "modalities", "families", "years", "licenses", "status"} <= set(facets) | |
| 81 | + assert any(o["slug"] == "anthropic" for o in facets["organizations"]) | |
| 82 | + assert (await client.get("/api/v1/models", params={"sort": "nope"})).status_code == 400 | |
| 83 | + assert (await client.get("/api/v1/models", params={"limit": 999})).status_code == 422 | |
| 84 | + | |
| 85 | + | |
| 86 | +async def test_entity_detail_blocks(client: AsyncClient) -> None: | |
| 87 | + r = await client.get(f"/api/v1/entities/{MODEL}") | |
| 88 | + assert r.status_code == 200 | |
| 89 | + d = r.json() | |
| 90 | + assert d["slug"] == MODEL and d["entity_type"] == "model" | |
| 91 | + for block in ("attributes", "provenance", "aliases", "identifiers", "relations", "sources", "timeline", "prices", "price_history", "results", "lineage", | |
| 92 | + "providers", "quality", "counts"): | |
| 93 | + assert block in d, block | |
| 94 | + assert d["provenance"] and all({"tier", "confidence", "extractor", "observed_at"} <= set(v) for v in d["provenance"].values()) | |
| 95 | + assert d["prices"] and d["prices"][0]["provider"]["slug"] and d["prices"][0]["input_per_mtok"] is not None | |
| 96 | + assert d["timeline"] and any(e["event_type"] == "NEW_MODEL" for e in d["timeline"]) | |
| 97 | + assert d["sources"] and {"url", "tier", "doc_type", "snapshots"} <= set(d["sources"][0]) | |
| 98 | + assert d["organization"]["slug"] == "anthropic" | |
| 99 | + assert {"ancestors", "descendants", "quantizations"} == set(d["lineage"]) | |
| 100 | + # type-scoped aliases | |
| 101 | + assert (await client.get(f"/api/v1/models/{MODEL}")).status_code == 200 | |
| 102 | + assert (await client.get(f"/api/v1/companies/{MODEL}")).status_code == 404 | |
| 103 | + assert (await client.get("/api/v1/entities/definitely-not-an-entity")).status_code == 404 | |
| 104 | + # resolution by id | |
| 105 | + by_id = await client.get(f"/api/v1/entities/{d['id']}") | |
| 106 | + assert by_id.status_code == 200 and by_id.json()["slug"] == MODEL | |
| 107 | + | |
| 108 | + | |
| 109 | +async def test_company_detail(client: AsyncClient) -> None: | |
| 110 | + r = await client.get("/api/v1/companies/anthropic") | |
| 111 | + assert r.status_code == 200 | |
| 112 | + d = r.json() | |
| 113 | + assert d["entity_type"] in ("company", "organization", "lab", "university") | |
| 114 | + assert d["models"]["total"] > 0 and d["models"]["items"][0]["entity_type"] == "model" | |
| 115 | + assert d["timeline"] # includes events of the models it develops | |
| 116 | + | |
| 117 | + | |
| 118 | +async def test_entity_subresources(client: AsyncClient) -> None: | |
| 119 | + tl = await client.get(f"/api/v1/entities/{MODEL}/timeline", params={"limit": 5}) | |
| 120 | + assert tl.status_code == 200 and "items" in tl.json() | |
| 121 | + hist = await client.get(f"/api/v1/entities/{MODEL}/history", params={"property": "context_length"}) | |
| 122 | + assert hist.status_code == 200 and hist.json()["items"] and hist.json()["items"][0]["property"] == "context_length" | |
| 123 | + graph = await client.get(f"/api/v1/entities/{MODEL}/graph", params={"depth": 1}) | |
| 124 | + assert graph.status_code == 200 and graph.json()["nodes"] and "edges" in graph.json() | |
| 125 | + src = await client.get(f"/api/v1/entities/{MODEL}/sources") | |
| 126 | + assert src.status_code == 200 and src.json()["items"] | |
| 127 | + rel = await client.get(f"/api/v1/entities/{MODEL}/related", params={"limit": 5}) | |
| 128 | + assert rel.status_code == 200 and len(rel.json()["items"]) <= 5 | |
| 129 | + | |
| 130 | + | |
| 131 | +async def test_asof(client: AsyncClient) -> None: | |
| 132 | + r = await client.get(f"/api/v1/entities/{MODEL}/asof", params={"date": "2020-01-01"}) | |
| 133 | + assert r.status_code == 200 | |
| 134 | + body = r.json() | |
| 135 | + assert body["existed"] is False and body["attributes"] == {} and body["claims"] == [] | |
| 136 | + now = await client.get(f"/api/v1/entities/{MODEL}/asof", params={"date": "2999-12-31"}) | |
| 137 | + assert now.status_code == 200 and now.json()["existed"] is True and now.json()["attributes"] | |
| 138 | + assert (await client.get(f"/api/v1/entities/{MODEL}/asof", params={"date": "not-a-date"})).status_code == 400 | |
| 139 | + | |
| 140 | + | |
| 141 | +async def test_changes_pagination(client: AsyncClient) -> None: | |
| 142 | + first = await client.get("/api/v1/changes", params={"limit": 3}) | |
| 143 | + assert first.status_code == 200 | |
| 144 | + body = first.json() | |
| 145 | + assert {"items", "total", "limit", "offset"} <= set(body) and len(body["items"]) == 3 and body["total"] >= 3 | |
| 146 | + assert all(e["event_type"] != "DOCUMENT_CHANGED" for e in body["items"]) | |
| 147 | + ev = body["items"][0] | |
| 148 | + assert {"id", "event_type", "category", "summary", "importance", "observed_at", "entity", "meta"} <= set(ev) | |
| 149 | + cursor = body["next_before"] | |
| 150 | + assert cursor == body["items"][-1]["observed_at"] | |
| 151 | + second = await client.get("/api/v1/changes", params={"limit": 3, "before": cursor}) | |
| 152 | + assert second.status_code == 200 | |
| 153 | + ids1 = {e["id"] for e in body["items"]} | |
| 154 | + assert all(e["id"] not in ids1 for e in second.json()["items"]) | |
| 155 | + assert all(e["observed_at"] < cursor for e in second.json()["items"]) | |
| 156 | + daily = await client.get("/api/v1/changes/daily") | |
| 157 | + assert daily.status_code == 200 and {"date", "counts", "sections", "new_models"} <= set(daily.json()) | |
| 158 | + cats = await client.get("/api/v1/changes/categories", params={"days": 30}) | |
| 159 | + assert cats.status_code == 200 and "items" in cats.json() | |
| 160 | + | |
| 161 | + | |
| 162 | +async def test_compare(client: AsyncClient) -> None: | |
| 163 | + r = await client.get("/api/v1/compare", params={"ids": f"{MODEL},{MODEL_B}"}) | |
| 164 | + assert r.status_code == 200 | |
| 165 | + body = r.json() | |
| 166 | + assert body["entity_type"] == "model" and len(body["items"]) == 2 | |
| 167 | + keys = {d["key"] for d in body["dimensions"]} | |
| 168 | + assert {"parameter_count", "context_length", "openness", "release_date", "best_input_per_mtok"} <= keys | |
| 169 | + for item in body["items"]: | |
| 170 | + assert {"entity", "values", "provenance", "prices", "results"} <= set(item) | |
| 171 | + assert item["values"]["context_length"] is not None | |
| 172 | + assert (await client.get("/api/v1/compare", params={"ids": MODEL})).status_code == 400 | |
| 173 | + assert (await client.get("/api/v1/compare", params={"ids": f"{MODEL},anthropic"})).status_code == 400 | |
| 174 | + | |
| 175 | + | |
| 176 | +async def test_listings_and_misc(client: AsyncClient) -> None: | |
| 177 | + for path in ("/api/v1/companies?facets=1", "/api/v1/papers", "/api/v1/providers", "/api/v1/prices", "/api/v1/prices/index?days=14", "/api/v1/benchmarks", | |
| 178 | + "/api/v1/hardware", "/api/v1/hardware/fit?memory_gb=24", "/api/v1/explore/types", "/api/v1/explore/hardware", "/api/v1/timeline?limit=10", | |
| 179 | + "/api/v1/diff?a=2026-01-01&b=2026-12-31", "/api/v1/sources", "/api/v1/methodology", "/api/v1/trending", "/api/v1/sitemap?limit=5", | |
| 180 | + f"/api/v1/prices/history?model={MODEL}"): | |
| 181 | + r = await client.get(path) | |
| 182 | + assert r.status_code == 200, (path, r.text[:200]) | |
| 183 | + prices = (await client.get("/api/v1/prices", params={"model": MODEL})).json() | |
| 184 | + assert prices["items"] and prices["items"][0]["model"]["slug"] == MODEL | |
| 185 | + fit = (await client.get("/api/v1/hardware/fit", params={"memory_gb": 24})).json() | |
| 186 | + assert fit["estimated"] is True and fit["assumptions"] and "items" in fit | |
| 187 | + assert (await client.get("/api/v1/hardware/fit", params={"memory_gb": 24, "quant": "2bit"})).status_code == 400 | |
| 188 | + assert (await client.get("/api/v1/explore/nonsense")).status_code == 404 | |
| 189 | + assert (await client.get("/api/v1/api-keys/me")).status_code == 401 | |
| 190 | + view = await client.post("/api/v1/views", json={"path": "/models/claude-opus-5"}) | |
| 191 | + assert view.status_code == 200 and view.json() == {"ok": True} | |
| 192 | + assert (await client.post("/api/v1/views", json={"path": "/x"})).status_code == 429 # 1 req/s/IP | |
| 193 | + | |
| 194 | + | |
| 195 | +async def test_cache_roundtrip_is_stable(client: AsyncClient) -> None: | |
| 196 | + a = (await client.get("/api/v1/models", params={"limit": 3})).json() | |
| 197 | + b = (await client.get("/api/v1/models", params={"limit": 3})).json() | |
| 198 | + assert a == b | |
| 199 | + | |
| 200 | + | |
| 201 | +async def test_admin_auth(client: AsyncClient) -> None: | |
| 202 | + assert (await client.get("/api/v1/admin/overview")).status_code == 401 | |
| 203 | + assert (await client.get("/api/v1/admin/overview", headers={"x-aia-admin-token": "wrong"})).status_code == 401 | |
| 204 | + r = await client.get("/api/v1/admin/overview", headers=ADMIN) | |
| 205 | + assert r.status_code == 200 | |
| 206 | + body = r.json() | |
| 207 | + assert {"stats", "queue", "heartbeats", "connectors", "review_pending", "recent_errors", "llm", "archive"} <= set(body) | |
| 208 | + saved = settings.admin_token | |
| 209 | + settings.admin_token = "" | |
| 210 | + try: | |
| 211 | + assert (await client.get("/api/v1/admin/overview", headers=ADMIN)).status_code == 503 | |
| 212 | + finally: | |
| 213 | + settings.admin_token = saved | |
| 214 | + | |
| 215 | + | |
| 216 | +async def test_merge_entities_rolled_back(client: AsyncClient) -> None: | |
| 217 | + """Exercise the curation SQL end-to-end on live rows, then roll back — the database is left untouched.""" | |
| 218 | + from aiatlas.db import engine, fetch_one, fetch_val | |
| 219 | + from aiatlas.services.merge import merge_entities | |
| 220 | + | |
| 221 | + async with engine().connect() as conn: | |
| 222 | + trans = await conn.begin() | |
| 223 | + try: | |
| 224 | + src = await fetch_one(conn, "select id from entities where slug = :s", s=MODEL_B) | |
| 225 | + dst = await fetch_one(conn, "select id from entities where slug = :s", s=MODEL) | |
| 226 | + assert src and dst | |
| 227 | + res = await merge_entities(conn, src["id"], dst["id"]) | |
| 228 | + assert res["target_id"] == dst["id"] and res["moved"]["claims"] > 0 | |
| 229 | + merged = await fetch_one(conn, "select status, merged_into from entities where id = :id", id=src["id"]) | |
| 230 | + assert merged["status"] == "merged" and merged["merged_into"] == dst["id"] | |
| 231 | + assert await fetch_val(conn, "select count(*) from claims where entity_id = :id", id=src["id"]) == 0 | |
| 232 | + assert await fetch_val(conn, "select count(*) from entity_aliases where entity_id = :t and kind = 'former_name'", t=dst["id"]) >= 1 | |
| 233 | + with pytest.raises(ValueError): | |
| 234 | + await merge_entities(conn, src["id"], dst["id"]) # already merged | |
| 235 | + finally: | |
| 236 | + await trans.rollback() | |
| 237 | + async with engine().connect() as conn: | |
| 238 | + assert (await fetch_one(conn, "select status from entities where slug = :s", s=MODEL_B))["status"] != "merged" | |
| 239 | + | |
| 240 | + | |
| 241 | +async def test_admin_read_routes(client: AsyncClient) -> None: | |
| 242 | + review = await client.get("/api/v1/admin/review", headers=ADMIN) | |
| 243 | + assert review.status_code == 200 and {"items", "total", "by_kind"} <= set(review.json()) | |
| 244 | + for path in ("/api/v1/admin/connectors", "/api/v1/admin/runs?limit=2", "/api/v1/admin/errors?limit=2", "/api/v1/admin/documents?limit=2", "/api/v1/admin/jobs", | |
| 245 | + "/api/v1/admin/llm-jobs?limit=2", "/api/v1/admin/entities/duplicates?type=model", "/api/v1/admin/infrastructure"): | |
| 246 | + r = await client.get(path, headers=ADMIN) | |
| 247 | + assert r.status_code == 200, (path, r.text[:200]) | |
| 248 | + docs = (await client.get("/api/v1/admin/documents?limit=1", headers=ADMIN)).json() | |
| 249 | + if docs["items"]: | |
| 250 | + doc = (await client.get(f"/api/v1/admin/documents/{docs['items'][0]['id']}", headers=ADMIN)).json() | |
| 251 | + assert "snapshots" in doc | |
| 252 | + if doc["snapshots"]: | |
| 253 | + snap = (await client.get(f"/api/v1/admin/snapshots/{doc['snapshots'][0]['id']}", headers=ADMIN)).json() | |
| 254 | + assert "raw_path" not in snap and "text_path" not in snap | |
| 255 | + assert {"structured", "diff", "text", "claims"} <= set(snap) | |
| 256 | + assert snap["text"] is None or len(snap["text"]) <= 20 * 1024 | |
| 257 | + assert (await client.post("/api/v1/admin/connectors/does-not-exist/run", headers=ADMIN, json={})).status_code == 404 | |
| 258 | + assert (await client.get("/api/v1/admin/review/nope", headers=ADMIN)).status_code in (404, 405) | |
| 259 | ||