arborist/tests/test_search_vec.py

"""Tests for the optional sqlite-vec semantic-retrieval backend (#000039).

Skips entirely when the ``[vec]`` extra (sqlite-vec) is not installed.
Uses a deterministic *stub* embedder (a hash → unit-vector map) so the
suite exercises the sqlite-vec plumbing — extension load, ``chunk_vecs``
schema, ingest, KNN query, JOIN back to chunks/documents, Hit shape —
without pulling in the heavy fastembed ONNX model. Real semantic
quality is demonstrated end-to-end on a corpus shard, not unit-tested
(embedding quality isn't a property a unit test can pin).
"""

from __future__ import annotations

import hashlib

import pytest

from arborist.search import VEC_AVAILABLE

pytestmark = pytest.mark.skipif(
    not VEC_AVAILABLE, reason="sqlite-vec not installed — pip install 'arborist[vec]'"
)

if VEC_AVAILABLE:
    from arborist.search.base import AuditMode
    from arborist.search.vec import (
        EMBED_DIM,
        VEC_BACKEND_VERSION,
        VecBackend,
        embed_documents,
        ensure_chunk_vecs_table,
    )
    from arborist.store import connect


# --- deterministic stub embedder -------------------------------------


def _stub_vec(text: str) -> list[float]:
    """Map text → a fixed unit vector in R^EMBED_DIM, deterministically.

    Expand SHA-256(text) into EMBED_DIM bytes via counter mode, scale to
    [-1, 1), then L2-normalize. Same text → same vector (so a query that
    equals a chunk's text gets distance 0); different texts → (almost
    surely) different vectors.
    """
    raw = bytearray()
    i = 0
    while len(raw) < EMBED_DIM:
        raw += hashlib.sha256(text.encode("utf-8") + i.to_bytes(4, "little")).digest()
        i += 1
    vals = [2.0 * (b / 256.0) - 1.0 for b in raw[:EMBED_DIM]]
    norm = sum(v * v for v in vals) ** 0.5 or 1.0
    return [v / norm for v in vals]


def _stub_embedder(texts: list[str]) -> list[list[float]]:
    return [_stub_vec(t) for t in texts]


# --- fixtures --------------------------------------------------------


@pytest.fixture()
def db(tmp_path):
    conn = connect(tmp_path / "vec-test.db")
    # Three minimal documents, one chunk each.
    docs = [
        ("root_a" * 8, "uri://a", "Alpha doc", "alpha content about combinations"),
        ("root_b" * 8, "uri://b", "Beta doc", "beta content about permutations"),
        ("root_c" * 8, "uri://c", "Gamma doc", "gamma content about pigeonhole"),
    ]
    for root, uri, title, content in docs:
        conn.execute(
            "INSERT INTO documents("
            "  document_root, document_uri, source_type, kind, title,"
            "  chunking_version, canonicalization_version, schema_version, ingest_ts"
            ") VALUES (?, ?, 'test', 'surface', ?, 'tok-512-v1', 'norm-v1', 'v9.8.0', 0)",
            (root, uri, title),
        )
        conn.execute(
            "INSERT INTO chunks(document_root, idx, leaf_hash, content, tier) "
            "VALUES (?, 0, ?, ?, 'hot')",
            (root, hashlib.sha256(content.encode()).hexdigest(), content),
        )
    yield conn, docs
    conn.close()


# --- tests -----------------------------------------------------------


def test_ensure_table_records_backend_version(db):
    conn, _ = db
    ensure_chunk_vecs_table(conn)
    row = conn.execute(
        "SELECT value FROM vec_meta WHERE key='backend_version'"
    ).fetchone()
    assert row[0] == VEC_BACKEND_VERSION


def test_embed_then_search_roundtrips(db):
    conn, docs = db
    n = embed_documents(conn, embedder=_stub_embedder)
    assert n == 3, "all three chunks should embed (none cold)"

    backend = VecBackend(conn, embedder=_stub_embedder)
    assert backend.populated()

    # Query with text == doc B's content → distance 0 → top hit is doc B.
    target_content = docs[1][3]
    hits = backend.search(target_content, limit=3)
    assert hits, "search should return hits"
    assert hits[0].document_root == docs[1][0]
    assert hits[0].audit_mode is AuditMode.UNGROUNDED
    assert hits[0].title == "Beta doc"
    assert hits[0].chunk_idx == 0
    # Score is the negated distance: exact match → ~0; ordering is
    # monotone (best first), so score is non-increasing down the list.
    assert hits[0].score >= hits[-1].score
    assert hits[0].snippet  # non-empty snippet rendered


def test_search_respects_limit(db):
    conn, _ = db
    embed_documents(conn, embedder=_stub_embedder)
    backend = VecBackend(conn, embedder=_stub_embedder)
    hits = backend.search("anything", limit=2)
    assert len(hits) <= 2


def test_search_empty_query_returns_nothing(db):
    conn, _ = db
    embed_documents(conn, embedder=_stub_embedder)
    backend = VecBackend(conn, embedder=_stub_embedder)
    assert backend.search("   ", limit=5) == []


def test_unpopulated_backend_returns_nothing(db):
    conn, _ = db
    backend = VecBackend(conn, embedder=_stub_embedder)
    assert not backend.populated()
    assert backend.search("anything", limit=5) == []


def test_embed_limit_caps_count(db):
    conn, _ = db
    n = embed_documents(conn, embedder=_stub_embedder, limit=2)
    assert n == 2
    cnt = conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0]
    assert cnt == 2


def test_embed_is_idempotent(db):
    conn, _ = db
    embed_documents(conn, embedder=_stub_embedder)
    embed_documents(conn, embedder=_stub_embedder)  # re-run; INSERT OR REPLACE
    cnt = conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0]
    assert cnt == 3