1"""`EntityDetail` builder (docs/API.md): shared by `/entities/{slug}` and the type-scoped aliases."""2from __future__ import annotations34import asyncio5from typing import Any67from sqlalchemy.ext.asyncio import AsyncConnection89from 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 enrich_provenance,22 entity_summary,23 price_row,24 result_row,25)26from aiatlas.db import connection, fetch_all27from aiatlas.services import hardware_fit as hf2829LINEAGE_PREDICATES = ("derived_from", "fine_tuned_from", "distilled_from", "merged_from", "quantized_from")30RELATION_GROUP_LIMIT = 24313233async def relations_grouped(conn: AsyncConnection, entity_id: str) -> list[dict[str, Any]]:34 rows = await fetch_all(conn, f"""35 with rel as (36 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 null37 union all38 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),39 ranked as (select rel.*, row_number() over (partition by predicate, direction order by observed_at desc) as rn,40 count(*) over (partition by predicate, direction) as total from rel)41 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_id42 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)43 groups: dict[tuple[str, str], dict[str, Any]] = {}44 for r in rows:45 g = groups.setdefault((r["predicate"], r["direction"]), {"predicate": r["predicate"], "direction": r["direction"], "items": [], "total": int(r["total"])})46 g["items"].append(entity_summary(r))47 return list(groups.values())484950async def sources_of(conn: AsyncConnection, entity_id: str, limit: int = 60) -> list[dict[str, Any]]:51 """Documents attached to the entity + documents whose snapshots back its claims (deduped by URL — `documents.url` is unique)."""52 rows = await fetch_all(conn, """53 with docs as (select d.id from documents d where d.entity_id = :id54 union select s.document_id from claims c join snapshots s on s.id = c.snapshot_id where c.entity_id = :id55 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)56 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,57 s.domain, s.tier, (select count(*) from snapshots x where x.document_id = d.id) as snapshots58 from docs join documents d on d.id = docs.id left join sources s on s.id = d.source_id59 order by s.tier nulls last, last_observed_at desc nulls last limit :lim""", id=entity_id, lim=limit)60 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"],61 "title": r.get("title"), "last_observed_at": r["last_observed_at"], "snapshots": int(r["snapshots"] or 0)} for r in rows]626364def _scope_sql(is_org: bool) -> str:65 if is_org:66 return ("(ev.entity_id = :id or ev.entity_id in (select id from entities where organization_id = :id union "67 "select object_id from relations where subject_id = :id and predicate in ('develops','owns','operates','published') and valid_to is null))")68 return "ev.entity_id = :id"697071async def timeline_of(conn: AsyncConnection, entity_id: str, entity_type: str, *, limit: int = 30, before: Any = None,72 include_documents: bool = False) -> list[dict[str, Any]]:73 where = [_scope_sql(entity_type in COMPANY_TYPES)]74 params: dict[str, Any] = {"id": entity_id, "lim": limit}75 if before is not None:76 where.append("ev.observed_at < :before")77 params["before"] = before78 if not include_documents:79 where.append("ev.event_type <> 'DOCUMENT_CHANGED'")80 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)81 return [change_event(r) for r in rows]828384async def prices_of_model(conn: AsyncConnection, model_id: str, *, current_only: bool) -> list[dict[str, Any]]:85 cond = "and p.valid_to is null" if current_only else ""86 order = "p.input_per_mtok nulls last, pv.canonical_name" if current_only else "p.valid_from, p.id"87 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)88 return [price_row(r) for r in rows]899091async def prices_of_provider(conn: AsyncConnection, provider_id: str) -> list[dict[str, Any]]:92 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)93 return [price_row(r) for r in rows]949596async def results_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]:97 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)98 return [result_row(r) for r in rows]99100101async def leaderboard(conn: AsyncConnection, benchmark_id: str, *, limit: int = 100, offset: int = 0, config: str | None = None, history: bool = False) -> list[dict[str, Any]]:102 where = ["r.benchmark_id = :id"]103 params: dict[str, Any] = {"id": benchmark_id, "lim": limit, "off": offset}104 if not history:105 where.append("r.valid_to is null")106 if config:107 where.append("r.config::text ilike :cfg")108 params["cfg"] = f"%{config}%"109 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)110 return [result_row(r) for r in rows]111112113async def lineage_of(conn: AsyncConnection, model_id: str) -> dict[str, list[dict[str, Any]]]:114 preds = list(LINEAGE_PREDICATES)115 desc_preds = [p for p in preds if p != "quantized_from"]116 ancestors = await fetch_all(conn, f"""117 with recursive up as (118 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[]))119 union120 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)121 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_id122 where e.id <> :id and e.merged_into is null order by e.id, up.depth""", id=model_id, preds=preds)123 descendants = await fetch_all(conn, f"""124 with recursive down as (125 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[]))126 union127 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)128 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_id129 where e.id <> :id and e.merged_into is null order by e.id, down.depth""", id=model_id, preds=desc_preds)130 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_id131 where r.object_id = :id and r.predicate = 'quantized_from' and r.valid_to is null and e.merged_into is null132 order by e.canonical_name limit 100""", id=model_id)133 key = lambda r: (r.get("depth", 0), r["canonical_name"] or "")134 return {"ancestors": [entity_summary(r) for r in sorted(ancestors, key=key)], "descendants": [entity_summary(r) for r in sorted(descendants, key=key)],135 "quantizations": [entity_summary(r) for r in quants]}136137138async def providers_of_model(conn: AsyncConnection, model_id: str) -> list[dict[str, Any]]:139 rows = await fetch_all(conn, f"""140 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 null141 union select p.provider_id from prices p where p.model_id = :id and p.valid_to is null)142 select {ENTITY_COLS} from ids join entities e on e.id = ids.id left join entities eo on eo.id = e.organization_id143 where e.merged_into is null order by e.canonical_name limit 100""", id=model_id)144 return [entity_summary(r) for r in rows]145146147async def hardware_fit_of_model(conn: AsyncConnection, attrs: dict[str, Any]) -> list[dict[str, Any]] | None:148 params = hf.parameter_count(attrs)149 if params is None:150 return None151 context = 8192152 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")153 out: list[dict[str, Any]] = []154 for r in rows:155 mem = hf.hardware_memory_gb(r["attributes"])156 if mem is None:157 continue158 f = hf.fit(params, mem, "4bit", context)159 out.append({"hardware": entity_summary(r), "memory_gb": mem, **f})160 out.sort(key=lambda x: (not x["fits"], -x["headroom_gb"]))161 return out162163164async 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]:165 """Entities of `etype` linked to `entity_id` by any live relation (either direction) — optionally also those whose organization is `entity_id`."""166 org_sql = "union select id from entities where organization_id = :id" if include_org_children else ""167 rows = await fetch_all(conn, f"""168 with ids as (select r.object_id as id from relations r where r.subject_id = :id and r.valid_to is null169 union select r.subject_id from relations r where r.object_id = :id and r.valid_to is null {org_sql})170 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_id171 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 desc172 limit :lim offset :off""", id=entity_id, t=etype, lim=limit, off=offset)173 return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0174175176async def entity_detail(row: dict[str, Any]) -> dict[str, Any]:177 """Full detail. Blocks run on separate pooled connections concurrently (read-only, so no transaction needed)."""178 eid, etype = row["id"], row["entity_type"]179180 async def one(fn: Any, *args: Any, **kw: Any) -> Any:181 async with connection() as conn:182 return await fn(conn, *args, **kw)183184 provenance = dict(row.get("provenance") or {})185186 async def base(conn: AsyncConnection) -> dict[str, Any]:187 aliases = await fetch_all(conn, "select alias from entity_aliases where entity_id = :id order by kind, alias limit 200", id=eid)188 idents = await fetch_all(conn, "select scheme, value from entity_identifiers where entity_id = :id order by scheme, value limit 200", id=eid)189 await enrich_provenance(conn, provenance)190 return {"aliases": [a["alias"] for a in aliases], "identifiers": [{"scheme": i["scheme"], "value": i["value"]} for i in idents]}191192 tasks: dict[str, Any] = {"base": one(base), "relations": one(relations_grouped, eid), "sources": one(sources_of, eid), "timeline": one(timeline_of, eid, etype)}193 if etype == "model":194 tasks["prices"] = one(prices_of_model, eid, current_only=True)195 tasks["price_history"] = one(prices_of_model, eid, current_only=False)196 tasks["results"] = one(results_of_model, eid)197 tasks["lineage"] = one(lineage_of, eid)198 tasks["providers"] = one(providers_of_model, eid)199 tasks["hardware_fit"] = one(hardware_fit_of_model, row["attributes"] or {})200 tasks["papers"] = one(related_by_type, eid, "paper", limit=24)201 tasks["repositories"] = one(related_by_type, eid, "repository", limit=24)202 elif etype in COMPANY_TYPES:203 tasks["models"] = one(related_by_type, eid, "model", limit=50, include_org_children=True)204 tasks["papers"] = one(related_by_type, eid, "paper", limit=24, include_org_children=True)205 tasks["repositories"] = one(related_by_type, eid, "repository", limit=24, include_org_children=True)206 elif etype == "provider":207 tasks["prices"] = one(prices_of_provider, eid)208 tasks["models"] = one(_provider_models, eid)209 elif etype == "benchmark":210 tasks["results"] = one(leaderboard, eid, limit=100)211 elif etype == "hardware":212 tasks["models"] = one(related_by_type, eid, "model", limit=50)213 elif etype == "framework":214 tasks["repositories"] = one(related_by_type, eid, "repository", limit=24)215216 keys = list(tasks)217 values = await asyncio.gather(*(tasks[k] for k in keys))218 blocks = dict(zip(keys, values, strict=True))219220 detail = entity_summary(row) or {}221 detail["attributes"] = row.get("attributes") or {}222 detail["provenance"] = provenance223 detail.update(blocks.pop("base"))224 detail["relations"] = blocks.pop("relations")225 detail["sources"] = blocks.pop("sources")226 detail["timeline"] = blocks.pop("timeline")227 for k, v in blocks.items():228 if k in ("models",) and isinstance(v, tuple):229 items, total = v230 detail[k] = {"items": items, "total": total, "limit": 50, "offset": 0}231 elif k in ("papers", "repositories") and isinstance(v, tuple):232 detail[k] = v[0]233 elif v is not None:234 detail[k] = v235 if etype == "model" and "hardware_fit" in detail:236 detail["hardware_fit_assumptions"] = hf.ASSUMPTIONS237 return detail238239240async def _provider_models(conn: AsyncConnection, provider_id: str) -> tuple[list[dict[str, Any]], int]:241 rows = await fetch_all(conn, f"""242 with ids as (select p.model_id as id from prices p where p.provider_id = :id and p.valid_to is null243 union select r.subject_id from relations r where r.object_id = :id and r.predicate = 'available_through' and r.valid_to is null)244 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_id245 where e.merged_into is null order by e.canonical_name limit 50""", id=provider_id)246 return [entity_summary(r) for r in rows], int(rows[0]["total"]) if rows else 0247248249__all__ = [250 "entity_detail",251 "hardware_fit_of_model",252 "leaderboard",253 "lineage_of",254 "prices_of_model",255 "prices_of_provider",256 "providers_of_model",257 "related_by_type",258 "relations_grouped",259 "results_of_model",260 "sources_of",261 "timeline_of",262]263