HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""Local embeddings (via the LLM gateway's embedding endpoint) stored in pgvector. Semantic search is optional: FTS works without it."""2from __future__ import annotations34import hashlib5import logging6from typing import Any78from aiatlas.config import settings9from aiatlas.db import execute, fetch_all, fetch_one, transaction10from aiatlas.services.llm import LLMUnavailable, gateway1112log = logging.getLogger(__name__)131415def entity_text(row: dict[str, Any]) -> str:16 attrs = row.get("attributes") or {}17 bits = [row["canonical_name"], row["entity_type"].replace("_", " ")]18 for k in ("family", "architecture", "openness", "license", "modalities", "country", "kind", "category"):19 v = attrs.get(k)20 if v:21 bits.append(f"{k}: {', '.join(v) if isinstance(v, list) else v}")22 if row.get("organization_name"):23 bits.append(f"by {row['organization_name']}")24 if row.get("description"):25 bits.append(row["description"][:1200])26 return "\n".join(str(b) for b in bits)272829async def embed_entities(entity_ids: list[str]) -> dict[str, Any]:30 if not gateway.available:31 raise LLMUnavailable("embedding engine not configured")32 async with transaction() as conn:33 has_vector = await fetch_one(conn, "select 1 from pg_extension where extname = 'vector'")34 if not has_vector:35 return {"skipped": "pgvector missing"}36 rows = await fetch_all(conn, """select e.id, e.entity_type, e.canonical_name, e.description, e.attributes, o.canonical_name as organization_name37 from entities e left join entities o on o.id = e.organization_id where e.id = any(cast(:ids as text[]))""", ids=entity_ids)38 todo = []39 for r in rows:40 text = entity_text(r)41 h = hashlib.sha256((settings.embedding_model + text).encode()).hexdigest()42 async with transaction() as conn:43 existing = await fetch_one(conn, "select text_hash from entity_embeddings where entity_id = :id", id=r["id"])44 if existing and existing["text_hash"] == h:45 continue46 todo.append((r["id"], text, h))47 done = 048 for i in range(0, len(todo), 16):49 batch = todo[i:i + 16]50 vectors = await gateway.embed([t for _, t, _ in batch])51 async with transaction() as conn:52 for (eid, _, h), vec in zip(batch, vectors, strict=False):53 if len(vec) != settings.embedding_dim:54 log.warning("embedding dimension mismatch", extra={"got": len(vec), "want": settings.embedding_dim})55 continue56 await execute(conn, """insert into entity_embeddings (entity_id, model, embedding, text_hash) values (:id, :m, cast(:v as vector), :h)57 on conflict (entity_id) do update set model = excluded.model, embedding = excluded.embedding, text_hash = excluded.text_hash, created_at = now()""",58 id=eid, m=settings.embedding_model, v="[" + ",".join(f"{x:.6f}" for x in vec) + "]", h=h)59 done += 160 return {"embedded": done, "skipped": len(rows) - len(todo)}616263async def embed_query(text: str) -> list[float] | None:64 if not gateway.available:65 return None66 try:67 return (await gateway.embed([text]))[0]68 except Exception as exc: # noqa: BLE00169 log.info("query embedding failed", extra={"error": str(exc)})70 return None717273async def pending_entity_ids(limit: int = 200) -> list[str]:74 async with transaction() as conn:75 if not await fetch_one(conn, "select 1 from pg_extension where extname = 'vector'"):76 return []77 rows = await fetch_all(conn, """select e.id from entities e left join entity_embeddings x on x.entity_id = e.id78 where e.merged_into is null and e.entity_type in ('model','company','paper','provider','benchmark','hardware','framework','tool','dataset')79 and (x.entity_id is null or x.created_at < e.updated_at - interval '1 day') order by e.updated_at desc limit :n""", n=limit)80 return [r["id"] for r in rows]818283__all__ = ["embed_entities", "embed_query", "entity_text", "pending_entity_ids"]84