feat(search): pgvector substrate — vector(384) + HNSW for semantic search
Move semantic_search_notes off the full-table Python cosine scan onto a native pgvector column: indexed ORDER BY embedding <=> :q LIMIT k (HNSW, cosine). Migration 0067 enables the extension, converts the JSONB embedding column to vector(384) (stale-dim rows dropped and regenerated by the startup backfill), and builds the HNSW cosine index. Postgres image moves postgres:16-alpine -> pgvector/pgvector:pg17 across prod, quickstart, and CI. Scribe: project 2, milestone 93, task 1031. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Xz4j1H7pjYSjKsEpgcNH5E
This commit is contained in:
@@ -1,11 +1,17 @@
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import DateTime, ForeignKey, Integer
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from scribe.models import Base
|
||||
|
||||
# bge-small-en-v1.5 produces 384-dim unit-normalized vectors. The column is a
|
||||
# native pgvector `vector(384)` (see migration 0067) so similarity search runs
|
||||
# as an indexed `ORDER BY embedding <=> :q LIMIT k` in Postgres rather than a
|
||||
# full-table Python cosine scan.
|
||||
EMBEDDING_DIM = 384
|
||||
|
||||
|
||||
class NoteEmbedding(Base):
|
||||
"""Stores the embedding vector for a note, used for semantic search."""
|
||||
@@ -18,7 +24,7 @@ class NoteEmbedding(Base):
|
||||
primary_key=True,
|
||||
)
|
||||
user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
|
||||
embedding: Mapped[list] = mapped_column(JSONB, nullable=False)
|
||||
embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True),
|
||||
default=lambda: datetime.now(timezone.utc),
|
||||
|
||||
@@ -28,6 +28,10 @@ logger = logging.getLogger(__name__)
|
||||
# loosely-related results that pad the sidebar without adding real value.
|
||||
_SIMILARITY_THRESHOLD = 0.45
|
||||
|
||||
# Public alias so callers (and telemetry) can record the effective default
|
||||
# threshold without reaching for the underscored name.
|
||||
DEFAULT_SIMILARITY_THRESHOLD = _SIMILARITY_THRESHOLD
|
||||
|
||||
_MODEL_NAME = "BAAI/bge-small-en-v1.5"
|
||||
_CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", "/data/fastembed-cache")
|
||||
|
||||
@@ -115,6 +119,14 @@ async def semantic_search_notes(
|
||||
|
||||
Scores are cosine similarities in [-1, 1]; only notes at or above
|
||||
*threshold* are returned, sorted highest-first.
|
||||
|
||||
Ranking and the top-k cut happen in Postgres via pgvector's cosine-distance
|
||||
operator (`<=>`, exposed as ``Vector.cosine_distance``) backed by the HNSW
|
||||
index from migration 0067 — so this is an indexed ``ORDER BY ... LIMIT k``
|
||||
rather than a full-table scan. Cosine distance is ``1 - cosine_similarity``,
|
||||
so a similarity floor of *threshold* is a distance ceiling of
|
||||
``1 - threshold`` and similarity is recovered as ``1 - distance``.
|
||||
|
||||
Returns an empty list if the embedder is unavailable or on any error.
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
@@ -125,10 +137,17 @@ async def semantic_search_notes(
|
||||
logger.debug("Semantic search skipped — embedder unavailable")
|
||||
return []
|
||||
|
||||
# Distance ceiling equivalent to the similarity floor. Clamp to the valid
|
||||
# cosine-distance range [0, 2] so a threshold of, say, -1 doesn't produce a
|
||||
# nonsensical ceiling.
|
||||
max_distance = min(2.0, max(0.0, 1.0 - threshold))
|
||||
distance = NoteEmbedding.embedding.cosine_distance(query_vec)
|
||||
|
||||
try:
|
||||
async with async_session() as session:
|
||||
stmt = (
|
||||
select(NoteEmbedding, Note)
|
||||
select(Note, distance.label("distance"))
|
||||
.select_from(NoteEmbedding)
|
||||
.join(Note, NoteEmbedding.note_id == Note.id)
|
||||
.where(NoteEmbedding.user_id == user_id, Note.deleted_at.is_(None))
|
||||
)
|
||||
@@ -142,30 +161,14 @@ async def semantic_search_notes(
|
||||
stmt = stmt.where(Note.status.is_(None))
|
||||
if exclude_ids:
|
||||
stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids))
|
||||
stmt = stmt.where(distance <= max_distance).order_by(distance.asc()).limit(limit)
|
||||
rows = list((await session.execute(stmt)).all())
|
||||
except Exception:
|
||||
logger.warning("Failed to query note embeddings", exc_info=True)
|
||||
return []
|
||||
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
def _score() -> list[tuple[float, Note]]:
|
||||
out: list[tuple[float, Note]] = []
|
||||
for ne, note in rows:
|
||||
try:
|
||||
sim = _cosine_similarity(query_vec, ne.embedding)
|
||||
except Exception:
|
||||
continue
|
||||
if sim >= threshold:
|
||||
out.append((sim, note))
|
||||
out.sort(key=lambda x: x[0], reverse=True)
|
||||
return out[:limit]
|
||||
|
||||
# Offload the O(rows) cosine scoring off the event loop so a large corpus
|
||||
# doesn't stall other requests while ranking. Results are unchanged; the
|
||||
# real scaling fix (ORDER BY / LIMIT in pgvector) is a separate effort.
|
||||
return await asyncio.to_thread(_score)
|
||||
# Recover similarity (1 - distance) and preserve the highest-first contract.
|
||||
return [(1.0 - float(dist), note) for note, dist in rows]
|
||||
|
||||
|
||||
async def backfill_note_embeddings() -> None:
|
||||
|
||||
Reference in New Issue
Block a user