HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""/search · /search/suggest — natural-language compiler v2 + FTS/trigram (+ embeddings when the gateway is reachable)."""2from __future__ import annotations34import asyncio5import time6from typing import Any78from fastapi import APIRouter, Depends, Query, Request9from sqlalchemy.exc import DBAPIError1011from 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, suggest, verify_organization, where_for1819router = 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. All casts are regex-guarded in `where_for`."""28 import re2930 where, params = where_for(q)31 text = (q.residual or "").strip() or ("" if q.has_structure else q.text)32 if q.benchmark and q.entity_type == "benchmark":33 text = q.benchmark34 if q.provider and q.entity_type == "provider":35 text = q.provider36 if text:37 params["q"] = text38 params["qlike"] = f"%{text}%"39 params["qprefix"] = " & ".join(f"{w}:*" for w in re.findall(r"\w+", text)[:8]) or text40 where.append("(e.search @@ to_tsquery('simple', :qprefix) or e.canonical_name ilike :qlike or e.canonical_name % :q "41 "or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")42 return f"select count(*) from (select 1 from entities e where {' and '.join(where)} limit {TOTAL_CAP}) t", params434445@router.get("", dependencies=[Depends(rate_limit("search"))])46@cached(120)47async def search(request: Request, q: str = Query("", max_length=300), type: str | None = Query(None, alias="type"),48 limit: int = Query(30, ge=1, le=200), offset: int = Query(0, ge=0)) -> dict[str, Any]:49 compiled = compile_query(q)50 if type:51 compiled.entity_type = type52 if not q.strip() and not type:53 raise ApiError(400, "q is required")54 async with connection() as conn:55 org_ok = None56 if compiled.organization:57 org_ok = await verify_organization(conn, compiled.organization)58 if org_ok is None:59 # not an organization we know: demote to free text and say so60 name = compiled.organization61 compiled.organization = None62 compiled.compiled = [c for c in compiled.compiled if c["filter"] != "organization"]63 compiled.residual = " ".join(x for x in [compiled.residual, name] if x)64 compiled.filters["residual"] = compiled.residual65 compiled.unrecognised = sorted(set(compiled.unrecognised) | {name})66 else:67 compiled.organization = org_ok["canonical_name"]68 for c in compiled.compiled:69 if c["filter"] == "organization":70 c["value"] = {"id": org_ok["id"], "slug": org_ok["slug"], "name": org_ok["canonical_name"]}71 c["label"] = f"Organization: {org_ok['canonical_name']}"72 embedding = None73 residual = (compiled.residual or "").strip()74 if gateway.available and residual and not await cache.cache_get(EMBED_DEGRADED_KEY):75 try:76 embedding = await asyncio.wait_for(embed_query(residual), timeout=EMBED_TIMEOUT_S)77 except Exception: # noqa: BLE001 — never wait on the LLM78 embedding = None79 if embedding is None: # back off: FTS-only for a while instead of paying the timeout on every query80 await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time()}, EMBED_BACKOFF_S)81 if embedding is not None:82 # the vector join needs `entity_embeddings` (pgvector); when the table is missing the statement aborts the connection's83 # transaction, so the FTS-only fallback runs on a fresh pooled connection and embeddings are backed off for a while84 try:85 async with connection() as vconn:86 hits = await search_entities(vconn, compiled, limit=limit, offset=offset, embedding=embedding)87 except DBAPIError:88 embedding = None89 await cache.cache_set(EMBED_DEGRADED_KEY, {"since": time.time(), "reason": "entity_embeddings unavailable"}, EMBED_BACKOFF_S)90 if embedding is None:91 hits = await search_entities(conn, compiled, limit=limit, offset=offset, embedding=None)92 ids = [h["id"] for h in hits]93 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 []94 count_sql, count_params = _count_sql(compiled)95 total = await fetch_val(conn, count_sql, **count_params)96 by_id = {r["id"]: r for r in full}97 items = []98 for h in hits:99 row = by_id.get(h["id"])100 s = entity_summary(row) if row else None101 if s:102 s["rank"] = float(h["rank"] or 0)103 items.append(s)104 query = {**compiled.as_dict(), "semantic": embedding is not None, "version": 2}105 if compiled.memory_gb is not None:106 query["note"] = "memory_gb is converted to an ESTIMATED parameter bound (4-bit weights, 8K context, 2 GB reserved) — see /methodology hardware_fit"107 return {"query": query, "items": items, "total": int(total or 0), "limit": limit, "offset": offset}108109110@router.get("/suggest", dependencies=[Depends(rate_limit("search"))])111@cached(120)112async def search_suggest(request: Request, q: str = Query("", max_length=120), limit: int = Query(8, ge=1, le=20)) -> dict[str, Any]:113 prefix = q.strip()114 if len(prefix) < 1:115 return {"items": []}116 async with connection() as conn:117 rows = await suggest(conn, prefix, limit=limit)118 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]}119