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
    cnt = conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0]
    assert cnt == 3


def test_embed_incremental_only_embeds_new_chunks(db):
    """Default (incremental=True): a second embed pass after ingesting a
    new doc embeds only the new chunk — re-running is cheap (#000039)."""
    conn, _ = db
    first = embed_documents(conn, embedder=_stub_embedder)
    assert first == 3
    # Simulate a follow-up ingest: one new document + chunk.
    new_root = "root_d" * 8
    conn.execute(
        "INSERT INTO documents("
        "  document_root, document_uri, source_type, kind, title,"
        "  chunking_version, canonicalization_version, schema_version, ingest_ts"
        ") VALUES (?, 'uri://d', 'test', 'surface', 'Delta doc',"
        "  'tok-512-v1', 'norm-v1', 'v9.8.0', 0)",
        (new_root,),
    )
    conn.execute(
        "INSERT INTO chunks(document_root, idx, leaf_hash, content, tier) "
        "VALUES (?, 0, ?, 'delta content about derangements', 'hot')",
        (new_root, hashlib.sha256(b"delta").hexdigest()),
    )
    second = embed_documents(conn, embedder=_stub_embedder)  # incremental
    assert second == 1, "only the new chunk should embed on the second pass"
    assert conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0] == 4
    # A third pass with nothing new is a no-op.
    third = embed_documents(conn, embedder=_stub_embedder)
    assert third == 0


def test_embed_rebuild_re_embeds_all(db):
    """rebuild=True: DROP + recreate + full pass — every chunk re-embedded
    (the clean VEC_BACKEND_VERSION-bump path)."""
    conn, _ = db
    embed_documents(conn, embedder=_stub_embedder)
    assert conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0] == 3
    n = embed_documents(conn, embedder=_stub_embedder, rebuild=True)
    assert n == 3, "rebuild re-embeds all chunks"
    assert conn.execute("SELECT COUNT(*) FROM chunk_vecs").fetchone()[0] == 3
    # vec_meta still records the backend version after the recreate.
    from arborist.search.vec import VEC_BACKEND_VERSION
    row = conn.execute(
        "SELECT value FROM vec_meta WHERE key='backend_version'"
    ).fetchone()
    assert row[0] == VEC_BACKEND_VERSION


def test_int8_quant_roundtrips(db):
    """quant='int8': table is int8[384], vec_meta records the int8 version,
    search round-trips (a query equal to a chunk's content → that chunk on
    top, since same float vec → same int8 vec → distance ~0)."""
    from arborist.search.vec import existing_quant, vec_backend_version
    conn, docs = db
    n = embed_documents(conn, embedder=_stub_embedder, quant="int8", rebuild=True)
    assert n == 3
    assert existing_quant(conn) == "int8"
    sql = conn.execute(
        "SELECT sql FROM sqlite_master WHERE name='chunk_vecs'"
    ).fetchone()[0]
    assert "int8[384]" in sql
    assert conn.execute(
        "SELECT value FROM vec_meta WHERE key='backend_version'"
    ).fetchone()[0] == vec_backend_version("int8")

    backend = VecBackend(conn, embedder=_stub_embedder)  # quant inferred from table
    assert backend.quant == "int8"
    hits = backend.search(docs[2][3], limit=3)  # query == doc C's content
    assert hits and hits[0].document_root == docs[2][0]
    assert hits[0].audit_mode is AuditMode.UNGROUNDED


def test_quant_mismatch_requires_rebuild(db):
    """Switching quant on an existing chunk_vecs without --rebuild errors —
    the vec0 column element type can't be altered in place."""
    conn, _ = db
    embed_documents(conn, embedder=_stub_embedder, quant="float32")  # default
    with pytest.raises(ValueError, match="rebuild"):
        embed_documents(conn, embedder=_stub_embedder, quant="int8")  # no rebuild
    # ...but with rebuild it switches cleanly.
    embed_documents(conn, embedder=_stub_embedder, quant="int8", rebuild=True)
    from arborist.search.vec import existing_quant
    assert existing_quant(conn) == "int8"


@pytest.mark.parametrize("bad", ["fp16", "binary", "int4", "", None])
def test_invalid_quant_rejected(db, bad):
    conn, _ = db
    with pytest.raises(ValueError, match="quant"):
        embed_documents(conn, embedder=_stub_embedder, quant=bad)  # type: ignore[arg-type]