spb/uqo-research-db

1# ============================================================2# SPID — embeddings & semantic search3#4# Local sentence-transformers (default all-MiniLM-L6-v2) over5# project descriptions; cosine similarity search.  Vectors are6# stored in DuckDB (FLOAT[]) — pgvector/Qdrant are drop-in7# alternatives via the same vector column.8# ============================================================9from __future__ import annotations1011import numpy as np1213from spid.config import get_logger, get_settings14from spid.store import Store1516log = get_logger("embed")1718_MODEL = None192021def _model(name: str):22    global _MODEL23    if _MODEL is None:24        from sentence_transformers import SentenceTransformer25        log.info("Loading embedding model %s ...", name)26        _MODEL = SentenceTransformer(name)27    return _MODEL282930def _doc(row: dict) -> str:31    parts = [row.get("project_name") or "", row.get("project_type") or "",32             row.get("description") or ""]33    techs = row.get("technologies") or []34    if techs:35        parts.append("technologies: " + ", ".join(techs))36    return ". ".join(p for p in parts if p)373839def build_embeddings(store: Store, batch: int = 256) -> dict:40    s = get_settings()41    rows = store.con.execute(42        "SELECT project_id, project_name, project_type, description, technologies FROM projects"43    ).to_arrow_table().to_pylist()44    if not rows:45        return {"embedded": 0}46    model = _model(s.embedding_model)47    docs = [_doc(r) for r in rows]48    vecs = model.encode(docs, batch_size=batch, show_progress_bar=False,49                        normalize_embeddings=True)50    dim = int(vecs.shape[1])51    store.con.execute("DELETE FROM project_embeddings")52    payload = [(rows[i]["project_id"], s.embedding_model, dim,53                [float(x) for x in vecs[i]]) for i in range(len(rows))]54    store.con.executemany(55        "INSERT OR REPLACE INTO project_embeddings VALUES (?,?,?,?)", payload)56    log.info("Embedded %d projects (dim=%d, model=%s).", len(rows), dim, s.embedding_model)57    return {"embedded": len(rows), "dim": dim, "model": s.embedding_model}585960def semantic_search(store: Store, query: str, k: int = 10,61                    where_sql: str | None = None) -> list[dict]:62    s = get_settings()63    rows = store.con.execute(64        "SELECT project_id, vector FROM project_embeddings"65    ).fetchall()66    if not rows:67        return []68    ids = [r[0] for r in rows]69    mat = np.asarray([r[1] for r in rows], dtype=np.float32)70    qv = _model(s.embedding_model).encode([query], normalize_embeddings=True)[0]71    sims = mat @ qv  # vectors already normalised72    order = np.argsort(-sims)7374    # join back to project metadata, applying optional filter75    top_ids = [ids[i] for i in order[: max(k * 5, k)]]76    quoted = ",".join("'" + i + "'" for i in top_ids)77    filt = f" AND {where_sql}" if where_sql else ""78    meta = {r["project_id"]: r for r in store.con.execute(79        f"SELECT * FROM projects WHERE project_id IN ({quoted}){filt}"80    ).to_arrow_table().to_pylist()}81    out = []82    sim_by_id = {ids[i]: float(sims[i]) for i in range(len(ids))}83    for pid in top_ids:84        if pid in meta:85            row = dict(meta[pid]); row["score"] = round(sim_by_id[pid], 4)86            out.append(row)87        if len(out) >= k:88            break89    return out90