bvandeusenandClaude Opus 5.5
4c150cbc2f
Twenty-eight test files described the code they guard by what it did before a fix. Each is now the property and the failure it prevents, in the present tense: the backfill snapshot, the per-dial calibration query, the rule scope, the lesson trigger filter, the subshell global, the trash owner scope, the version readout fallback, and others. A docstring line that narrated the backup suite's own history is removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
"""Every semantic search scopes first, then ranks.
|
|
|
|
Ordered straight off a `*_embeddings` table, the planner walks the HNSW index:
|
|
it takes ~`hnsw.ef_search` (40) nearest chunks from every owner and project
|
|
and filters them afterwards, so an in-scope record behind 40 nearer ones the
|
|
caller cannot see is silently dropped. The remedy is one shape, `_rank_scoped`,
|
|
and these pin that every search goes through it — fixed in one search alone,
|
|
the others keep the fault, because nothing says they are the same query.
|
|
|
|
The integration lane's crowd tests show the behaviour on a real index; these
|
|
are the guard that fails on a rank-then-scope query whatever plan Postgres
|
|
picks.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from scribe.models.embedding import NoteEmbedding
|
|
from scribe.models.note import Note
|
|
from scribe.services import embeddings as emb
|
|
from tests.helpers import compiled_sql
|
|
|
|
_SOURCE = Path("src/scribe/services/embeddings.py")
|
|
|
|
|
|
def _searches() -> dict[str, ast.AsyncFunctionDef]:
|
|
tree = ast.parse(_SOURCE.read_text())
|
|
return {
|
|
node.name: node for node in tree.body
|
|
if isinstance(node, ast.AsyncFunctionDef)
|
|
and node.name.startswith("semantic_search_")
|
|
}
|
|
|
|
|
|
def _calls(fn: ast.AST) -> list[str]:
|
|
return [
|
|
getattr(n.func, "attr", None) or getattr(n.func, "id", None)
|
|
for n in ast.walk(fn) if isinstance(n, ast.Call)
|
|
]
|
|
|
|
|
|
def test_every_semantic_search_ranks_through_the_scoped_shape():
|
|
searches = _searches()
|
|
# The four corpora with an embedding table: if one is renamed or a fifth
|
|
# is added, this has to be looked at rather than silently passing.
|
|
assert set(searches) == {
|
|
"semantic_search_notes", "semantic_search_rules",
|
|
"semantic_search_milestones", "semantic_search_systems",
|
|
}
|
|
for name, fn in searches.items():
|
|
calls = _calls(fn)
|
|
assert "_rank_scoped" in calls, f"{name} does not rank through _rank_scoped"
|
|
assert "_scoped_chunks" in calls, f"{name} does not build its scope as chunks"
|
|
# Its own ORDER BY is the rank-then-scope shape: a distance ordered on the
|
|
# embedding table, which the index serves before the scope applies.
|
|
assert "order_by" not in calls, f"{name} orders by distance itself"
|
|
|
|
|
|
def test_a_rank_then_scope_query_is_caught():
|
|
"""Replay the guard against a query that ranks first and scopes after."""
|
|
old = ast.parse(
|
|
"async def semantic_search_x():\n"
|
|
" rows = await session.execute(select(Note, distance).select_from(E)"
|
|
".join(Note, E.note_id == Note.id).where(scope)"
|
|
".order_by(distance).limit(k))\n"
|
|
).body[0]
|
|
calls = _calls(old)
|
|
assert "_rank_scoped" not in calls and "order_by" in calls
|
|
|
|
|
|
def test_the_scope_is_materialized_and_the_order_is_over_it():
|
|
# Column against column, so there is no vector literal to inline: the
|
|
# shape is the subject here, not the query.
|
|
distance = NoteEmbedding.embedding.cosine_distance(NoteEmbedding.embedding)
|
|
scoped = (
|
|
emb._scoped_chunks(NoteEmbedding, NoteEmbedding.note_id, distance)
|
|
.join(Note, NoteEmbedding.note_id == Note.id)
|
|
.where(Note.user_id == 1)
|
|
)
|
|
sql = compiled_sql(
|
|
emb._rank_scoped(Note, scoped, name="scoped_x", limit=5),
|
|
dialect=postgresql.dialect(),
|
|
)
|
|
# MATERIALIZED is what keeps the planner from inlining the CTE and walking
|
|
# the index again; the ORDER BY must name the CTE's column, not the table's.
|
|
assert "scoped_x AS MATERIALIZED" in sql
|
|
assert "ORDER BY scoped_x.distance" in sql
|
|
assert "note_embeddings.embedding <=>" in sql.split("ORDER BY")[0]
|
|
# The scope sits inside the CTE, where the ranking draws from.
|
|
cte_body = sql.split("AS MATERIALIZED", 1)[1].split("SELECT notes", 1)[0]
|
|
assert "notes.user_id" in cte_body
|