1"""/search · /search/suggest — natural-language compiler + FTS/trigram (+ embeddings when the gateway is reachable)."""2from __future__ import annotations34import asyncio5import re6import time7from typing import Any89from fastapi import APIRouter, Depends, Query, Request1011from aiatlas.api.common import ENTITY_COLS, ENTITY_FROM, ApiError, cached, entity_summary, rate_limit12from aiatlas.db import connection, fetch_all, fetch_val13from aiatlas.services import cache14from aiatlas.services.embeddings import embed_query15from aiatlas.services.llm import gateway16from aiatlas.services.search import Query as SearchQuery17from aiatlas.services.search import compile_query, search_entities, suggest1819router = APIRouter(prefix="/api/v1/search", tags=["search"])20TOTAL_CAP = 10_00021EMBED_TIMEOUT_S = 2.022EMBED_BACKOFF_S = 30023EMBED_DEGRADED_KEY = "search:embed-degraded"242526def _count_sql(q: SearchQuery) -> tuple[str, dict[str, Any]]:27 """Same filters as `search_entities`, without ranking — capped estimate."""28 where = ["e.merged_into is null"]29 params: dict[str, Any] = {}30 if q.entity_type:31 where.append("e.entity_type = :etype")32 params["etype"] = q.entity_type33 if q.openness == "open":34 where.append("(e.attributes->>'openness' in ('open-weights','open-source','open') or e.attributes->>'weights_availability' = 'open')")35 elif q.openness == "proprietary":36 where.append("e.attributes->>'openness' in ('proprietary','closed')")37 if q.year_from:38 where.append("left(e.attributes->>'release_date', 4) >= :yf")39 params["yf"] = str(q.year_from)40 if q.year_to:41 where.append("left(e.attributes->>'release_date', 4) <= :yt")42 params["yt"] = str(q.year_to)43 if q.params_min:44 where.append("(e.attributes->>'parameter_count')::double precision >= :pmin")45 params["pmin"] = float(q.params_min)46 if q.params_max:47 where.append("(e.attributes->>'parameter_count')::double precision <= :pmax")48 params["pmax"] = float(q.params_max)49 if q.context_min:50 where.append("(e.attributes->>'context_length')::double precision >= :cmin")51 params["cmin"] = float(q.context_min)52 for i, mod in enumerate(q.modalities):53 where.append(f"e.attributes->'modalities' ? :mod{i}")54 params[f"mod{i}"] = mod55 if q.organization:56 where.append("exists (select 1 from entities o where o.id = e.organization_id and o.canonical_name ilike :org)")57 params["org"] = f"%{q.organization}%"58 text = (q.filters.get("residual") or "").strip() or ("" if any([q.entity_type, q.openness, q.year_from, q.params_min, q.context_min, q.modalities]) else q.text)59 if text:60 params["q"] = text61 params["qlike"] = f"%{text}%"62 params["qprefix"] = " & ".join(f"{w}:*" for w in re.findall(r"\w+", text)[:8]) or text63 where.append("(e.search @@ to_tsquery('simple', :qprefix) or e.canonical_name ilike :qlike or e.canonical_name % :q "64 "or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")65 return f"select count(*) from (select 1 from entities e where {' and '.join(where)} limit {TOTAL_CAP}) t", params666768@router.get("", dependencies=[Depends(rate_limit("search"))])69@cached(120)70async def search(request: Request, q: str = Query("", max_length=300), type: str | None = Query(None, alias="type"),71 limit: int = Query(30, ge=1, le=200), offset: int = Query(0, ge=0)) -> dict[str, Any]:72 compiled = compile_query(q)73 if type:74 compiled.entity_type = type75 if not q.strip() and not type:76 raise ApiError(400, "q is required")77 embedding = None78 residual = (compiled.filters.get("residual") or "").strip()79 if gateway.available and residual and not await cache.cache_get(EMBED_DEGRADED_KEY):80 try:81 embedding = await asyncio.wait_for(embed_query(residual), timeout=EMBED_TIMEOUT_S)82 except Exception: # noqa: BLE001 — never wait on the LLM83 embedding = None84 if embedding is None: # back off: FTS-only for a while instead of paying the timeout on every query85 await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time()}, EMBED_BACKOFF_S)86 async with connection() as conn:87 hits = await search_entities(conn, compiled, limit=limit, offset=offset, embedding=embedding)88 ids = [h["id"] for h in hits]89 full = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where e.id = any(cast(:ids as text[]))", ids=ids) if ids else []90 count_sql, count_params = _count_sql(compiled)91 total = await fetch_val(conn, count_sql, **count_params)92 by_id = {r["id"]: r for r in full}93 items = []94 for h in hits:95 row = by_id.get(h["id"])96 s = entity_summary(row) if row else None97 if s:98 s["rank"] = float(h["rank"] or 0)99 items.append(s)100 return {"query": {**compiled.as_dict(), "semantic": embedding is not None}, "items": items, "total": int(total or 0), "limit": limit, "offset": offset}101102103@router.get("/suggest", dependencies=[Depends(rate_limit("search"))])104@cached(120)105async def search_suggest(request: Request, q: str = Query("", max_length=120), limit: int = Query(8, ge=1, le=20)) -> dict[str, Any]:106 prefix = q.strip()107 if len(prefix) < 1:108 return {"items": []}109 async with connection() as conn:110 rows = await suggest(conn, prefix, limit=limit)111 return {"items": [{"id": r["id"], "entity_type": r["entity_type"], "slug": r["slug"], "name": r["canonical_name"], "organization_name": r["organization_name"]} for r in rows]}112