1"""Search: Postgres FTS + trigram (+ pgvector when embeddings exist) and a deterministic natural-language → filter compiler.23 "open models released in 2026 with more than 100B parameters and 128k context"4 → {entity_type: model, openness: open, year_from: 2026, params_min: 1e11, context_min: 128000}5"""6from __future__ import annotations78import re9from dataclasses import dataclass, field10from typing import Any1112from sqlalchemy.ext.asyncio import AsyncConnection1314from aiatlas.db import fetch_all15from aiatlas.sdk.extract.numbers import parse_context_length, parse_param_count1617TYPE_WORDS = {18 "model": ("model", "models", "llm", "llms", "language model"), "company": ("company", "companies", "startup", "startups", "lab", "labs", "organization"),19 "paper": ("paper", "papers", "research", "publication", "preprint", "arxiv"), "provider": ("provider", "providers", "inference provider", "api provider"),20 "benchmark": ("benchmark", "benchmarks", "leaderboard", "eval", "evals"), "hardware": ("gpu", "gpus", "chip", "chips", "hardware", "accelerator", "npu", "tpu"),21 "framework": ("framework", "frameworks", "library", "libraries", "runtime", "runtimes"), "dataset": ("dataset", "datasets"),22 "tool": ("tool", "tools", "agent", "agents", "app", "application"), "repository": ("repo", "repos", "repository", "repositories"),23}24MODALITY_WORDS = {"vision": "image", "multimodal": "multimodal", "image": "image", "audio": "audio", "speech": "audio", "video": "video", "code": "code",25 "coding": "code", "embedding": "embedding", "embeddings": "embedding"}262728@dataclass29class Query:30 text: str = ""31 entity_type: str | None = None32 openness: str | None = None33 year_from: int | None = None34 year_to: int | None = None35 params_min: int | None = None36 params_max: int | None = None37 context_min: int | None = None38 modalities: list[str] = field(default_factory=list)39 organization: str | None = None40 filters: dict[str, Any] = field(default_factory=dict)4142 def as_dict(self) -> dict[str, Any]:43 return {k: v for k, v in self.__dict__.items() if v not in (None, "", [], {})}444546def compile_query(q: str) -> Query:47 s = q.strip()48 low = f" {s.lower()} "49 out = Query(text=s)50 for etype, words in TYPE_WORDS.items():51 if any(f" {w} " in low for w in words):52 out.entity_type = etype53 break54 if re.search(r"\b(open[- ]?(weight|weights|source)|open models?)\b", low):55 out.openness = "open"56 elif re.search(r"\b(proprietary|closed)\b", low):57 out.openness = "proprietary"58 m = re.search(r"\b(?:released|launched|published|from|since|after|in)\s+(20\d\d)\b", low)59 if m:60 out.year_from = int(m.group(1))61 if re.search(r"\bin\s+" + m.group(1), low):62 out.year_to = out.year_from63 m = re.search(r"\b(?:before|until)\s+(20\d\d)\b", low)64 if m:65 out.year_to = int(m.group(1))66 m = re.search(r"(?:more than|over|above|>|at least|≥)\s*(\d+(?:\.\d+)?\s*[bBmMtT])\b(?:\s*param)?", s)67 if m:68 out.params_min = parse_param_count(m.group(1) + " params")69 m = re.search(r"(?:less than|under|below|<|at most|≤)\s*(\d+(?:\.\d+)?\s*[bBmMtT])\b", s)70 if m:71 out.params_max = parse_param_count(m.group(1) + " params")72 m = re.search(r"(?:context|ctx)\s*(?:window\s*)?(?:>|over|above|of at least|at least|≥)?\s*(\d+\s*[kKmM])", s) or \73 re.search(r"(\d+\s*[kKmM])\s*(?:\+\s*)?(?:token\s*)?context", s)74 if m:75 out.context_min = parse_context_length(m.group(1))76 for w, mod in MODALITY_WORDS.items():77 if f" {w} " in low and mod not in out.modalities:78 out.modalities.append(mod)79 m = re.search(r"\b(?:by|from)\s+([A-Z][\w.&-]+(?:\s+[A-Z][\w.&-]+)?)", s)80 if m:81 out.organization = m.group(1)82 # residual free-text (remove the structured bits)83 residual = re.sub(r"\b(more than|over|above|less than|under|below|at least|at most|released|launched|published|since|after|before|until|with|and|in|from|by|the|a|an|context|window|params?|parameters)\b", " ", s, flags=re.I)84 residual = re.sub(r"\d+(?:\.\d+)?\s*[bBmMkKtT]\b|\b20\d\d\b|[<>≤≥]", " ", residual)85 for words in TYPE_WORDS.values():86 for w in words:87 residual = re.sub(rf"\b{re.escape(w)}\b", " ", residual, flags=re.I)88 residual = re.sub(r"\b(open[- ]?(weight|weights|source)|proprietary|closed)\b", " ", residual, flags=re.I)89 out.filters["residual"] = " ".join(residual.split())90 return out919293async def search_entities(conn: AsyncConnection, q: Query, *, limit: int = 30, offset: int = 0, embedding: list[float] | None = None) -> list[dict[str, Any]]:94 where = ["e.merged_into is null"]95 params: dict[str, Any] = {"limit": limit, "offset": offset}96 if q.entity_type:97 where.append("e.entity_type = :etype")98 params["etype"] = q.entity_type99 if q.openness == "open":100 where.append("(e.attributes->>'openness' in ('open-weights','open-source','open') or e.attributes->>'weights_availability' = 'open')")101 elif q.openness == "proprietary":102 where.append("e.attributes->>'openness' in ('proprietary','closed')")103 if q.year_from:104 where.append("left(e.attributes->>'release_date', 4) >= :yf")105 params["yf"] = str(q.year_from)106 if q.year_to:107 where.append("left(e.attributes->>'release_date', 4) <= :yt")108 params["yt"] = str(q.year_to)109 if q.params_min:110 where.append("(e.attributes->>'parameter_count')::double precision >= :pmin")111 params["pmin"] = float(q.params_min)112 if q.params_max:113 where.append("(e.attributes->>'parameter_count')::double precision <= :pmax")114 params["pmax"] = float(q.params_max)115 if q.context_min:116 where.append("(e.attributes->>'context_length')::double precision >= :cmin")117 params["cmin"] = float(q.context_min)118 for i, mod in enumerate(q.modalities):119 where.append(f"e.attributes->'modalities' ? :mod{i}")120 params[f"mod{i}"] = mod121 if q.organization:122 where.append("exists (select 1 from entities o where o.id = e.organization_id and o.canonical_name ilike :org)")123 params["org"] = f"%{q.organization}%"124 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)125 rank = "coalesce((e.quality->>'score')::float, 0) / 100.0"126 if text:127 params["q"] = text128 params["qlike"] = f"%{text}%"129 params["qprefix"] = " & ".join(f"{w}:*" for w in re.findall(r"\w+", text)[:8]) or text130 where.append("(e.search @@ to_tsquery('simple', :qprefix) or e.canonical_name ilike :qlike or e.canonical_name % :q "131 "or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :qlike))")132 rank = ("(ts_rank_cd(e.search, to_tsquery('simple', :qprefix)) * 2 + similarity(e.canonical_name, :q) * 3 "133 "+ case when e.canonical_name ilike :qlike then 1 else 0 end + coalesce((e.quality->>'score')::float, 0) / 200.0 "134 "+ case e.entity_type when 'model' then 0.3 when 'company' then 0.3 when 'provider' then 0.2 else 0 end)")135 if embedding is not None:136 params["vec"] = "[" + ",".join(f"{x:.6f}" for x in embedding) + "]"137 rank = f"({rank}) + coalesce(1 - (x.embedding <=> cast(:vec as vector)), 0) * 2"138 join = "left join entity_embeddings x on x.entity_id = e.id"139 else:140 join = ""141 sql = f"""select e.id, e.entity_type, e.canonical_name, e.slug, left(e.description, 240) as description, e.status, e.attributes, e.quality,142 o.canonical_name as organization_name, o.slug as organization_slug, {rank} as rank143 from entities e left join entities o on o.id = e.organization_id {join}144 where {' and '.join(where)} order by rank desc, e.updated_at desc limit :limit offset :offset"""145 return await fetch_all(conn, sql, **params)146147148async def suggest(conn: AsyncConnection, prefix: str, *, limit: int = 8) -> list[dict[str, Any]]:149 return await fetch_all(conn, """select e.id, e.entity_type, e.canonical_name, e.slug, o.canonical_name as organization_name150 from entities e left join entities o on o.id = e.organization_id151 where e.merged_into is null and (e.canonical_name ilike :p or exists (select 1 from entity_aliases a where a.entity_id = e.id and a.alias ilike :p))152 order by case e.entity_type when 'model' then 0 when 'company' then 1 when 'provider' then 2 else 3 end,153 coalesce((e.quality->>'score')::float, 0) desc, length(e.canonical_name) limit :n""", p=f"{prefix}%", n=limit)154155156__all__ = ["Query", "compile_query", "search_entities", "suggest"]157