spb/ai-atlas

1"""/graph/explore — neighbourhood exploration filtered by predicate sets per mode (API 1.1). `/entities/{slug}/graph` (v1) is unchanged."""2from __future__ import annotations34from typing import Any56from fastapi import APIRouter, Query, Request78from aiatlas.api.common import ApiError, cached, resolve_entity9from aiatlas.db import connection, fetch_all1011router = APIRouter(prefix="/api/v1/graph", tags=["graph"])1213MODES: dict[str, tuple[str, ...]] = {14    "lineage": ("artifact_of", "quantized_from", "fine_tuned_from", "distilled_from", "merged_from", "derived_from", "superseded_by", "member_of_family", "develops"),15    "research": ("authored", "described_by", "published_by", "uses_dataset", "evaluates_on", "evaluated_on"),16    "company": ("develops", "owns", "operates", "manufactures", "member_of_family"),17    "benchmark": ("evaluated_on", "evaluates_on", "variant_of"),18    "dataset": ("uses_dataset", "derived_from", "subset_of"),19    "provider": ("available_through",),20    "hardware": ("runs_on", "uses", "manufactures"),21}22KEY_ATTRS = ("parameter_count", "context_length", "release_date", "openness", "published_at", "category", "memory_gb", "kind")232425@router.get("/explore")26@cached(300)27async def explore(request: Request, node: str = Query(...), mode: str = Query("lineage"), depth: int = Query(1, ge=1, le=2), limit: int = Query(150, ge=2, le=150)) -> dict[str, Any]:28    if mode not in MODES:29        raise ApiError(400, f"mode must be one of {', '.join(MODES)}")30    preds = list(MODES[mode])31    async with connection() as conn:32        root_row = await resolve_entity(conn, node)33        root = root_row["id"]34        seen: dict[str, int] = {root: 0}35        frontier = [root]36        edges: dict[tuple[str, str, str], dict[str, Any]] = {}37        truncated = False38        for level in range(1, depth + 1):39            if not frontier:40                break41            rows = await fetch_all(conn, """select r.subject_id, r.predicate, r.object_id, r.attributes, r.tier from relations r42                                            where r.valid_to is null and r.predicate = any(cast(:preds as text[]))43                                              and (r.subject_id = any(cast(:ids as text[])) or r.object_id = any(cast(:ids as text[])))44                                            order by r.tier, r.observed_at desc limit :lim""", preds=preds, ids=frontier, lim=limit * 4)45            nxt: list[str] = []46            for r in rows:47                other = r["object_id"] if r["subject_id"] in seen else r["subject_id"]48                if other not in seen:49                    if len(seen) >= limit:50                        truncated = True51                        continue52                    seen[other] = level53                    nxt.append(other)54                edges[(r["subject_id"], r["object_id"], r["predicate"])] = {"source": r["subject_id"], "target": r["object_id"], "predicate": r["predicate"], "attributes": r["attributes"] or {}, "tier": r["tier"]}55            frontier = nxt56        # organisation edges for the company mode come from entities.organization_id as well57        if mode == "company":58            org_rows = await fetch_all(conn, "select id, organization_id from entities where organization_id is not null and merged_into is null and (id = any(cast(:ids as text[])) or organization_id = any(cast(:ids as text[]))) limit :lim",59                                       ids=list(seen), lim=limit * 4)60            for r in org_rows:61                for x in (r["id"], r["organization_id"]):62                    if x not in seen:63                        if len(seen) >= limit:64                            truncated = True65                            break66                        seen[x] = depth67                edges.setdefault((r["organization_id"], r["id"], "develops"), {"source": r["organization_id"], "target": r["id"], "predicate": "develops", "attributes": {"from": "organization_id"}, "tier": None})68        nodes = await fetch_all(conn, """select e.id, e.slug, e.canonical_name, e.entity_type, e.attributes, e.merged_into, e.artifact_kind, o.canonical_name as org_name, o.slug as org_slug69                                         from entities e left join entities o on o.id = e.organization_id where e.id = any(cast(:ids as text[]))""", ids=list(seen))70    node_ids = {n["id"] for n in nodes}71    out_edges = [e for e in edges.values() if e["source"] in node_ids and e["target"] in node_ids]72    counts: dict[str, int] = {}73    for n in nodes:74        counts[n["entity_type"]] = counts.get(n["entity_type"], 0) + 175    return {"root": root, "mode": mode, "depth": depth, "predicates": preds,76            "nodes": [{"id": n["id"], "slug": n["slug"], "name": n["canonical_name"], "entity_type": n["entity_type"], "org": n["org_name"], "org_slug": n["org_slug"], "level": seen.get(n["id"], 0),77                       "artifact_kind": n["artifact_kind"], "attributes": {k: (n["attributes"] or {}).get(k) for k in KEY_ATTRS if (n["attributes"] or {}).get(k) not in (None, "", [])}} for n in nodes],78            "edges": out_edges, "truncated": truncated, "counts": {"nodes": len(nodes), "edges": len(out_edges), "by_type": counts}}79