FabledGrimoire/tests/test_search_scope_shape.py

"""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 grimoire.models.embedding import NoteEmbedding
from grimoire.models.note import Note
from grimoire.services import embeddings as emb
from tests.helpers import compiled_sql

_SOURCE = Path("src/grimoire/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