chore: import upstream snapshot with attribution

tests/integration/vector/test_valkey_search.py

@@ -0,0 +1,344 @@
+"""Integration tests for the Valkey Search VDB backend.
+
+These are SLOW, real-server tests.  They are gated on ``TEST_VALKEY_URL`` and
+skipped when it is unset (same precedent as the PostgreSQL migration tests).
+
+Run locally against valkey/valkey-bundle:9.1.0::
+
+    podman run -d --name valkey-test-langbot -p 6380:6379 valkey/valkey-bundle:9.1.0
+    TEST_VALKEY_URL=valkey://localhost:6380 \\
+        uv run pytest tests/integration/vector/test_valkey_search.py -m slow -q
+
+The default upstream fast CI lane (``-m "not slow"``) skips these; the local
+supervisor validator MUST run them.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import uuid
+from types import SimpleNamespace
+from urllib.parse import urlparse
+
+import pytest
+
+pytestmark = [pytest.mark.integration, pytest.mark.slow]
+
+
+def _parse_valkey_url(url: str) -> tuple[str, int, int]:
+    """Parse ``valkey://host:port/db`` into ``(host, port, db)``."""
+    parsed = urlparse(url)
+    host = parsed.hostname or 'localhost'
+    port = parsed.port or 6379
+    db = 0
+    if parsed.path and parsed.path.strip('/'):
+        try:
+            db = int(parsed.path.strip('/'))
+        except ValueError:
+            db = 0
+    return host, port, db
+
+
+@pytest.fixture
+def valkey_config():
+    url = os.environ.get('TEST_VALKEY_URL')
+    if not url:
+        pytest.skip('TEST_VALKEY_URL not set')
+    host, port, db = _parse_valkey_url(url)
+    return {
+        'host': host,
+        'port': port,
+        'db': db,
+        'password': '',
+        'username': '',
+        'tls': False,
+        'index_algorithm': 'HNSW',
+        'distance_metric': 'COSINE',
+    }
+
+
+def _make_ap(valkey_config):
+    """Build a minimal fake ``ap`` with the config + a no-op logger."""
+    logger = SimpleNamespace(
+        info=lambda *a, **k: None,
+        warning=lambda *a, **k: None,
+        error=lambda *a, **k: None,
+        debug=lambda *a, **k: None,
+    )
+    instance_config = SimpleNamespace(data={'vdb': {'valkey_search': valkey_config}})
+    return SimpleNamespace(instance_config=instance_config, logger=logger)
+
+
+@pytest.fixture
+async def backend(valkey_config):
+    """Create a Valkey Search backend, skip if module/server unavailable."""
+    from langbot.pkg.vector.vdbs.valkey_search import (
+        ValkeySearchVectorDatabase,
+        VALKEY_SEARCH_AVAILABLE,
+    )
+
+    if not VALKEY_SEARCH_AVAILABLE:
+        pytest.skip('valkey-glide not installed')
+
+    from glide import ft
+
+    ap = _make_ap(valkey_config)
+    db = ValkeySearchVectorDatabase(ap)
+    client = await db._ensure_client()
+
+    # Module-presence gate: FT.LIST must be available (Search module loaded).
+    try:
+        await ft.list(client)
+    except Exception as exc:  # noqa: BLE001
+        await client.close()
+        pytest.skip(f'Valkey Search module not available: {exc}')
+
+    collection = f'test_{uuid.uuid4().hex[:12]}'
+    yield db, collection
+
+    # Cleanup
+    try:
+        await db.delete_collection(collection)
+    except Exception:
+        pass
+    if db._client is not None:
+        await db._client.close()
+
+
+async def _poll_until(coro_factory, predicate, timeout=5.0, interval=0.2):
+    """Poll an async result until predicate is true (indexer is async)."""
+    deadline = asyncio.get_event_loop().time() + timeout
+    result = await coro_factory()
+    while not predicate(result) and asyncio.get_event_loop().time() < deadline:
+        await asyncio.sleep(interval)
+        result = await coro_factory()
+    return result
+
+
+def _sample_docs():
+    ids = ['d1', 'd2', 'd3']
+    embeddings = [
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.9, 0.1, 0.0, 0.0],
+    ]
+    metadatas = [
+        {'file_id': 'fileA', 'topic': 'cats'},
+        {'file_id': 'fileB', 'topic': 'dogs'},
+        {'file_id': 'fileA', 'topic': 'cats'},
+    ]
+    documents = [
+        'the quick brown fox',
+        'lazy dogs sleeping',
+        'foxes and cats playing',
+    ]
+    return ids, embeddings, metadatas, documents
+
+
+@pytest.mark.asyncio
+async def test_add_and_vector_search(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+
+    result = await _poll_until(
+        lambda: db.search(collection, [1.0, 0.0, 0.0, 0.0], k=3, search_type='vector'),
+        lambda r: len(r['ids'][0]) >= 1,
+    )
+    assert len(result['ids'][0]) >= 1
+    # Closest to [1,0,0,0] should be d1.
+    assert result['ids'][0][0] == 'd1'
+    assert all(isinstance(d, float) for d in result['distances'][0])
+
+
+@pytest.mark.asyncio
+async def test_full_text_search(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+
+    result = await _poll_until(
+        lambda: db.search(collection, [0.0, 0.0, 0.0, 0.0], k=5, search_type='full_text', query_text='dogs'),
+        lambda r: len(r['ids'][0]) >= 1,
+    )
+    assert 'd2' in result['ids'][0]
+
+
+@pytest.mark.asyncio
+async def test_hybrid_filter_then_knn(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+
+    result = await _poll_until(
+        lambda: db.search(
+            collection,
+            [1.0, 0.0, 0.0, 0.0],
+            k=5,
+            search_type='hybrid',
+            query_text='cats',
+            filter={'file_id': 'fileA'},
+        ),
+        lambda r: len(r['ids'][0]) >= 1,
+    )
+    # Only fileA docs (d1, d3) should be candidates.
+    assert set(result['ids'][0]).issubset({'d1', 'd3'})
+
+
+@pytest.mark.asyncio
+async def test_vector_weight_not_honored(backend):
+    """Passing different vector_weight values must NOT change ranking."""
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+
+    common = dict(
+        collection=collection, query_embedding=[1.0, 0.0, 0.0, 0.0], k=3, search_type='hybrid', query_text='cats'
+    )
+    await _poll_until(lambda: db.search(**common), lambda r: len(r['ids'][0]) >= 1)
+
+    r_low = await db.search(**common, vector_weight=0.1)
+    r_high = await db.search(**common, vector_weight=0.9)
+    assert r_low['ids'][0] == r_high['ids'][0]
+
+
+@pytest.mark.asyncio
+async def test_filter_operators(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+
+    # Wait for indexing.
+    await _poll_until(
+        lambda: db.list_by_filter(collection, limit=10),
+        lambda r: r[1] >= 3,
+    )
+
+    # $eq
+    items, total = await db.list_by_filter(collection, filter={'file_id': 'fileA'})
+    assert total == 2
+    assert {it['id'] for it in items} == {'d1', 'd3'}
+
+    # $ne
+    items, total = await db.list_by_filter(collection, filter={'file_id': {'$ne': 'fileA'}})
+    assert {it['id'] for it in items} == {'d2'}
+
+    # $in
+    items, total = await db.list_by_filter(collection, filter={'file_id': {'$in': ['fileA', 'fileB']}})
+    assert total == 3
+
+    # $nin
+    items, total = await db.list_by_filter(collection, filter={'file_id': {'$nin': ['fileB']}})
+    assert {it['id'] for it in items} == {'d1', 'd3'}
+
+
+@pytest.mark.asyncio
+async def test_delete_by_file_id(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+    await _poll_until(lambda: db.list_by_filter(collection, limit=10), lambda r: r[1] >= 3)
+
+    await db.delete_by_file_id(collection, 'fileA')
+    items, total = await _poll_until(
+        lambda: db.list_by_filter(collection, limit=10),
+        lambda r: r[1] <= 1,
+    )
+    assert {it['id'] for it in items} == {'d2'}
+
+
+@pytest.mark.asyncio
+async def test_delete_by_filter_returns_count(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+    await _poll_until(lambda: db.list_by_filter(collection, limit=10), lambda r: r[1] >= 3)
+
+    deleted = await db.delete_by_filter(collection, filter={'file_id': 'fileA'})
+    assert deleted == 2
+
+
+@pytest.mark.asyncio
+async def test_list_by_filter_pagination(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+    await _poll_until(lambda: db.list_by_filter(collection, limit=10), lambda r: r[1] >= 3)
+
+    page1, total = await db.list_by_filter(collection, limit=2, offset=0)
+    assert total == 3
+    assert len(page1) == 2
+
+    page2, total = await db.list_by_filter(collection, limit=2, offset=2)
+    assert total == 3
+    assert len(page2) == 1
+
+
+@pytest.mark.asyncio
+async def test_delete_collection(backend):
+    db, collection = backend
+    ids, embeddings, metadatas, documents = _sample_docs()
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+    await _poll_until(lambda: db.list_by_filter(collection, limit=10), lambda r: r[1] >= 3)
+
+    await db.delete_collection(collection)
+
+    # After dropping, search on a missing index returns empty.
+    result = await db.search(collection, [1.0, 0.0, 0.0, 0.0], k=3, search_type='vector')
+    assert result['ids'][0] == []
+
+
+@pytest.mark.asyncio
+async def test_adversarial_filter_and_query_input(backend):
+    """Crafted FT special chars in file_id / query_text must not break out.
+
+    Guarantees locked in here:
+    * A file_id full of injection-style chars (quotes, parens, ``|``, ``@``,
+      ``:``, spaces, dashes) only ever matches its own row — the payload is
+      escaped to literal TAG content, never interpreted as extra clauses.
+    * A query_text full of FT operators does not raise and does not widen the
+      result set.
+    * A file_id containing FT-unsafe chars (``{`` / ``}`` / ``*``) is
+      percent-encoded, so it round-trips correctly: an exact match returns ONLY
+      its own row and never widens to an unrelated row, and the query does not
+      raise.
+    """
+    db, collection = backend
+
+    # Injection-style file_id WITHOUT FT-unsafe chars (the realistic surface).
+    injection_fid = 'evil") @file_id (".id|x-y:z'
+    # file_id WITH FT-unsafe chars that previously could not be queried.
+    brace_fid = 'x} @file_id:{*'
+    ids = ['adv1', 'benign2', 'brace3']
+    embeddings = [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]]
+    metadatas = [{'file_id': injection_fid}, {'file_id': 'plainB'}, {'file_id': brace_fid}]
+    documents = ['payload row content', 'unrelated benign content', 'brace row content']
+    await db.add_embeddings(collection, ids, embeddings, metadatas, documents)
+    await _poll_until(lambda: db.list_by_filter(collection, limit=10), lambda r: r[1] >= 3)
+
+    # Exact-match on the crafted file_id returns ONLY its own row.
+    items, total = await db.list_by_filter(collection, filter={'file_id': injection_fid})
+    assert total == 1
+    assert {it['id'] for it in items} == {'adv1'}
+
+    # A query_text packed with FT operators must not raise and must not match
+    # the benign row (escaped to literal terms, none of which it contains).
+    result = await db.search(
+        collection,
+        [0.0, 0.0, 0.0, 0.0],
+        k=5,
+        search_type='full_text',
+        query_text='@document:{*} | -()~ "evil"',
+    )
+    assert 'benign2' not in result['ids'][0]
+
+    # The brace/star-bearing file_id is encoded, so it round-trips: exact match
+    # returns ONLY its own row and never widens. No RequestError is raised.
+    b_items, b_total = await db.list_by_filter(collection, filter={'file_id': brace_fid})
+    assert b_total == 1
+    assert {it['id'] for it in b_items} == {'brace3'}
+
+    # And deletion by that file_id removes exactly its own row.
+    deleted = await db.delete_by_filter(collection, filter={'file_id': brace_fid})
+    assert deleted == 1