SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
13.1 KB · 213 lines python
Raw Blame History
1"""/entities/{slug_or_id} and its sub-resources (timeline, history, asof, graph, sources, related, claims, provenance)."""2from __future__ import annotations34from datetime import UTC, datetime5from datetime import time as dtime6from typing import Any78from fastapi import APIRouter, Query, Request910from 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)23from aiatlas.api.detail import entity_detail, sources_of, timeline_of24from aiatlas.db import connection, fetch_all, fetch_one, fetch_val2526router = APIRouter(prefix="/api/v1/entities", tags=["entities"])272829@router.get("/{slug_or_id}")30@cached(300)31async 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)353637@router.get("/{slug_or_id}/timeline")38@cached(60)39async 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), include_backfill: int = Query(0, ge=0, le=1),41                          date_field: str = Query("occurred", pattern="^(occurred|observed)$")) -> dict[str, Any]:42    before_ts = parse_ts(before, "before")43    async with connection() as conn:44        row = await resolve_entity(conn, slug_or_id)45        items = await timeline_of(conn, row["id"], row["entity_type"], limit=limit, before=before_ts, include_documents=bool(include_documents),46                                  include_backfill=bool(include_backfill), date_field=date_field)47    cursor_key = "observed_at" if date_field == "observed" else "occurred_at"48    return {"items": items, "next_before": items[-1][cursor_key] if len(items) == limit else None, "date_field": date_field, "include_backfill": bool(include_backfill)}495051@router.get("/{slug_or_id}/history")52@cached(120)53async 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]:54    async with connection() as conn:55        row = await resolve_entity(conn, slug_or_id)56        where = "c.entity_id = :id" + (" and c.property = :p" if property else "")57        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",58                               id=row["id"], p=property, lim=limit)59    return {"items": [claim_row(r) for r in rows]}606162@router.get("/{slug_or_id}/claims")63@cached(120)64async def entity_claims(request: Request, slug_or_id: str, property: str | None = Query(None, max_length=120), status: str | None = Query(None, max_length=40),65                        limit: int = Query(200, ge=1, le=2000), offset: int = Query(0, ge=0)) -> dict[str, Any]:66    """Public claim list (current by default; `status=all|superseded|conflicting|retracted`), newest first, with claim ids for `/claims/{id}`."""67    async with connection() as conn:68        row = await resolve_entity(conn, slug_or_id)69        where = ["c.entity_id = :id"]70        params: dict[str, Any] = {"id": row["id"], "lim": limit, "off": offset}71        if property:72            where.append("c.property = :p")73            params["p"] = property74        st = status or "current"75        if st != "all":76            where.append("c.status = :st")77            params["st"] = st78        rows = await fetch_all(conn, f"select {CLAIM_COLS}, c.entity_id, c.snapshot_id, c.run_id, c.extractor_version, c.value_raw from {CLAIM_FROM} where {' and '.join(where)} "79                                     f"order by c.property, c.valid_from desc limit :lim offset :off", **params)80        total = await fetch_val(conn, f"select count(*) from claims c where {' and '.join(where)}", **{k: v for k, v in params.items() if k not in ('lim', 'off')})81    items = [{**claim_row(r), "snapshot_id": r.get("snapshot_id"), "run_id": r.get("run_id"), "value_raw": r.get("value_raw")} for r in rows]82    return {"entity": entity_summary(row), "items": items, "total": int(total or 0), "limit": limit, "offset": offset, "status": st}838485@router.get("/{slug_or_id}/provenance/{property}")86@cached(120)87async def entity_provenance(request: Request, slug_or_id: str, property: str) -> dict[str, Any]:88    """Evidence-drawer payload for one property: current value, source, tier, extractor, confidence, conflicts, history count, snapshot id (no raw content)."""89    async with connection() as conn:90        row = await resolve_entity(conn, slug_or_id)91        cur = await fetch_one(conn, f"select {CLAIM_COLS}, c.snapshot_id, c.run_id, c.extractor_version, c.value_raw, s.tier as source_tier, s.domain as source_domain, "92                                    f"s.id as source_id from {CLAIM_FROM} where c.entity_id = :id and c.property = :p and c.status = 'current' order by c.tier, c.valid_from desc limit 1",93                              id=row["id"], p=property)94        conflicts = await fetch_all(conn, f"select {CLAIM_COLS}, c.snapshot_id from {CLAIM_FROM} where c.entity_id = :id and c.property = :p and c.status = 'conflicting' "95                                          f"order by c.tier, c.observed_at desc limit 20", id=row["id"], p=property)96        history = await fetch_val(conn, "select count(*) from claims c where c.entity_id = :id and c.property = :p", id=row["id"], p=property)97        snap = None98        if cur and cur.get("snapshot_id"):99            snap = await fetch_one(conn, "select s.id, s.observed_at, d.url as document_url, d.title, s.raw_path is not null as archived from snapshots s join documents d on d.id = s.document_id where s.id = :id",100                                   id=cur["snapshot_id"])101    attrs = row.get("attributes") or {}102    prov = (row.get("provenance") or {}).get(property) or {}103    if not cur and property not in attrs:104        raise ApiError(404, f"no claim for property {property!r}")105    return {"entity": {"id": row["id"], "slug": row["slug"], "name": row["canonical_name"], "entity_type": row["entity_type"]}, "property": property,106            "value": cur["value"] if cur else attrs.get(property), "value_raw": cur.get("value_raw") if cur else None, "unit": (cur or {}).get("unit") or prov.get("unit"),107            "source": {"id": (cur or {}).get("source_id") or prov.get("source_id"), "name": (cur or {}).get("source_name") or prov.get("source_name"),108                       "domain": (cur or {}).get("source_domain"), "url": (cur or {}).get("source_url") or prov.get("url")},109            "tier": (cur or {}).get("tier") or prov.get("tier"), "confidence": (cur or {}).get("confidence") or prov.get("confidence"),110            "extractor": (cur or {}).get("extractor") or prov.get("extractor"), "extractor_version": (cur or {}).get("extractor_version"),111            "observed_at": (cur or {}).get("observed_at") or prov.get("observed_at"), "effective_at": (cur or {}).get("effective_at"),112            "valid_since": (cur or {}).get("valid_from"), "claim_id": (cur or {}).get("id"), "run_id": (cur or {}).get("run_id"),113            "snapshot_id": (cur or {}).get("snapshot_id"), "snapshot": snap, "conflicts": [claim_row(c) for c in conflicts], "history_count": int(history or 0),114            "note": None if cur else "value materialised in attributes without a current claim row (curated or inherited)"}115116117@router.get("/{slug_or_id}/asof")118@cached(300)119async def entity_asof(request: Request, slug_or_id: str, date: str = Query(..., description="YYYY-MM-DD")) -> dict[str, Any]:120    d = parse_date(date, "date")121    assert d is not None122    at = datetime.combine(d, dtime.max, UTC)  # end of that day, UTC123    async with connection() as conn:124        row = await resolve_entity(conn, slug_or_id)125        existed = row["first_seen_at"] is not None and row["first_seen_at"] <= at126        rows = await fetch_all(conn, f"""select distinct on (c.property) {CLAIM_COLS} from {CLAIM_FROM}127                                         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)128                                         order by c.property, c.tier, c.valid_from desc""", id=row["id"], at=at) if existed else []129    claims = [claim_row(r) for r in rows]130    return {"id": row["id"], "slug": row["slug"], "name": row["canonical_name"], "entity_type": row["entity_type"], "existed": bool(existed),131            "first_seen_at": row["first_seen_at"], "date": d.isoformat(), "attributes": {c["property"]: c["value"] for c in claims}, "claims": claims}132133134@router.get("/{slug_or_id}/graph")135@cached(300)136async 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]:137    async with connection() as conn:138        row = await resolve_entity(conn, slug_or_id)139        root = row["id"]140        seen = {root}141        frontier = [root]142        edges: list[dict[str, str]] = []143        for _ in range(depth):144            if not frontier or len(seen) >= limit:145                break146            rows = await fetch_all(conn, """select r.subject_id, r.predicate, r.object_id from relations r147                                            where r.valid_to is null and (r.subject_id = any(cast(:ids as text[])) or r.object_id = any(cast(:ids as text[])))148                                            order by r.observed_at desc limit :lim""", ids=frontier, lim=limit * 3)149            nxt: list[str] = []150            for r in rows:151                other = r["object_id"] if r["subject_id"] in seen else r["subject_id"]152                if other not in seen:153                    if len(seen) >= limit:154                        continue155                    seen.add(other)156                    nxt.append(other)157                edges.append({"source": r["subject_id"], "target": r["object_id"], "predicate": r["predicate"]})158            frontier = nxt159        nodes = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.id = any(cast(:ids as text[]))", ids=list(seen))160    node_ids = {n["id"] for n in nodes}161    uniq = {(e["source"], e["target"], e["predicate"]): e for e in edges if e["source"] in node_ids and e["target"] in node_ids}162    return {"root": root, "nodes": [{"id": n["id"], "slug": n["slug"], "name": n["canonical_name"], "entity_type": n["entity_type"],163                                     "organization_name": n["organization_name"]} for n in nodes], "edges": list(uniq.values())}164165166@router.get("/{slug_or_id}/sources")167@cached(300)168async def entity_sources(request: Request, slug_or_id: str) -> dict[str, Any]:169    async with connection() as conn:170        row = await resolve_entity(conn, slug_or_id)171        items = await sources_of(conn, row["id"], limit=200)172    return {"items": items}173174175@router.get("/{slug_or_id}/related")176@cached(300)177async def entity_related(request: Request, slug_or_id: str, limit: int = Query(12, ge=1, le=60)) -> dict[str, Any]:178    async with connection() as conn:179        row = await resolve_entity(conn, slug_or_id)180        fam = (row.get("attributes") or {}).get("family")181        rows = await fetch_all(conn, f"""182            with cand as (183                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 <> :id184                union all185                select e.id, 4 from entities e where :fam <> '' and e.entity_type = :t and e.attributes->>'family' = :fam and e.id <> :id186                union all187                select e.id, 5 from entities e where :fid <> '' and e.entity_type = :t and e.family_id = :fid and e.id <> :id188                union all189                select r2.subject_id, 2 from relations r1 join relations r2 on r2.object_id = r1.object_id and r2.predicate = r1.predicate190                    where r1.subject_id = :id and r1.valid_to is null and r2.valid_to is null and r2.subject_id <> :id191                union all192                select r2.object_id, 1 from relations r1 join relations r2 on r2.subject_id = r1.subject_id and r2.predicate = r1.predicate193                    where r1.object_id = :id and r1.valid_to is null and r2.valid_to is null and r2.object_id <> :id),194            scored as (select id, sum(w) as w from cand group by id)195            select {ENTITY_COLS} from scored join entities e on e.id = scored.id left join entities eo on eo.id = e.organization_id196            where e.merged_into is null order by scored.w desc, coalesce((e.quality->>'score')::float, 0) desc, e.updated_at desc limit :lim""",197            id=row["id"], t=row["entity_type"], org=row.get("organization_id") or "", fam=str(fam or ""), fid=row.get("family_id") or "", lim=limit)198    return {"items": [entity_summary(r) for r in rows]}199200201async def detail_for_type(slug: str, types: tuple[str, ...]) -> dict[str, Any]:202    async with connection() as conn:203        row = await resolve_entity(conn, slug, types)204    return await entity_detail(row)205206207async def entity_exists(slug: str) -> dict[str, Any] | None:208    async with connection() as conn:209        return await fetch_one(conn, "select id, entity_type from entities where slug = :s or id = :s limit 1", s=slug)210211212__all__ = ["ApiError", "detail_for_type", "router"]213