unsloth/tests/python/test_rag_bm25.py

"""FTS5 lexical index lifecycle tests (sqlite-vec backed rag.db)."""

import sys
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
    sys.path.insert(0, str(STUDIO_BACKEND))

pytest.importorskip("sqlite_vec")


@pytest.fixture
def isolated_rag_db(tmp_path, monkeypatch):
    """Point rag.db at tmp_path and reset the cached connection so each
    test gets a fresh database."""
    monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
    from core.rag import db as rag_db

    rag_db._reset_for_tests()
    yield tmp_path
    rag_db._reset_for_tests()


def _chunk(cid: str, text: str, document_id: str = "doc1", index: int = 0) -> dict:
    return {
        "id": cid,
        "vector": [1.0, 0.0, 0.0, 0.0],
        "payload": {
            "document_id": document_id,
            "chunk_index": index,
            "kind": "text",
            "text": text,
        },
    }


def test_lexical_search_roundtrip(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_test"
    vector_store.upsert_chunks(
        scope,
        [
            _chunk("c1", "the quick brown fox jumps over the lazy dog", index = 0),
            _chunk(
                "c2", "machine learning models predict outputs from inputs", index = 1
            ),
            _chunk("c3", "fox terriers are small dogs", index = 2),
        ],
    )
    results = vector_store.search_lexical(scope, "fox", k = 3)
    ids = [cid for cid, _ in results]
    assert "c1" in ids
    assert "c3" in ids
    assert "c2" not in ids
    # Scores are flipped to higher-is-better.
    assert all(score >= 0 for _, score in results)


def test_lexical_empty_returns_empty(isolated_rag_db):
    from core.rag import vector_store

    assert vector_store.search_lexical("kb_nonexistent", "anything", k = 5) == []


def test_lexical_blank_query_returns_empty(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_blank"
    vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")])
    # No word tokens → no MATCH expression → empty (not a syntax error).
    assert vector_store.search_lexical(scope, "!!! ???", k = 5) == []


def test_lexical_query_operators_are_defused(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_ops"
    vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")])
    # Bare FTS5 operators would raise without sanitization.
    results = vector_store.search_lexical(scope, "alpha OR NOT (beta)", k = 5)
    assert [cid for cid, _ in results] == ["a"]


def test_lexical_scope_isolation(isolated_rag_db):
    from core.rag import vector_store

    vector_store.upsert_chunks("kb_one", [_chunk("x", "shared keyword here")])
    vector_store.upsert_chunks("kb_two", [_chunk("y", "shared keyword here")])
    results = vector_store.search_lexical("kb_one", "keyword", k = 5)
    assert [cid for cid, _ in results] == ["x"]


def test_lexical_delete_scope(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_del"
    vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta gamma")])
    assert vector_store.search_lexical(scope, "alpha", k = 1)
    vector_store.delete_scope(scope)
    assert vector_store.search_lexical(scope, "alpha", k = 1) == []


def test_lexical_delete_document(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_doc_del"
    vector_store.upsert_chunks(
        scope,
        [
            _chunk("keep1", "alpha keyword", document_id = "keep"),
            _chunk("drop1", "beta keyword", document_id = "drop"),
        ],
    )
    vector_store.delete_document(scope, "drop")
    ids = [cid for cid, _ in vector_store.search_lexical(scope, "keyword", k = 5)]
    assert ids == ["keep1"]


def test_lexical_reingest_is_idempotent(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_reingest"
    vector_store.upsert_chunks(scope, [_chunk("a", "alpha beta")])
    vector_store.upsert_chunks(scope, [_chunk("a", "gamma delta")])
    # No duplicate FTS row; old text no longer matches.
    assert vector_store.search_lexical(scope, "alpha", k = 5) == []
    ids = [cid for cid, _ in vector_store.search_lexical(scope, "gamma", k = 5)]
    assert ids == ["a"]


def test_lexical_caption_kind_is_indexed(isolated_rag_db):
    from core.rag import vector_store

    scope = "kb_caption"
    point = _chunk("cap", "Figure 1: a diagram of the pipeline")
    point["payload"]["kind"] = "caption"
    vector_store.upsert_chunks(scope, [point])
    ids = [cid for cid, _ in vector_store.search_lexical(scope, "diagram", k = 5)]
    assert ids == ["cap"]


def test_fts_backfill_from_existing_vectors(isolated_rag_db):
    """rag.db files created before FTS5 have vectors but an empty FTS table;
    opening the connection backfills it once."""
    import sqlite_vec

    from core.rag import db as rag_db

    # Seed vectors directly, bypassing upsert_chunks' incremental FTS insert.
    conn = rag_db.get_rag_connection()
    conn.execute(
        """
        INSERT INTO rag_vectors
            (chunk_id, scope, document_id, chunk_index, kind, dim, vector, payload_json)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            "legacy",
            "kb_legacy",
            "doc",
            0,
            "text",
            4,
            sqlite_vec.serialize_float32([1.0, 0.0, 0.0, 0.0]),
            '{"text": "backfilled lexical content", "kind": "text"}',
        ),
    )
    conn.execute("DELETE FROM rag_chunks_fts")
    conn.commit()
    # Force a fresh connection so _ensure_schema runs the guarded backfill.
    rag_db._reset_for_tests()

    from core.rag import vector_store

    ids = [cid for cid, _ in vector_store.search_lexical("kb_legacy", "lexical", k = 5)]
    assert ids == ["legacy"]