@pentatonic-ai/ai-agent-sdk 0.9.4 → 0.9.6

package/packages/memory-engine/tests/test_l2_qmd_vec_search.py

@@ -0,0 +1,280 @@
+"""Tests for the sqlite-vec-backed QMD search path in l2-hybridrag-proxy.
+
+Validates the migration from the legacy Python-cosine-over-JSON path
+(which had a silent `ORDER BY id LIMIT 2000` correctness bug — only
+the OLDEST 2000 chunks were ever considered) to native sqlite-vec
+KNN MATCH over a vec0 virtual table.
+
+Pure-Python tests — no Neo4j, no Milvus. The proxy module is loaded
+via importlib so we can call helpers and handlers directly, and
+QMD_DB_PATH is overridden to a tmp_path file per test.
+
+Run:
+
+    cd packages/memory-engine
+    .venv/bin/python -m pytest tests/test_l2_qmd_vec_search.py -v
+
+The tests skip cleanly when ``sqlite_vec`` is not importable — useful
+for unit-only runs on machines that don't have the wheel installed.
+"""
+from __future__ import annotations
+
+import importlib.util
+import json
+import struct
+import sys
+from pathlib import Path
+
+import pytest
+
+try:
+    import sqlite_vec  # noqa: F401
+    _SQLITE_VEC_OK = True
+except ImportError:
+    _SQLITE_VEC_OK = False
+
+_skip_no_sqlite_vec = pytest.mark.skipif(
+    not _SQLITE_VEC_OK,
+    reason="sqlite_vec wheel not installed in this venv",
+)
+
+
+ENGINE_ROOT = Path(__file__).resolve().parent.parent / "engine" / "services" / "l2"
+sys.path.insert(0, str(ENGINE_ROOT))
+
+
+@pytest.fixture(scope="module")
+def proxy_module():
+    """Load l2-hybridrag-proxy as a module. Same pattern as
+    test_channel_stat_reader / test_people_list_reader so the
+    module-load failure mode (missing deps) skips cleanly rather than
+    erroring."""
+    spec = importlib.util.spec_from_file_location(
+        "l2_proxy_module_qmd_vec",
+        ENGINE_ROOT / "l2-hybridrag-proxy.py",
+    )
+    assert spec and spec.loader
+    try:
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+    except ImportError:
+        pytest.skip("l2 proxy deps unavailable in this venv (fine for unit-only runs)")
+    return mod
+
+
+@pytest.fixture
+def qmd_db(tmp_path, proxy_module, monkeypatch):
+    """Per-test qmd.sqlite at a tmp path, with the proxy module pointed
+    at it. Yields the path so tests can run their own asserting queries
+    against it."""
+    db_path = tmp_path / "qmd.sqlite"
+    monkeypatch.setattr(proxy_module, "QMD_DB_PATH", str(db_path))
+    return db_path
+
+
+def _make_vec(seed: int, dim: int) -> list[float]:
+    """Deterministic synthetic embedding — small enough to test fast,
+    structured enough that nearest-neighbour relationships are stable
+    across runs. The first slot dominates the cosine direction so we
+    can build orthogonal-ish clusters by varying its sign + magnitude."""
+    import random as _r
+    rng = _r.Random(seed)
+    return [rng.gauss(0.0, 1.0) for _ in range(dim)]
+
+
+# ---------------------------------------------------------------------------
+# 1. vec_index MATCH semantics — sanity check the SDK glue against sqlite-vec.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_sqlite_vec
+def test_vec_index_match_returns_top_k(qmd_db, proxy_module) -> None:
+    """Insert N known vectors with a planted ringer, query with the
+    ringer's vector, assert the ringer is the top hit. This is the
+    minimum signal that ``_ensure_vec_index`` + native MATCH actually
+    work end-to-end against the dim our proxy is configured for."""
+    conn = proxy_module._open_qmd_conn()
+    proxy_module._ensure_vec_index(conn)
+    dim = proxy_module.EMBED_DIM
+    # 20 rows of noise + 1 planted ringer at id=999. Planted vector is
+    # near-orthogonal to the noise (which uses positive-slot dominance)
+    # by flipping the first slot's sign — confirms the cosine MATCH
+    # actually orders by similarity, not by row id.
+    for i in range(20):
+        v = _make_vec(seed=i + 1, dim=dim)
+        v[0] = abs(v[0]) + 10.0  # bias positive
+        conn.execute(
+            "INSERT INTO vec_index(id, embedding) VALUES (?, ?)",
+            (i + 1, struct.pack(f"{dim}f", *v)),
+        )
+    ringer = _make_vec(seed=999, dim=dim)
+    ringer[0] = -abs(ringer[0]) - 10.0  # bias negative — opposite cluster
+    conn.execute(
+        "INSERT INTO vec_index(id, embedding) VALUES (?, ?)",
+        (999, struct.pack(f"{dim}f", *ringer)),
+    )
+    conn.commit()
+    qbytes = struct.pack(f"{dim}f", *ringer)
+    rows = conn.execute(
+        """
+        SELECT id, distance
+        FROM vec_index
+        WHERE embedding MATCH ? AND k = ?
+        ORDER BY distance
+        """,
+        (qbytes, 5),
+    ).fetchall()
+    conn.close()
+    assert len(rows) == 5
+    top_id, top_dist = rows[0]
+    assert top_id == 999, f"expected ringer id=999, got {top_id} ({rows!r})"
+    # Cosine distance = 1 - cos_sim, so identity vector → ~0 distance.
+    # Ringer-vs-itself is exact, so we expect ~0 here; allow float32
+    # round-trip slop.
+    assert top_dist < 1e-3, f"ringer-vs-itself should be ~0, got {top_dist}"
+
+
+# ---------------------------------------------------------------------------
+# 2. search_qmd_informed uses vec_index, not the legacy JSON-cosine path.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_sqlite_vec
+def test_search_qmd_informed_uses_vec_index(qmd_db, proxy_module, monkeypatch) -> None:
+    """Full search path test: seed chunks + vec_index, mock
+    ``get_embedding`` to return a vector that matches the ringer,
+    assert the returned results are sourced from the vec_index JOIN
+    (which preserves path/text from chunks) and ranked by similarity.
+
+    This is the test that would fail if someone reverted the search
+    body to the legacy ``ORDER BY id LIMIT 2000`` path — because the
+    ringer's id is 999 (well outside the 2000-row prefix), the legacy
+    path would never see it."""
+    import sqlite3
+    conn = proxy_module._open_qmd_conn()
+    proxy_module._ensure_vec_index(conn)
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS chunks (
+            id INTEGER PRIMARY KEY,
+            path TEXT,
+            text TEXT,
+            embedding TEXT,
+            embedding_model TEXT,
+            embedding_dim INTEGER,
+            chunk_index INTEGER,
+            created_at TEXT
+        )
+        """
+    )
+    dim = proxy_module.EMBED_DIM
+    # Noise rows 1..20 + planted ringer id=999. Same orthogonal-cluster
+    # setup as test 1 — guarantees the ringer wins on cosine.
+    for i in range(20):
+        v = _make_vec(seed=i + 1, dim=dim)
+        v[0] = abs(v[0]) + 10.0
+        conn.execute(
+            "INSERT INTO chunks(id, path, text, embedding) VALUES (?, ?, ?, ?)",
+            (i + 1, f"noise/{i}.md", f"noise text {i}", json.dumps(v)),
+        )
+        conn.execute(
+            "INSERT INTO vec_index(id, embedding) VALUES (?, ?)",
+            (i + 1, struct.pack(f"{dim}f", *v)),
+        )
+    ringer = _make_vec(seed=999, dim=dim)
+    ringer[0] = -abs(ringer[0]) - 10.0
+    conn.execute(
+        "INSERT INTO chunks(id, path, text, embedding) VALUES (?, ?, ?, ?)",
+        (999, "ringer/needle.md", "needle in the haystack", json.dumps(ringer)),
+    )
+    conn.execute(
+        "INSERT INTO vec_index(id, embedding) VALUES (?, ?)",
+        (999, struct.pack(f"{dim}f", *ringer)),
+    )
+    conn.commit()
+    conn.close()
+
+    # Mock get_embedding to return the ringer's vector for any query.
+    monkeypatch.setattr(proxy_module, "get_embedding", lambda *_a, **_kw: ringer)
+
+    out = proxy_module.search_qmd_informed(
+        "any query — get_embedding is mocked",
+        {"graph_entities": []},
+        limit=3,
+    )
+    assert out, "search returned empty; vec_index path must surface ringer"
+    top = out[0]
+    assert top["path"] == "ringer/needle.md", (
+        f"top hit should be the ringer at row 999; got {top['path']}. "
+        f"If this fails, the search may have reverted to the LIMIT 2000 "
+        f"legacy path which never sees row 999."
+    )
+    assert top["text"] == "needle in the haystack"
+    assert top["base_similarity"] > 0.9
+    assert top["source"] == "vector"
+
+
+# ---------------------------------------------------------------------------
+# 3. Backfill is idempotent — second run on a populated vec_index is no-op.
+# ---------------------------------------------------------------------------
+
+
+@_skip_no_sqlite_vec
+def test_backfill_idempotent(qmd_db, proxy_module) -> None:
+    """Seed chunks with JSON embeddings only (no vec_index rows),
+    call ``_backfill_vec_index`` twice, assert:
+      1. First call copies all rows into vec_index.
+      2. Second call observes vec_n >= chunks_n and is a no-op (no
+         duplicate inserts, no errors).
+
+    Catches the failure mode where a missing idempotency check would
+    INSERT duplicate ids on the second invocation, blow up the UNIQUE
+    constraint, and corrupt the index."""
+    import asyncio
+    conn = proxy_module._open_qmd_conn()
+    proxy_module._ensure_vec_index(conn)
+    conn.execute(
+        """
+        CREATE TABLE IF NOT EXISTS chunks (
+            id INTEGER PRIMARY KEY,
+            path TEXT,
+            text TEXT,
+            embedding TEXT,
+            embedding_model TEXT,
+            embedding_dim INTEGER,
+            chunk_index INTEGER,
+            created_at TEXT
+        )
+        """
+    )
+    dim = proxy_module.EMBED_DIM
+    N = 7
+    for i in range(N):
+        v = _make_vec(seed=i + 100, dim=dim)
+        conn.execute(
+            "INSERT INTO chunks(id, path, text, embedding) VALUES (?, ?, ?, ?)",
+            (i + 1, f"p/{i}.md", f"t{i}", json.dumps(v)),
+        )
+    conn.commit()
+    pre_chunks = conn.execute("SELECT count(*) FROM chunks").fetchone()[0]
+    pre_vec = conn.execute("SELECT count(*) FROM vec_index").fetchone()[0]
+    conn.close()
+    assert pre_chunks == N and pre_vec == 0, (
+        f"setup mismatch: chunks={pre_chunks}, vec={pre_vec}"
+    )
+
+    # First run — should copy all N rows.
+    asyncio.run(proxy_module._backfill_vec_index())
+    conn = proxy_module._open_qmd_conn()
+    mid_vec = conn.execute("SELECT count(*) FROM vec_index").fetchone()[0]
+    conn.close()
+    assert mid_vec == N, f"first backfill should copy all {N} rows, got {mid_vec}"
+
+    # Second run — must no-op cleanly. No exception, no duplicate inserts.
+    asyncio.run(proxy_module._backfill_vec_index())
+    conn = proxy_module._open_qmd_conn()
+    final_vec = conn.execute("SELECT count(*) FROM vec_index").fetchone()[0]
+    conn.close()
+    assert final_vec == N, (
+        f"second backfill should be no-op; got {final_vec} rows instead of {N}"
+    )