bvandeusen
7827b4ce63
CI & Build / Python lint (push) Successful in 5s
CI & Build / Plugin hooks (push) Successful in 16s
CI & Build / TypeScript typecheck (push) Successful in 37s
CI & Build / integration (push) Successful in 38s
CI & Build / Python tests (push) Failing after 55s
CI & Build / Build & push image (push) Skipped
An embedding refresh replaces a record's vectors as delete-then-insert, which takes the chunk rows first and the parent row second (via the insert's foreign key). A cascading delete of the parent takes exactly those two locks in the other order. Postgres calls the cycle a deadlock and kills one side: sometimes the detached embedder, silently, and sometimes the user's delete, as a 500 on an operation that should have worked. Both upserts now claim the parent row with FOR KEY SHARE NOWAIT before touching any chunk row. That removes the cycle instead of narrowing it — either the embedder is first and the delete queues behind it, or the delete already holds the row and the embedder loses at once, which is the side designed to lose. FOR KEY SHARE is the lock the insert would take anyway, so an ordinary edit is unaffected. The note twin, recorded as unverified on the issue, has the same shape and the same fix; a trash purge is the hard delete that reaches it. Unit tests pin the ORDER and the lock mode by compiling the statement; the integration pair holds a real delete open in one transaction and proves the embedder returns having written nothing, with a deadline so a regression fails instead of hanging.
859 lines
37 KiB
Python
859 lines
37 KiB
Python
"""Semantic note search via fastembed (in-process ONNX, no external service).
|
||
|
||
Embeddings are stored as JSONB lists in the note_embeddings table (one row per
|
||
note). All search operations degrade gracefully — if the embedder fails to
|
||
initialize the callers fall back to keyword search.
|
||
|
||
Model: BAAI/bge-small-en-v1.5 (384-dim). The first call downloads the model
|
||
into `FASTEMBED_CACHE_DIR` (defaults to /data/fastembed-cache, a mounted
|
||
volume so subsequent boots are instant).
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import math
|
||
import os
|
||
|
||
from collections.abc import Sequence
|
||
|
||
from typing import TYPE_CHECKING
|
||
|
||
from sqlalchemy import delete, or_, select
|
||
|
||
from scribe.models import async_session
|
||
from scribe.models.embedding import NoteEmbedding, RuleEmbedding
|
||
from scribe.models.note import Note
|
||
from scribe.services.access import notes_visibility_clause
|
||
|
||
if TYPE_CHECKING: # resolves the Rule forward ref without importing at runtime
|
||
from scribe.models.rulebook import Rule
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Minimum cosine similarity to include a note in context results.
|
||
# bge-small-en-v1.5 produces unit-normalized vectors, so range is [-1, 1].
|
||
# 0.45 keeps only genuinely relevant notes; lower values like 0.30 let in
|
||
# loosely-related results that pad the sidebar without adding real value.
|
||
_SIMILARITY_THRESHOLD = 0.45
|
||
|
||
# The floor for INTERACTIVE, human-facing feeds — REST /api/search, Browse
|
||
# search, and the list views' semantic `q`. Deliberately looser than the agent
|
||
# default above: a human scanning a result list gets value from loosely-related
|
||
# hits an agent would be misled by. One constant, because this number was
|
||
# written as `0.3` in two files with the reasoning attached to only one of
|
||
# them — the exact shape where a later tuner moves one and not the other
|
||
# (#2463 finding 3).
|
||
INTERACTIVE_SEARCH_THRESHOLD = 0.3
|
||
|
||
# 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")
|
||
|
||
_model = None # lazy singleton; first call downloads model files
|
||
_model_lock = asyncio.Lock()
|
||
|
||
|
||
async def _get_model():
|
||
"""Return the singleton fastembed.TextEmbedding instance, loading on first call."""
|
||
global _model
|
||
if _model is None:
|
||
async with _model_lock:
|
||
if _model is None:
|
||
# Defer the import so module import doesn't pull in onnxruntime
|
||
# for non-embedding code paths (cheaper cold-start for tests etc.)
|
||
from fastembed import TextEmbedding
|
||
_model = await asyncio.to_thread(
|
||
TextEmbedding,
|
||
model_name=_MODEL_NAME,
|
||
cache_dir=_CACHE_DIR,
|
||
)
|
||
logger.info("Loaded fastembed model %s (cache: %s)", _MODEL_NAME, _CACHE_DIR)
|
||
return _model
|
||
|
||
|
||
async def get_embedding(text: str) -> list[float]:
|
||
"""Get an embedding vector for the given text.
|
||
|
||
Raises if the fastembed model fails to load. Callers should catch and
|
||
degrade to keyword search.
|
||
"""
|
||
return (await get_embeddings([text]))[0]
|
||
|
||
|
||
async def get_embeddings(texts: list[str]) -> list[list[float]]:
|
||
"""Embed several texts in one model call (the chunked write path).
|
||
|
||
fastembed batches internally, so N chunks cost far less than N single
|
||
calls. Raises like get_embedding; callers catch and degrade.
|
||
"""
|
||
embedder = await _get_model()
|
||
# embed() is synchronous CPU work; offload so we don't block the event loop.
|
||
vecs = await asyncio.to_thread(lambda: list(embedder.embed(texts)))
|
||
return [v.tolist() for v in vecs]
|
||
|
||
|
||
def _cosine_similarity(a: list[float], b: list[float]) -> float:
|
||
"""Cosine similarity between two vectors. Returns 0 for zero-length or
|
||
mismatched-length inputs (defensive — mixed-dim vectors can sneak in
|
||
across the migration boundary)."""
|
||
if not a or not b or len(a) != len(b):
|
||
return 0.0
|
||
dot = sum(x * y for x, y in zip(a, b))
|
||
mag_a = math.sqrt(sum(x * x for x in a))
|
||
mag_b = math.sqrt(sum(x * x for x in b))
|
||
if mag_a == 0.0 or mag_b == 0.0:
|
||
return 0.0
|
||
return dot / (mag_a * mag_b)
|
||
|
||
|
||
# How much a superseded record is pushed down the ranking (#278).
|
||
#
|
||
# Chosen against a measurement, not by feel. On 2026-08-07 dev-log #2420 sat at
|
||
# 0.6120 on a query made of its own title phrase, 8th, behind #1759 at 0.6506 —
|
||
# a deficit of 0.039 to the top and ~0.014 to its nearest neighbours. A penalty
|
||
# of 0.05 clears that whole band, so demoting a cluster's stale members actually
|
||
# reorders it rather than shuffling within a tie.
|
||
#
|
||
# It is deliberately NOT large. Supersession is a claim about SOME of a record's
|
||
# content, so a superseded note that strongly answers a question nothing else
|
||
# answers should still surface — just behind anything comparable that is
|
||
# current. A penalty big enough to bury it outright would be hiding by another
|
||
# name, which is the thing the operator ruled out.
|
||
_SUPERSESSION_PENALTY = 0.05
|
||
|
||
# Candidates fetched per requested result when a re-rank follows. Three ranks of
|
||
# headroom is far more than a 0.05 penalty can move anything through in a corpus
|
||
# whose neighbours sit ~0.01-0.02 apart.
|
||
_SUPERSESSION_OVERFETCH = 3
|
||
|
||
# Chunk rows fetched per requested result (#280). The HNSW top-k runs at CHUNK
|
||
# grain — several chunks of one strong note can occupy consecutive ranks, and
|
||
# each collapses into a single result. Four ranks of headroom per result keeps
|
||
# the top-k indexed while making it effectively impossible for collapsing to
|
||
# starve the result list: that would need every requested note to be shadowed
|
||
# by four chunks of notes ranked above it.
|
||
_CHUNK_OVERFETCH = 4
|
||
|
||
|
||
async def _apply_supersession_penalty(
|
||
scored: list[tuple[float, "Note"]], limit: int
|
||
) -> list[tuple[float, "Note"]]:
|
||
"""Push superseded records below their equals, then take the top `limit`.
|
||
|
||
The penalty is applied to the RANKING score and the returned score, so
|
||
downstream gates see the adjusted value — the auto-inject margin band in
|
||
particular, which exists to stop near-ties dragging in neighbours and would
|
||
otherwise re-tie exactly what this just separated.
|
||
|
||
It is NOT applied to the relevance threshold: the floor decides whether a
|
||
record is relevant at all, the penalty decides which relevant record comes
|
||
first. Applying it to the floor would drop a superseded record out of the
|
||
results entirely — hiding, which is the one thing this must not do.
|
||
|
||
Stable within a tie: Python's sort preserves the distance order the database
|
||
already established, so equal-scoring records keep their original sequence
|
||
rather than reshuffling per call.
|
||
"""
|
||
if not scored:
|
||
return []
|
||
from scribe.services.supersession import superseded_ids
|
||
|
||
try:
|
||
stale = await superseded_ids([int(note.id) for _score, note in scored])
|
||
except Exception:
|
||
# Fail OPEN, and the direction matters: ranking without the penalty is
|
||
# the behaviour that shipped for months. Returning nothing, or raising,
|
||
# would turn a supersession-lookup hiccup into a broken search.
|
||
logger.warning("Supersession lookup failed — ranking unpenalised", exc_info=True)
|
||
return scored[:limit]
|
||
|
||
if not stale:
|
||
return scored[:limit]
|
||
adjusted = [
|
||
(score - _SUPERSESSION_PENALTY if int(note.id) in stale else score, note)
|
||
for score, note in scored
|
||
]
|
||
adjusted.sort(key=lambda pair: pair[0], reverse=True)
|
||
return adjusted[:limit]
|
||
|
||
|
||
def embedding_text(title: str | None, body: str | None) -> str:
|
||
"""`title\\n{body}` — the atomic join every embedded document is built from.
|
||
|
||
One definition, deliberately. This was written out three times — the write
|
||
path (`notes.embed_note`), the recurring-task spawn, and the startup
|
||
backfill — and identical copies of a formatting rule are three chances to
|
||
change one and not the others. The spawn path is the dangerous one: a
|
||
recurring task embedded to a different shape than everything else would be
|
||
ranked against a corpus it doesn't match, and nothing would report it.
|
||
|
||
Since the chunking build (#280) this is a BUILDING BLOCK, not the whole
|
||
story: the document shape a record is embedded as is `chunk_document`
|
||
below, which calls this once per chunk. Callers that want "the text this
|
||
note is embedded as" want `chunk_document`; this stays public because the
|
||
two functions are one contract and the guard in test_embedding_text pins
|
||
both.
|
||
"""
|
||
title = title or ""
|
||
body = body or ""
|
||
return f"{title}\n{body}".strip() if body else title
|
||
|
||
|
||
# --- chunking (#280): the document shape ------------------------------------
|
||
#
|
||
# bge-small reads at most 512 tokens and fastembed silently truncates the rest,
|
||
# so a single whole-document vector loses everything past ~400 words — for a
|
||
# long dev-log, three quarters of the record was PERMANENTLY invisible to
|
||
# search. The fix is the document shape: one vector per meaningful chunk, and a
|
||
# record is as findable as its best-matching section.
|
||
|
||
# Bumped whenever chunk_document's output can change for the same input. Stored
|
||
# on every note_embeddings row so the startup backfill can re-embed exactly the
|
||
# notes whose stored shape is stale — a version comparison instead of the table
|
||
# wipe migrations 0067/0077 had to do.
|
||
CHUNKER_VERSION = 1
|
||
|
||
# Character budget approximating the model window. Tokens-per-char varies by
|
||
# content — ~4 chars/token for prose, closer to 3 for code and tables — so 1400
|
||
# chars sits at roughly 350-470 tokens, leaving headroom for the title prefixed
|
||
# to every chunk. Deliberately conservative: our own measurement (#2485) says
|
||
# shorter, single-topic documents embed SHARPER, so the cost of over-splitting
|
||
# is a few extra cheap vectors while the cost of under-splitting is truncation —
|
||
# the exact data loss this exists to end.
|
||
_CHUNK_CHAR_BUDGET = 1400
|
||
|
||
_HEADING_RE = None # compiled lazily below to keep re import local
|
||
|
||
|
||
def _split_sections(body: str) -> list[str]:
|
||
"""Split a markdown body at heading lines, fence-aware.
|
||
|
||
Each section is a heading line plus everything under it; text before the
|
||
first heading is its own section. Heading-looking lines inside ``` / ~~~
|
||
code fences do not split — a commented `# step` in a recorded shell snippet
|
||
is content, not structure.
|
||
"""
|
||
import re
|
||
global _HEADING_RE
|
||
if _HEADING_RE is None:
|
||
_HEADING_RE = re.compile(r"^#{1,6}\s")
|
||
|
||
sections: list[list[str]] = [[]]
|
||
in_fence = False
|
||
for line in body.splitlines():
|
||
if line.lstrip().startswith(("```", "~~~")):
|
||
in_fence = not in_fence
|
||
if not in_fence and _HEADING_RE.match(line) and sections[-1]:
|
||
sections.append([line])
|
||
else:
|
||
sections[-1].append(line)
|
||
return ["\n".join(chunk).strip() for chunk in sections if any(s.strip() for s in chunk)]
|
||
|
||
|
||
def _split_paragraphs(section: str, budget: int) -> list[str]:
|
||
"""Break one oversize section into budget-sized pieces at paragraph
|
||
boundaries, hard-splitting only a single paragraph that alone exceeds the
|
||
budget (a monster table or code block — split at line boundaries so no
|
||
content is dropped, which is the entire point of this module)."""
|
||
pieces: list[str] = []
|
||
current = ""
|
||
for para in section.split("\n\n"):
|
||
while len(para) > budget:
|
||
# Hard split: prefer the last newline inside the budget so lines
|
||
# stay whole, then the last space so words do; a clean char cut is
|
||
# the final resort for one enormous unbroken token.
|
||
cut = para.rfind("\n", 0, budget)
|
||
if cut <= 0:
|
||
cut = para.rfind(" ", 0, budget)
|
||
if cut <= 0:
|
||
cut = budget
|
||
head, para = para[:cut], para[cut:].lstrip("\n ")
|
||
if current:
|
||
pieces.append(current)
|
||
current = ""
|
||
pieces.append(head.strip())
|
||
if not para.strip():
|
||
continue
|
||
candidate = f"{current}\n\n{para}" if current else para
|
||
if len(candidate) > budget and current:
|
||
pieces.append(current)
|
||
current = para
|
||
else:
|
||
current = candidate
|
||
if current:
|
||
pieces.append(current)
|
||
return pieces
|
||
|
||
|
||
def chunk_document(title: str | None, body: str | None) -> list[str]:
|
||
"""The document(s) a record is embedded AS — one string per chunk.
|
||
|
||
The contract every retrieval surface builds on:
|
||
|
||
- A record that fits the model window yields EXACTLY ONE chunk, identical
|
||
to the historical `title\\nbody` shape — snippets and reference notes,
|
||
the corpus's sharpest records, are byte-for-byte unaffected.
|
||
- A longer record is split at markdown heading boundaries (fence-aware),
|
||
small neighbouring sections merged, oversize sections split at paragraph
|
||
boundaries, so every chunk fits the window. NOTHING is dropped: every
|
||
line of the body lands in some chunk.
|
||
- Every chunk is prefixed with the record's title — each vector carries
|
||
its own topical anchor, the property that makes snippets discriminative
|
||
(#2485). Pieces sub-split from one section also repeat that section's
|
||
heading line, so "which part of which topic" survives the split.
|
||
- An empty record yields [] (callers gate on falsiness to skip embedding).
|
||
|
||
Bump CHUNKER_VERSION when changing anything observable here.
|
||
"""
|
||
single = embedding_text(title, body)
|
||
if not single:
|
||
return []
|
||
if len(single) <= _CHUNK_CHAR_BUDGET:
|
||
return [single]
|
||
|
||
title = title or ""
|
||
# Budget for section content, net of the title prefix added to every chunk.
|
||
budget = max(200, _CHUNK_CHAR_BUDGET - len(title) - 1)
|
||
|
||
# Merge small adjacent sections upward so tiny sections don't each spend a
|
||
# vector, then split anything still over budget at paragraph boundaries.
|
||
merged: list[str] = []
|
||
for section in _split_sections(body or ""):
|
||
if merged and len(merged[-1]) + 2 + len(section) <= budget:
|
||
merged[-1] = f"{merged[-1]}\n\n{section}"
|
||
else:
|
||
merged.append(section)
|
||
|
||
chunks: list[str] = []
|
||
for section in merged:
|
||
if len(section) <= budget:
|
||
chunks.append(embedding_text(title, section))
|
||
continue
|
||
pieces = _split_paragraphs(section, budget)
|
||
first_line = section.split("\n", 1)[0]
|
||
heading = first_line if first_line.lstrip().startswith("#") else ""
|
||
for i, piece in enumerate(pieces):
|
||
# Repeat the section heading on continuation pieces so each vector
|
||
# still knows what topic it is part of.
|
||
if i > 0 and heading and not piece.startswith(heading):
|
||
piece = f"{heading}\n{piece}"
|
||
chunks.append(embedding_text(title, piece))
|
||
return chunks
|
||
|
||
|
||
async def _claim_parent_row(session, id_column, row_id: int, label: str) -> bool:
|
||
"""Lock the record a vector belongs to BEFORE rewriting that vector (#3262).
|
||
|
||
An embedding write and a cascading delete of the same record take the same
|
||
two row locks in OPPOSITE orders. The embedder deletes the old chunk rows
|
||
and then, on INSERT, needs the foreign key's lock on the parent; a delete
|
||
of the parent — or of the rulebook, topic or project above it — locks the
|
||
parent first and cascades down into the chunk rows. That is a cycle, and
|
||
Postgres breaks it by killing one side at random: sometimes the embedding
|
||
write, which is swallowed and invisible, and sometimes the operator's
|
||
delete, which surfaces as a 500 on an operation that should have worked.
|
||
|
||
Claiming the parent first REMOVES the cycle rather than narrowing it.
|
||
Either the embedder arrives first and the delete waits its turn behind it,
|
||
or the delete already holds the row and NOWAIT makes the embedder lose at
|
||
once. The embedder is the side that should lose: a skipped refresh costs a
|
||
stale vector until the next write or the startup backfill, and the other
|
||
outcome costs a person their request.
|
||
|
||
FOR KEY SHARE, not FOR UPDATE — it is precisely the lock the INSERT's
|
||
foreign key would take anyway, so it conflicts with a delete of the parent
|
||
and with nothing else. An ordinary edit of the same record, or a second
|
||
refresh racing this one, is unaffected.
|
||
|
||
Returns False when the row is locked or already gone; the caller skips.
|
||
"""
|
||
try:
|
||
held = (await session.execute(
|
||
select(id_column)
|
||
.where(id_column == row_id)
|
||
.with_for_update(key_share=True, nowait=True)
|
||
)).scalar_one_or_none()
|
||
except Exception:
|
||
# LockNotAvailable: this record is being deleted right now. Not an
|
||
# error — the delete wins by design.
|
||
logger.debug("Skipping embedding for %s %d — row is being deleted", label, row_id)
|
||
return False
|
||
if held is None:
|
||
logger.debug("Skipping embedding for %s %d — row is gone", label, row_id)
|
||
return False
|
||
return True
|
||
|
||
|
||
async def upsert_note_embedding(
|
||
note_id: int, user_id: int, title: str | None, body: str | None
|
||
) -> None:
|
||
"""Chunk, embed and persist a note's vectors. Safe to fire-and-forget.
|
||
|
||
Takes title/body rather than pre-built text so the chunking happens HERE —
|
||
one path for the write path, the recurrence spawn and the startup backfill,
|
||
which is the same single-definition discipline embedding_text existed for.
|
||
|
||
Replacement is atomic per note: old rows are deleted and the new chunk set
|
||
inserted in one transaction, so a concurrent read sees the old shape or the
|
||
new one, never a mixture.
|
||
"""
|
||
chunks = chunk_document(title, body)
|
||
try:
|
||
if not chunks:
|
||
# A record emptied of content should stop being findable by its
|
||
# old content — clear stale vectors rather than leaving them.
|
||
async with async_session() as session:
|
||
await session.execute(
|
||
delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id)
|
||
)
|
||
await session.commit()
|
||
return
|
||
except Exception:
|
||
logger.warning("Failed to clear embedding for note %d", note_id, exc_info=True)
|
||
return
|
||
|
||
try:
|
||
vectors = await get_embeddings(chunks)
|
||
except Exception:
|
||
logger.debug("Skipping embedding for note %d — embedder unavailable", note_id)
|
||
return
|
||
|
||
try:
|
||
async with async_session() as session:
|
||
if not await _claim_parent_row(session, Note.id, note_id, "note"):
|
||
return
|
||
await session.execute(
|
||
delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id)
|
||
)
|
||
for index, (chunk, vector) in enumerate(zip(chunks, vectors)):
|
||
session.add(
|
||
NoteEmbedding(
|
||
note_id=note_id,
|
||
chunk_index=index,
|
||
user_id=user_id,
|
||
embedding=vector,
|
||
chunk_text=chunk,
|
||
chunker_version=CHUNKER_VERSION,
|
||
)
|
||
)
|
||
await session.commit()
|
||
logger.debug("Upserted %d chunk embedding(s) for note %d", len(chunks), note_id)
|
||
except Exception:
|
||
logger.warning("Failed to persist embedding for note %d", note_id, exc_info=True)
|
||
|
||
|
||
async def semantic_search_notes(
|
||
user_id: int,
|
||
query: str,
|
||
exclude_ids: set[int] | None = None,
|
||
limit: int = 8,
|
||
threshold: float = _SIMILARITY_THRESHOLD,
|
||
project_id: int | None = None,
|
||
is_task: bool | None = None,
|
||
note_type: str | Sequence[str] | None = None,
|
||
task_kind: str | Sequence[str] | None = None,
|
||
orphan_only: bool = False,
|
||
scope: str = "own",
|
||
demote_superseded: bool = True,
|
||
system_id: int | None = None,
|
||
) -> list[tuple[float, Note]]:
|
||
"""Return up to *limit* (score, note) pairs most relevant to *query*.
|
||
|
||
Scores are cosine similarities in [-1, 1]; only notes at or above
|
||
*threshold* are returned, sorted highest-first.
|
||
|
||
`note_type` narrows to a record kind, or several (e.g. "snippet", or
|
||
("snippet", "note")), for callers that want prior art rather than everything
|
||
embedded.
|
||
|
||
`task_kind` restricts TASKS to the given kinds while leaving non-task notes
|
||
untouched. That asymmetry is the point: "recorded experience" is issues plus
|
||
dev-logs, and those differ on `is_task`, so neither `note_type` nor `is_task`
|
||
alone can express it. With `note_type="note", task_kind="issue"` a caller
|
||
gets fixed problems and durable notes without the open to-do list.
|
||
|
||
`scope` ("own" | "browse" | "read", see access.notes_visibility_clause)
|
||
decides how far this may see. It exists because this one function serves
|
||
three different kinds of act: an explicit search, which should reach
|
||
everything the caller may read; passive auto-injection, which must not pull
|
||
an unrequested record into their context; and the near-duplicate gate, whose
|
||
verdict must not depend on other people's notes at all. Defaults to "own" so
|
||
a caller that forgets is wrong in the safe direction.
|
||
|
||
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``.
|
||
|
||
`demote_superseded` applies the supersession penalty (#278): a record a
|
||
later note claims to have overtaken ranks below its equals. Callers asking
|
||
"what is the current answer" want it; the near-duplicate gate does NOT, and
|
||
passes False — a superseded record is still a duplicate of what you are
|
||
about to write, and demoting it there would let the same note be recorded
|
||
twice, the second time invisibly.
|
||
|
||
Returns an empty list if the embedder is unavailable or on any error.
|
||
"""
|
||
if not query or not query.strip():
|
||
return []
|
||
try:
|
||
query_vec = await get_embedding(query)
|
||
except Exception:
|
||
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:
|
||
# Scope on Note, not NoteEmbedding.user_id: the embedding row belongs
|
||
# to the note's owner, so filtering it would pin every scope to "own"
|
||
# and leave shared records unreachable by meaning.
|
||
stmt = (
|
||
select(Note, distance.label("distance"))
|
||
.select_from(NoteEmbedding)
|
||
.join(Note, NoteEmbedding.note_id == Note.id)
|
||
.where(
|
||
notes_visibility_clause(user_id, scope),
|
||
Note.deleted_at.is_(None),
|
||
)
|
||
)
|
||
if orphan_only:
|
||
stmt = stmt.where(Note.project_id.is_(None))
|
||
elif project_id is not None:
|
||
stmt = stmt.where(Note.project_id == project_id)
|
||
# Narrow to records tagged to one System (subsystem/area). An
|
||
# association filter, not a ranking signal — membership in the
|
||
# candidate set, decided before scoring, like project_id above.
|
||
if system_id is not None:
|
||
from scribe.models.system import RecordSystem
|
||
stmt = stmt.where(
|
||
select(RecordSystem.id)
|
||
.where(
|
||
RecordSystem.note_id == Note.id,
|
||
RecordSystem.system_id == system_id,
|
||
)
|
||
.exists()
|
||
)
|
||
if is_task is True:
|
||
stmt = stmt.where(Note.status.isnot(None))
|
||
elif is_task is False:
|
||
stmt = stmt.where(Note.status.is_(None))
|
||
# Narrow to one kind of record, or several. Composes with is_task
|
||
# rather than replacing it — 'snippet' is a non-task note_type, so a
|
||
# caller asking for prior art gets snippets and not the dev-log that
|
||
# mentions them.
|
||
if note_type:
|
||
kinds = [note_type] if isinstance(note_type, str) else list(note_type)
|
||
stmt = stmt.where(Note.note_type.in_(kinds))
|
||
# Restrict TASKS to certain kinds while leaving notes alone. A note
|
||
# has no task_kind that means anything, so a plain `.in_()` would
|
||
# drop every dev-log — which is exactly the record a caller asking
|
||
# for prior experience wants most.
|
||
if task_kind:
|
||
tkinds = [task_kind] if isinstance(task_kind, str) else list(task_kind)
|
||
stmt = stmt.where(
|
||
or_(Note.status.is_(None), Note.task_kind.in_(tkinds))
|
||
)
|
||
if exclude_ids:
|
||
stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids))
|
||
# OVER-FETCH when a re-rank follows, so the demotion can actually
|
||
# move something. Demoting after a LIMIT k would be theatre: the cut
|
||
# already happened, so a superseded record pushed down still sits in
|
||
# the results and the live record that should have replaced it was
|
||
# never fetched.
|
||
#
|
||
# Ordering stays on RAW distance so pgvector's HNSW index still
|
||
# serves it (migration 0067). Ordering by `distance + penalty`
|
||
# instead would be exact, and would turn an indexed top-k into a
|
||
# scan-and-sort of every embedded note.
|
||
#
|
||
# The cost of that trade, stated plainly: a live record outside the
|
||
# over-fetch window cannot be promoted into the results. With a
|
||
# penalty far smaller than the window's score spread, that case
|
||
# needs the true answer to be more than _SUPERSESSION_OVERFETCH
|
||
# ranks down, which no observed query comes close to.
|
||
fetch = limit * _CHUNK_OVERFETCH * (
|
||
_SUPERSESSION_OVERFETCH if demote_superseded else 1
|
||
)
|
||
stmt = (
|
||
stmt.where(distance <= max_distance)
|
||
.order_by(distance.asc())
|
||
.limit(fetch)
|
||
)
|
||
rows = list((await session.execute(stmt)).all())
|
||
except Exception:
|
||
logger.warning("Failed to query note embeddings", exc_info=True)
|
||
return []
|
||
|
||
# Collapse chunk rows to BEST-CHUNK-PER-NOTE (#280): rows arrive ordered by
|
||
# distance, so the first appearance of a note is its best chunk and later
|
||
# appearances are the same note matched less well. A note's relevance IS
|
||
# its best section's relevance — a query about one topic of a long record
|
||
# must find that record as strongly as if the topic were the whole record.
|
||
# Recover similarity (1 - distance); order stays highest-first.
|
||
scored: list[tuple[float, Note]] = []
|
||
seen: set[int] = set()
|
||
for note, dist in rows:
|
||
if int(note.id) in seen:
|
||
continue
|
||
seen.add(int(note.id))
|
||
scored.append((1.0 - float(dist), note))
|
||
if not demote_superseded:
|
||
return scored[:limit]
|
||
return await _apply_supersession_penalty(scored, limit)
|
||
|
||
|
||
async def backfill_note_embeddings() -> None:
|
||
"""(Re-)embed every note that is missing vectors OR whose stored vectors
|
||
were produced by an older chunker.
|
||
|
||
Runs as a background task at startup. Version-awareness is what makes a
|
||
document-shape change deployable: migration 0077 cleared the table once,
|
||
and every later CHUNKER_VERSION bump re-embeds the stale notes here — a
|
||
version comparison instead of another wipe. Adds a small sleep between
|
||
notes so a large backfill doesn't peg CPU.
|
||
"""
|
||
try:
|
||
async with async_session() as session:
|
||
current = {
|
||
row[0]
|
||
for row in (
|
||
await session.execute(
|
||
select(NoteEmbedding.note_id).where(
|
||
NoteEmbedding.chunker_version == CHUNKER_VERSION
|
||
)
|
||
)
|
||
).fetchall()
|
||
}
|
||
result = await session.execute(
|
||
select(Note.id, Note.user_id, Note.title, Note.body)
|
||
)
|
||
notes_to_embed = [
|
||
row for row in result.fetchall() if row[0] not in current
|
||
]
|
||
except Exception:
|
||
logger.warning("Embedding backfill: failed to query notes", exc_info=True)
|
||
return
|
||
|
||
if not notes_to_embed:
|
||
logger.info("Embedding backfill: all notes current at chunker v%d", CHUNKER_VERSION)
|
||
return
|
||
|
||
logger.info(
|
||
"Embedding backfill: embedding %d notes at chunker v%d",
|
||
len(notes_to_embed), CHUNKER_VERSION,
|
||
)
|
||
success = 0
|
||
for note_id, user_id, title, body in notes_to_embed:
|
||
if not chunk_document(title, body):
|
||
continue
|
||
await upsert_note_embedding(note_id, user_id, title, body)
|
||
success += 1
|
||
await asyncio.sleep(0.05) # gentle pacing
|
||
|
||
logger.info("Embedding backfill complete: %d/%d notes embedded", success, len(notes_to_embed))
|
||
|
||
|
||
# ── Rules (milestone 307, note 3026) ────────────────────────────────────
|
||
|
||
def rule_document(
|
||
title: str | None, statement: str | None, when_to_apply: str | None,
|
||
) -> tuple[str | None, str | None]:
|
||
"""The (title, body) a rule is EMBEDDED as — trigger first, `why` never.
|
||
|
||
Both halves of this are measured, not guessed (note 2485). That pass found
|
||
the snippet was the only sharp record in the corpus — a 0.153 top-to-second
|
||
gap against 0.010–0.023 for everything else — and that the cause was its
|
||
SHAPE: `{name} — {when_to_use}` as the title and `**When to use:** …`
|
||
repeated in the body, so purpose appears twice in a short document and
|
||
dominates the vector. This mirrors that exactly.
|
||
|
||
And it excludes `why` on the same evidence. `why` is dated incident
|
||
narrative — rule 46's runs to 4,300 characters of it — and long,
|
||
multi-topic prose is precisely what made sixteen dev-logs mutually
|
||
indistinguishable: the average lands on the centroid of "development",
|
||
which every one of them shares. Adding `why` would not give the vector more
|
||
to work with; it would give every rule the same thing to work with.
|
||
|
||
A rule with no trigger yet degrades to title + statement. It still embeds,
|
||
just less sharply — which is an argument for backfilling triggers, not an
|
||
argument for padding the document with whatever text is lying around.
|
||
"""
|
||
trigger = (when_to_apply or "").strip()
|
||
name = (title or "").strip()
|
||
body = (statement or "").strip()
|
||
if not trigger:
|
||
return name or None, body or None
|
||
return (
|
||
f"{name} — {trigger}" if name else trigger,
|
||
f"When to apply: {trigger}\n\n{body}" if body else f"When to apply: {trigger}",
|
||
)
|
||
|
||
|
||
async def upsert_rule_embedding(
|
||
rule_id: int, title: str | None, statement: str | None,
|
||
when_to_apply: str | None = None,
|
||
) -> None:
|
||
"""Chunk, embed and persist a rule's vectors. Safe to fire-and-forget.
|
||
|
||
The note twin's contract, for the same reasons: the document is built HERE
|
||
so the write path, the backfill and any re-embed share one definition, and
|
||
replacement is atomic per rule so a concurrent read sees the old chunk set
|
||
or the new one, never a mixture.
|
||
"""
|
||
from scribe.models.rulebook import Rule # runtime import: see TYPE_CHECKING above
|
||
|
||
doc_title, doc_body = rule_document(title, statement, when_to_apply)
|
||
chunks = chunk_document(doc_title, doc_body)
|
||
try:
|
||
if not chunks:
|
||
async with async_session() as session:
|
||
await session.execute(
|
||
delete(RuleEmbedding).where(RuleEmbedding.rule_id == rule_id)
|
||
)
|
||
await session.commit()
|
||
return
|
||
except Exception:
|
||
logger.warning("Failed to clear embedding for rule %d", rule_id, exc_info=True)
|
||
return
|
||
|
||
try:
|
||
vectors = await get_embeddings(chunks)
|
||
except Exception:
|
||
logger.debug("Skipping embedding for rule %d — embedder unavailable", rule_id)
|
||
return
|
||
|
||
try:
|
||
async with async_session() as session:
|
||
if not await _claim_parent_row(session, Rule.id, rule_id, "rule"):
|
||
return
|
||
await session.execute(
|
||
delete(RuleEmbedding).where(RuleEmbedding.rule_id == rule_id)
|
||
)
|
||
for index, (chunk, vector) in enumerate(zip(chunks, vectors)):
|
||
session.add(
|
||
RuleEmbedding(
|
||
rule_id=rule_id,
|
||
chunk_index=index,
|
||
embedding=vector,
|
||
chunk_text=chunk,
|
||
chunker_version=CHUNKER_VERSION,
|
||
)
|
||
)
|
||
await session.commit()
|
||
except Exception:
|
||
logger.warning("Failed to persist embedding for rule %d", rule_id, exc_info=True)
|
||
|
||
|
||
async def semantic_search_rules(
|
||
user_id: int,
|
||
query: str,
|
||
limit: int = 5,
|
||
threshold: float = _SIMILARITY_THRESHOLD,
|
||
tier: str | None = None,
|
||
) -> list[tuple[float, "Rule"]]:
|
||
"""Return up to *limit* (score, rule) pairs most relevant to *query*.
|
||
|
||
Scoped by OWNERSHIP — a rule is the caller's if they own its rulebook or
|
||
its project. Deliberately not filtered to what currently BINDS a given
|
||
project: this answers "is there a rule about this", which a person asking
|
||
wants answered across their whole rulebook. Deciding which rules bind where
|
||
is the surfacing question, and it has its own machinery
|
||
(get_applicable_rules) rather than a second, subtly different copy here.
|
||
|
||
`tier` narrows to one tier. The write-path hint passes "conditional",
|
||
because an always-on rule is ALREADY in the session — surfacing it again as
|
||
a suggestion is pure noise, and noise on a hint that fires on every write
|
||
is how a hint gets ignored.
|
||
|
||
Collapses to best-chunk-per-rule like the note search, so a long rule split
|
||
across chunks competes once rather than crowding the results with itself.
|
||
|
||
Returns an empty list if the embedder is unavailable or on any error.
|
||
"""
|
||
from scribe.models.project import Project
|
||
from scribe.models.rulebook import Rule, Rulebook, RulebookTopic
|
||
|
||
if not query or not query.strip():
|
||
return []
|
||
try:
|
||
query_vec = await get_embedding(query)
|
||
except Exception:
|
||
logger.debug("Rule search skipped — embedder unavailable")
|
||
return []
|
||
|
||
max_distance = min(2.0, max(0.0, 1.0 - threshold))
|
||
distance = RuleEmbedding.embedding.cosine_distance(query_vec)
|
||
|
||
try:
|
||
async with async_session() as session:
|
||
rows = (await session.execute(
|
||
select(Rule, distance.label("distance"))
|
||
.select_from(RuleEmbedding)
|
||
.join(Rule, RuleEmbedding.rule_id == Rule.id)
|
||
.outerjoin(RulebookTopic, Rule.topic_id == RulebookTopic.id)
|
||
.outerjoin(Rulebook, RulebookTopic.rulebook_id == Rulebook.id)
|
||
.outerjoin(Project, Rule.project_id == Project.id)
|
||
.where(
|
||
Rule.deleted_at.is_(None),
|
||
distance <= max_distance,
|
||
# topic_id XOR project_id, so exactly one arm can match.
|
||
or_(
|
||
Rulebook.owner_user_id == user_id,
|
||
Project.user_id == user_id,
|
||
),
|
||
*( [Rule.tier == tier] if tier else [] ),
|
||
)
|
||
# Overfetch so collapsing chunks to their best row still fills
|
||
# the page — the same reason the note search overfetches.
|
||
.order_by(distance)
|
||
.limit(limit * _CHUNK_OVERFETCH)
|
||
)).all()
|
||
except Exception:
|
||
logger.warning("Rule semantic search failed", exc_info=True)
|
||
return []
|
||
|
||
best: dict[int, tuple[float, object]] = {}
|
||
for rule, dist in rows:
|
||
score = 1.0 - float(dist)
|
||
if rule.id not in best or score > best[rule.id][0]:
|
||
best[rule.id] = (score, rule)
|
||
ranked = sorted(best.values(), key=lambda pair: pair[0], reverse=True)
|
||
return ranked[:limit]
|
||
|
||
|
||
async def backfill_rule_embeddings() -> None:
|
||
"""Embed rules that have no current vectors. Runs at startup beside the
|
||
note backfill; a CHUNKER_VERSION bump re-embeds rather than wiping."""
|
||
from scribe.models.rulebook import Rule
|
||
|
||
try:
|
||
async with async_session() as session:
|
||
current = select(RuleEmbedding.rule_id).where(
|
||
RuleEmbedding.chunker_version == CHUNKER_VERSION
|
||
)
|
||
stale = (await session.execute(
|
||
select(Rule.id, Rule.title, Rule.statement, Rule.when_to_apply)
|
||
.where(Rule.deleted_at.is_(None), Rule.id.notin_(current))
|
||
)).all()
|
||
except Exception:
|
||
logger.warning("Rule embedding backfill: failed to query rules", exc_info=True)
|
||
return
|
||
|
||
if not stale:
|
||
logger.info("Rule embedding backfill: all rules current at chunker v%d", CHUNKER_VERSION)
|
||
return
|
||
logger.info("Rule embedding backfill: embedding %d rule(s)", len(stale))
|
||
for rule_id, title, statement, when_to_apply in stale:
|
||
await upsert_rule_embedding(rule_id, title, statement, when_to_apply)
|