retrieval.py

"""pgvector semantic search over filing chunks."""

import threading

from config import get_settings
from db import connection

_model = None
_lock = threading.Lock()


def _embedder():
    global _model
    if _model is None:
        with _lock:
            if _model is None:
                from sentence_transformers import SentenceTransformer

                _model = SentenceTransformer(get_settings().embedding_model)
    return _model


def search_passages(
    query: str,
    ticker: str | None = None,
    form: str | None = None,
    section: str | None = None,
    k: int | None = None,
) -> list[dict]:
    k = k or get_settings().retrieval_k
    vector = _embedder().encode(query, normalize_embeddings=True)

    sql = """
        select ch.id, c.ticker, f.form, f.filing_date::text, f.accession,
               ch.section, ch.text, 1 - (ch.embedding <=> %s) as score
        from chunks ch
        join filings f on f.id = ch.filing_id
        join companies c on c.cik = ch.cik
    """
    where, params = [], [vector]
    if ticker:
        where.append("c.ticker = %s")
        params.append(ticker.upper())
    if form:
        where.append("f.form = %s")
        params.append(form.upper())
    if section:
        where.append("ch.section ilike %s")
        params.append(section)
    if where:
        sql += " where " + " and ".join(where)
    sql += " order by ch.embedding <=> %s limit %s"
    params += [vector, k]

    with connection() as conn:
        # HNSW returns global nearest neighbors before WHERE filtering; with a
        # ticker filter that can leave almost nothing. Widen the candidate pool
        # and let pgvector keep scanning until the limit is satisfied.
        conn.execute("set local hnsw.ef_search = 400")
        try:
            conn.execute("set local hnsw.iterative_scan = 'relaxed_order'")
        except Exception:
            pass  # pgvector < 0.8
        rows = conn.execute(sql, params).fetchall()

    return [
        {
            "chunk_id": row[0],
            "ticker": row[1],
            "form": row[2],
            "filing_date": row[3],
            "accession": row[4],
            "section": row[5],
            "text": row[6],
            "score": round(float(row[7]), 4),
        }
        for row in rows
    ]