spb/ai-atlas

1"""/companies listing (company, organization, lab, university) and /companies/{slug} alias."""2from __future__ import annotations34from typing import Any56from fastapi import APIRouter, Query, Request78from aiatlas.api.common import (9    COMPANY_TYPES,10    ENTITY_COLS,11    ENTITY_FROM,12    PAGINATION,13    ApiError,14    Pagination,15    cached,16    entity_summary,17    page,18)19from aiatlas.api.routers.entities import detail_for_type20from aiatlas.db import connection, fetch_all, fetch_val2122router = APIRouter(prefix="/api/v1/companies", tags=["companies"])2324MODEL_COUNT = "(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' and m.merged_into is null)"25PAPER_COUNT = ("(select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'paper' and m.merged_into is null) + "26               "(select count(*) from relations r join entities pp on pp.id = r.subject_id where r.object_id = e.id and r.valid_to is null and pp.entity_type = 'paper' and pp.organization_id is distinct from e.id)")27SORTS = {"models": "model_count desc, e.canonical_name", "name": "e.canonical_name asc", "updated": "e.updated_at desc",28         "quality": "coalesce((e.quality->>'score')::float, 0) desc, e.canonical_name", "papers": "paper_count desc, e.canonical_name"}293031@router.get("")32@cached(300)33async def list_companies(request: Request, p: Pagination = PAGINATION, q: str | None = Query(None, max_length=200), country: str | None = None, kind: str | None = None,34                         sort: str = "models", facets: int = Query(0, ge=0, le=1)) -> dict[str, Any]:35    if sort not in SORTS:36        raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")37    where = ["e.entity_type = any(cast(:types as text[]))", "e.merged_into is null"]38    params: dict[str, Any] = {"types": list(COMPANY_TYPES)}39    if q:40        where.append("(e.canonical_name ilike :qlike or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")41        params["qlike"] = f"%{q}%"42    if country:43        where.append("e.attributes->>'country' ilike :country")44        params["country"] = country45    if kind:46        where.append("(e.attributes->>'org_kind' ilike :kind or e.entity_type = :kind)")47        params["kind"] = kind48    where_sql = " and ".join(where)49    async with connection() as conn:50        rows = await fetch_all(conn, f"select {ENTITY_COLS}, {MODEL_COUNT} as model_count, {PAPER_COUNT} as paper_count from {ENTITY_FROM} where {where_sql} "51                                     f"order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params)52        total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params)53        out = page([{**(entity_summary(r) or {}), "model_count": int(r["model_count"]), "paper_count": int(r["paper_count"])} for r in rows], int(total or 0), p)54        if facets:55            countries = await fetch_all(conn, f"select e.attributes->>'country' as value, count(*) as count from {ENTITY_FROM} where {where_sql} and e.attributes ? 'country' group by 1 order by 2 desc limit 60", **params)56            kinds = await fetch_all(conn, f"select coalesce(e.attributes->>'org_kind', e.entity_type) as value, count(*) as count from {ENTITY_FROM} where {where_sql} group by 1 order by 2 desc", **params)57            out["facets"] = {"countries": [{"value": r["value"], "count": int(r["count"])} for r in countries], "kinds": [{"value": r["value"], "count": int(r["count"])} for r in kinds]}58    return out596061@router.get("/{slug}")62@cached(300)63async def get_company(request: Request, slug: str) -> dict[str, Any]:64    return await detail_for_type(slug, COMPANY_TYPES)65