"""`EntityDetail` builder (docs/API.md): shared by `/entities/{slug}` and the type-scoped aliases.""" from __future__ import annotations import asyncio from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection from aiatlas.api.common import ( COMPANY_TYPES, ENTITY_COLS, ENTITY_FROM, EVENT_COLS, EVENT_FROM, PRICE_COLS, PRICE_FROM, RESULT_COLS, RESULT_FROM, RESULT_ORDER, change_event, enrich_provenance, entity_summary, price_row, result_row, ) from aiatlas.db import connection, fetch_all from aiatlas.services import hardware_fit as hf LINEAGE_PREDICATES = ("derived_from", "fine_tuned_from", "distilled_from", "merged_from", "quantized_from") RELATION_GROUP_LIMIT = 24 async def relations_grouped(conn: AsyncConnection, entity_id: str) -> list[dict[str, Any]]: rows = await fetch_all(conn, f""" with rel as ( 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 union all 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), ranked as (select rel.*, row_number() over (partition by predicate, direction order by observed_at desc) as rn, count(*) over (partition by predicate, direction) as total from rel) 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 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) groups: dict[tuple[str, str], dict[str, Any]] = {} for r in rows: g = groups.setdefault((r["predicate"], r["direction"]), {"predicate": r["predicate"], "direction": r["direction"], "items": [], "total": int(r["total"])}) g["items"].append(entity_summary(r)) return list(groups.values()) async def sources_of(conn: AsyncConnection, entity_id: str, limit: int = 60) -> list[dict[str, Any]]: """Documents attached to the entity + documents whose snapshots back its claims (deduped by URL — `documents.url` is unique).""" rows = await fetch_all(conn, """ with docs as (select d.id from documents d where d.entity_id = :id union select s.document_id from claims c join snapshots s on s.id = c.snapshot_id where c.entity_id = :id 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) 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, s.domain, s.tier, (select count(*) from snapshots x where x.document_id = d.id) as snapshots from docs join documents d on d.id = docs.id left join sources s on s.id = d.source_id order by s.tier nulls last, last_observed_at desc nulls last limit :lim""", id=entity_id, lim=limit) 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"], "title": r.get("title"), "last_observed_at": r["last_observed_at"], "snapshots": int(r["snapshots"] or 0)} for r in rows] def _scope_sql(is_org: bool) -> str: if is_org: return ("(ev.entity_id = :id or ev.entity_id in (select id from entities where organization_id = :id union " "select object_id from relations where subject_id = :id and predicate in ('develops','owns','operates','published') and valid_to is null))") return "ev.entity_id = :id" async def timeline_of(conn: AsyncConnection, entity_id: str, entity_type: str, *, limit: int = 30, before: Any = None, include_documents: bool = False) -> list[dict[str, Any]]: where = [_scope_sql(entity_type in COMPANY_TYPES)] params: dict[str, Any] = {"id": entity_id, "lim": limit} if before is not None: where.append("ev.observed_at < :before") params["before"] = before if not include_documents: where.append("ev.event_type <> 'DOCUMENT_CHANGED'") 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) return [change_event(r) for r in rows] async def prices_of_model(conn: AsyncConnection, model_id: str, *, current_only: bool) -> list[dict[str, Any]]: cond = "and p.valid_to is null" if current_only else "" order = "p.input_per_mtok nulls last, pv.canonical_name" if current_only else "p.valid_from, p.id" 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) return [price_row(r) for r in rows] async def prices_of_provider(conn: AsyncConnection, provider_id: str) -> list[dict[str, Any]]: 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) return [price_row(r) for r in rows] async def results_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]: 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) return [result_row(r) for r in rows] 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]]: where = ["r.benchmark_id = :id"] params: dict[str, Any] = {"id": benchmark_id, "lim": limit, "off": offset} if not history: where.append("r.valid_to is null") if config: where.append("r.config::text ilike :cfg") params["cfg"] = f"%{config}%" 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) return [result_row(r) for r in rows] async def lineage_of(conn: AsyncConnection, model_id: str) -> dict[str, list[dict[str, Any]]]: preds = list(LINEAGE_PREDICATES) desc_preds = [p for p in preds if p != "quantized_from"] ancestors = await fetch_all(conn, f""" with recursive up as ( 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[])) union 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) 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 where e.id <> :id and e.merged_into is null order by e.id, up.depth""", id=model_id, preds=preds) descendants = await fetch_all(conn, f""" with recursive down as ( 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[])) union 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) 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 where e.id <> :id and e.merged_into is null order by e.id, down.depth""", id=model_id, preds=desc_preds) 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 where r.object_id = :id and r.predicate = 'quantized_from' and r.valid_to is null and e.merged_into is null order by e.canonical_name limit 100""", id=model_id) key = lambda r: (r.get("depth", 0), r["canonical_name"] or "") return {"ancestors": [entity_summary(r) for r in sorted(ancestors, key=key)], "descendants": [entity_summary(r) for r in sorted(descendants, key=key)], "quantizations": [entity_summary(r) for r in quants]} async def providers_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]: rows = await fetch_all(conn, f""" 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 union select p.provider_id from prices p where p.model_id = :id and p.valid_to is null) select {ENTITY_COLS} from ids join entities e on e.id = ids.id left join entities eo on eo.id = e.organization_id where e.merged_into is null order by e.canonical_name limit 100""", id=model_id) return [entity_summary(r) for r in rows] async def hardware_fit_of_model(conn: AsyncConnection, attrs: dict[str, Any]) -> list[dict[str, Any]] | None: params = hf.parameter_count(attrs) if params is None: return None context = 8192 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") out: list[dict[str, Any]] = [] for r in rows: mem = hf.hardware_memory_gb(r["attributes"]) if mem is None: continue f = hf.fit(params, mem, "4bit", context) out.append({"hardware": entity_summary(r), "memory_gb": mem, **f}) out.sort(key=lambda x: (not x["fits"], -x["headroom_gb"])) return out 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]: """Entities of `etype` linked to `entity_id` by any live relation (either direction) — optionally also those whose organization is `entity_id`.""" org_sql = "union select id from entities where organization_id = :id" if include_org_children else "" rows = await fetch_all(conn, f""" with ids as (select r.object_id as id from relations r where r.subject_id = :id and r.valid_to is null union select r.subject_id from relations r where r.object_id = :id and r.valid_to is null {org_sql}) 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 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 limit :lim offset :off""", id=entity_id, t=etype, lim=limit, off=offset) return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0 async def entity_detail(row: dict[str, Any]) -> dict[str, Any]: """Full detail. Blocks run on separate pooled connections concurrently (read-only, so no transaction needed).""" eid, etype = row["id"], row["entity_type"] async def one(fn: Any, *args: Any, **kw: Any) -> Any: async with connection() as conn: return await fn(conn, *args, **kw) provenance = dict(row.get("provenance") or {}) async def base(conn: AsyncConnection) -> dict[str, Any]: aliases = await fetch_all(conn, "select alias from entity_aliases where entity_id = :id order by kind, alias limit 200", id=eid) idents = await fetch_all(conn, "select scheme, value from entity_identifiers where entity_id = :id order by scheme, value limit 200", id=eid) await enrich_provenance(conn, provenance) return {"aliases": [a["alias"] for a in aliases], "identifiers": [{"scheme": i["scheme"], "value": i["value"]} for i in idents]} tasks: dict[str, Any] = {"base": one(base), "relations": one(relations_grouped, eid), "sources": one(sources_of, eid), "timeline": one(timeline_of, eid, etype)} if etype == "model": tasks["prices"] = one(prices_of_model, eid, current_only=True) tasks["price_history"] = one(prices_of_model, eid, current_only=False) tasks["results"] = one(results_of_model, eid) tasks["lineage"] = one(lineage_of, eid) tasks["providers"] = one(providers_of_model, eid) tasks["hardware_fit"] = one(hardware_fit_of_model, row["attributes"] or {}) tasks["papers"] = one(related_by_type, eid, "paper", limit=24) tasks["repositories"] = one(related_by_type, eid, "repository", limit=24) elif etype in COMPANY_TYPES: tasks["models"] = one(related_by_type, eid, "model", limit=50, include_org_children=True) tasks["papers"] = one(related_by_type, eid, "paper", limit=24, include_org_children=True) tasks["repositories"] = one(related_by_type, eid, "repository", limit=24, include_org_children=True) elif etype == "provider": tasks["prices"] = one(prices_of_provider, eid) tasks["models"] = one(_provider_models, eid) elif etype == "benchmark": tasks["results"] = one(leaderboard, eid, limit=100) elif etype == "hardware": tasks["models"] = one(related_by_type, eid, "model", limit=50) elif etype == "framework": tasks["repositories"] = one(related_by_type, eid, "repository", limit=24) keys = list(tasks) values = await asyncio.gather(*(tasks[k] for k in keys)) blocks = dict(zip(keys, values, strict=True)) detail = entity_summary(row) or {} detail["attributes"] = row.get("attributes") or {} detail["provenance"] = provenance detail.update(blocks.pop("base")) detail["relations"] = blocks.pop("relations") detail["sources"] = blocks.pop("sources") detail["timeline"] = blocks.pop("timeline") for k, v in blocks.items(): if k in ("models",) and isinstance(v, tuple): items, total = v detail[k] = {"items": items, "total": total, "limit": 50, "offset": 0} elif k in ("papers", "repositories") and isinstance(v, tuple): detail[k] = v[0] elif v is not None: detail[k] = v if etype == "model" and "hardware_fit" in detail: detail["hardware_fit_assumptions"] = hf.ASSUMPTIONS return detail async def _provider_models(conn: AsyncConnection, provider_id: str) -> tuple[list[dict[str, Any]], int]: rows = await fetch_all(conn, f""" with ids as (select p.model_id as id from prices p where p.provider_id = :id and p.valid_to is null union select r.subject_id from relations r where r.object_id = :id and r.predicate = 'available_through' and r.valid_to is null) 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 where e.merged_into is null order by e.canonical_name limit 50""", id=provider_id) return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0 __all__ = [ "entity_detail", "hardware_fit_of_model", "leaderboard", "lineage_of", "prices_of_model", "prices_of_provider", "providers_of_model", "related_by_type", "relations_grouped", "results_of_model", "sources_of", "timeline_of", ]