SPB Git

spb/ai-atlas

Public
7.8 KB · 153 lines python
1"""/entities/{slug_or_id} and its sub-resources (timeline, history, asof, graph, sources, related)."""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_one2526router = 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)) -> 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}464748@router.get("/{slug_or_id}/history")49@cached(120)50async 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]}575859@router.get("/{slug_or_id}/asof")60@cached(300)61async 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 None64    at = datetime.combine(d, dtime.max, UTC)  # end of that day, UTC65    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"] <= at68        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}747576@router.get("/{slug_or_id}/graph")77@cached(300)78async 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                break88            rows = await fetch_all(conn, """select r.subject_id, r.predicate, r.object_id from relations r89                                            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                        continue97                    seen.add(other)98                    nxt.append(other)99                edges.append({"source": r["subject_id"], "target": r["object_id"], "predicate": r["predicate"]})100            frontier = nxt101        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())}106107108@router.get("/{slug_or_id}/sources")109@cached(300)110async 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}115116117@router.get("/{slug_or_id}/related")118@cached(300)119async 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 <> :id126                union all127                select e.id, 4 from entities e where :fam <> '' and e.entity_type = :t and e.attributes->>'family' = :fam and e.id <> :id128                union all129                select r2.subject_id, 2 from relations r1 join relations r2 on r2.object_id = r1.object_id and r2.predicate = r1.predicate130                    where r1.subject_id = :id and r1.valid_to is null and r2.valid_to is null and r2.subject_id <> :id131                union all132                select r2.object_id, 1 from relations r1 join relations r2 on r2.subject_id = r1.subject_id and r2.predicate = r1.predicate133                    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_id136            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]}139140141async 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)145146147async 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)150151152__all__ = ["ApiError", "detail_for_type", "router"]153
Files · 9f1db171c786cf4a3b99436329f29b145c02aeec
Menu

Theme