"""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 ]