mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-14 08:26:07 +00:00
feat(vector): add Valkey Search vector database backend (#2276)
* feat(vector): add Valkey Search vector database backend Add a new opt-in VectorDatabase backend backed by the Valkey Search module (valkey/valkey-bundle), accessed via the official valkey-glide client's native ft command namespace. - Implements the full VectorDatabase ABC: VECTOR, FULL_TEXT and HYBRID search, all 8 metadata filter operators, and pagination with exact totals. - HYBRID uses filter-then-KNN (no app-side weighted fusion); vector_weight is accepted for interface parity but NOT honored (docstring + one-time warning + docs caveat). - Lazy connect so a down Valkey never blocks boot; mandatory client_name=langbot_vector_client; optional auth + TLS (never logged). - Registered via a single elif branch in vector/mgr.py; disabled by default (vdb.use stays chroma) for toC compatibility. - Adds valkey-glide>=2.4.1,<3.0.0; no protobuf/pydantic downgrade; no ORM change so no Alembic migration. - Unit tests (fast lane, no server) + slow-gated integration tests (TEST_VALKEY_URL, valkey/valkey-bundle:9.1.0) + integration doc. * fix(vector): paginate Valkey Search deletes and guard delete_by_filter Address self-review follow-ups for the Valkey Search VDB backend: - _search_keys now paginates through the full result set in batches of _DELETE_SCAN_BATCH instead of capping at a single hard-coded 10000-key page, so delete_by_file_id / delete_by_filter fully remove files and filters that match more than one page of chunks (no orphaned vectors). - Add unit regression tests for the delete_by_filter mass-deletion guard: a filter referencing only non-indexed fields must skip and return 0 (never fall back to match-all), and a supported filter still deletes matching keys. * refactor(vector): harden Valkey Search backend and add adversarial tests Address the self-review NICE-TO-HAVE items for the Valkey Search VDB backend: - Guard the username-without-password credential edge (skip auth + warn instead of building ServerCredentials(password=None, ...), which glide rejects). - Add an async close() teardown that closes the glide client and resets cached state (re-init is safe via the existing None guard). - Hoist 'import json' to module top (was imported inside three methods). - Document the FT TAG literal-brace limitation in _escape_tag (fails closed, never widens). Tests: - Add an adversarial-input integration test proving crafted file_id / query_text cannot break out of or widen a query (fail-closed on braces). - Add unit tests for close() and the credential-build guard. Signed-off-by: Daria Korenieva <daric2612@gmail.com> * fix(vector): make Valkey Search file_id TAG support arbitrary characters Valkey Search's FT TAG query parser cannot handle '{', '}' or '*' even when backslash-escaped, so a file_id containing those characters previously produced an unparseable query (it failed closed / raised). Percent-encode exactly those FT-unsafe characters (plus '%' for reversibility) in the file_id TAG value, applied identically at write time and query time, so an arbitrary file_id round-trips. For normal UUID/hash ids this is a no-op and the stored value is unchanged; the original file_id is always preserved verbatim in metadata_json. Strengthen the adversarial integration test to assert a brace/star-bearing file_id matches and deletes exactly its own row (no widening, no raise), and add unit tests for _encode_file_id and the filter encoding. Signed-off-by: Daria Korenieva <daric2612@gmail.com> * refactor(vector): address Valkey Search review feedback - Add configurable request_timeout (default 5000ms; glide default 250ms is too low for KNN); expose in config.yaml + docs table - Validate embedding dimension consistency in add_embeddings (fail fast on mixed lengths to avoid silent KNN corruption) - Use ft.info (O(1)) instead of ft.list (O(n)) for index existence checks in the query hot path; also closes the check-then-create TOCTOU window - Pipeline HSETs via a non-atomic Batch instead of N sequential awaits - Extract shared _iter_reply_docs to deduplicate reply parsing between _reply_to_chroma and list_by_filter - Parenthesize multi-condition pre-filters before the => KNN clause - Fail closed when a username is configured without a password - Catch only RequestError on ft.dropindex (let connection/auth errors surface) - Bound the delete_collection SCAN loop with a safety cap - Add VectorDatabase.close() (no-op default) + VectorDBManager.shutdown() - Simplify _MATCH_ALL literal; normalize typing to builtin generics * fix(vector/valkey_search): address round-2 review feedback - Serialize lazy client creation with an asyncio.Lock (double-checked) so concurrent first-use callers don't construct and leak duplicate clients. - Make the filter operator chain exhaustive: raise on an unhandled op rather than silently dropping the condition (which could widen delete_by_filter). - Cast numeric range (///) values to float, failing closed on non-numeric input and pre-empting a future NUMERIC-field injection surface. * refactor(vector): remove shutdown/close from base ABC per maintainer feedback Per maintainer request, interface changes to VectorDatabase ABC and VectorDBManager should be in a separate PR with implementation across all backends. The ValkeySearchVectorDatabase.close() method remains but does not override an ABC method. Signed-off-by: Daria Korenieva <daric2612@gmail.com> * docs(test): list valkey_search in vdb coverage exclusions Add valkey_search to the documented vector/vdbs/ coverage-exclusion list, matching the existing chroma/milvus/pgvector/qdrant/seekdb entries. These adapters require a live database instance and are covered by env-gated integration tests instead of unit tests. Signed-off-by: Daria Korenieva <daric2612@gmail.com> --------- Signed-off-by: Daria Korenieva <daric2612@gmail.com>
This commit is contained in:
Daria Korenieva
@@ -104,6 +104,17 @@ def create_minimal_config(tmpdir: Path, port: int = 15300) -> Path:
|
||||
'user': 'postgres',
|
||||
'password': 'postgres',
|
||||
},
|
||||
'valkey_search': {
|
||||
'host': 'localhost',
|
||||
'port': 6379,
|
||||
'db': 0,
|
||||
'password': '',
|
||||
'username': '',
|
||||
'tls': False,
|
||||
'index_algorithm': 'HNSW',
|
||||
'distance_metric': 'COSINE',
|
||||
'request_timeout': 5000,
|
||||
},
|
||||
},
|
||||
'storage': {
|
||||
'use': 'local',
|
||||
|
||||
@@ -0,0 +1,343 @@
|
||||
"""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,
|
||||
)
|
||||
from glide import ft
|
||||
|
||||
if not VALKEY_SEARCH_AVAILABLE:
|
||||
pytest.skip('valkey-glide not installed')
|
||||
|
||||
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
|
||||
@@ -27,7 +27,7 @@
|
||||
|
||||
### 4. 向量数据库 (`vector/vdbs/`)
|
||||
- **路径**: `src/langbot/pkg/vector/vdbs/`
|
||||
- **模块**: chroma, milvus, pgvector, qdrant, seekdb
|
||||
- **模块**: chroma, milvus, pgvector, qdrant, seekdb, valkey_search
|
||||
- **排除原因**: 需要真实向量数据库实例运行
|
||||
- **测试方式**: 需要 Docker 启动测试数据库或 mock
|
||||
- **状态**: 后续可补充 mock 测试
|
||||
|
||||
@@ -33,7 +33,7 @@ class TestVectorDBManagerInitialization:
|
||||
mocks['langbot.pkg.core.app'] = MagicMock()
|
||||
|
||||
# Mock all VDB backend implementations
|
||||
for backend in ['chroma', 'qdrant', 'seekdb', 'milvus', 'pgvector_db']:
|
||||
for backend in ['chroma', 'qdrant', 'seekdb', 'milvus', 'pgvector_db', 'valkey_search']:
|
||||
mocks[f'langbot.pkg.vector.vdbs.{backend}'] = MagicMock()
|
||||
|
||||
return mocks
|
||||
@@ -123,6 +123,25 @@ class TestVectorDBManagerInitialization:
|
||||
|
||||
mock_seekdb_class.assert_called_once_with(mock_app)
|
||||
|
||||
def test_initialize_valkey_search_backend(self):
|
||||
"""Valkey Search config uses ValkeySearchVectorDatabase backend."""
|
||||
vdb_config = {'use': 'valkey_search'}
|
||||
mock_app = self._create_mock_app(vdb_config)
|
||||
|
||||
mocks = self._make_vector_import_mocks()
|
||||
mock_valkey_class = MagicMock()
|
||||
mocks['langbot.pkg.vector.vdbs.valkey_search'].ValkeySearchVectorDatabase = mock_valkey_class
|
||||
|
||||
with isolated_sys_modules(mocks):
|
||||
from langbot.pkg.vector.mgr import VectorDBManager
|
||||
|
||||
mgr = VectorDBManager(mock_app)
|
||||
|
||||
import asyncio
|
||||
asyncio.get_event_loop().run_until_complete(mgr.initialize())
|
||||
|
||||
mock_valkey_class.assert_called_once_with(mock_app)
|
||||
|
||||
def test_initialize_milvus_backend_with_uri(self):
|
||||
"""Milvus config with custom URI."""
|
||||
vdb_config = {
|
||||
|
||||
@@ -0,0 +1,388 @@
|
||||
"""Unit tests for the Valkey Search VDB backend's pure helpers.
|
||||
|
||||
These tests exercise the filter-to-FT mapping, float32 packing, tag/text
|
||||
escaping, FT.SEARCH reply parsing and the import guard. They run in the fast
|
||||
CI lane and require NO running Valkey server.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import struct
|
||||
from importlib import import_module
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def get_valkey_module():
|
||||
"""Lazy import of the valkey_search backend module."""
|
||||
return import_module('langbot.pkg.vector.vdbs.valkey_search')
|
||||
|
||||
|
||||
def make_backend():
|
||||
"""Construct a backend instance without running its __init__.
|
||||
|
||||
The constructor needs a live ``ap`` + config; for pure-helper tests we
|
||||
only need a bare instance with the attributes the helpers touch.
|
||||
"""
|
||||
mod = get_valkey_module()
|
||||
backend = object.__new__(mod.ValkeySearchVectorDatabase)
|
||||
# _ensure_client serializes creation through this lock; set it here since
|
||||
# __init__ (which normally creates it) is bypassed.
|
||||
backend._client_lock = asyncio.Lock()
|
||||
return backend
|
||||
|
||||
|
||||
class TestFloat32Packing:
|
||||
"""Tests for _pack_vector little-endian float32 packing."""
|
||||
|
||||
def test_pack_round_trips(self):
|
||||
mod = get_valkey_module()
|
||||
vec = [0.1, -2.5, 3.0, 4.25]
|
||||
packed = mod.ValkeySearchVectorDatabase._pack_vector(vec)
|
||||
assert isinstance(packed, bytes)
|
||||
assert len(packed) == 4 * len(vec)
|
||||
unpacked = list(struct.unpack(f'<{len(vec)}f', packed))
|
||||
for original, restored in zip(vec, unpacked):
|
||||
assert restored == pytest.approx(original, rel=1e-6)
|
||||
|
||||
def test_pack_is_little_endian(self):
|
||||
mod = get_valkey_module()
|
||||
packed = mod.ValkeySearchVectorDatabase._pack_vector([1.0])
|
||||
assert packed == struct.pack('<f', 1.0)
|
||||
|
||||
|
||||
class TestTagEscaping:
|
||||
"""Tests for _escape_tag."""
|
||||
|
||||
def test_escapes_special_chars(self):
|
||||
mod = get_valkey_module()
|
||||
escaped = mod.ValkeySearchVectorDatabase._escape_tag('a-b c.d')
|
||||
assert '\\-' in escaped
|
||||
assert '\\ ' in escaped
|
||||
assert '\\.' in escaped
|
||||
|
||||
def test_plain_value_unchanged(self):
|
||||
mod = get_valkey_module()
|
||||
assert mod.ValkeySearchVectorDatabase._escape_tag('abc123') == 'abc123'
|
||||
|
||||
|
||||
class TestFileIdEncoding:
|
||||
"""Tests for _encode_file_id (FT-unsafe char percent-encoding)."""
|
||||
|
||||
def test_uuid_is_noop(self):
|
||||
mod = get_valkey_module()
|
||||
fid = '550e8400-e29b-41d4-a716-446655440000'
|
||||
assert mod.ValkeySearchVectorDatabase._encode_file_id(fid) == fid
|
||||
|
||||
def test_encodes_braces_star_and_percent(self):
|
||||
mod = get_valkey_module()
|
||||
enc = mod.ValkeySearchVectorDatabase._encode_file_id('a{b}c*d%e')
|
||||
# '{'=7B '}'=7D '*'=2A '%'=25
|
||||
assert enc == 'a%7Bb%7Dc%2Ad%25e'
|
||||
# No raw FT-unsafe char survives.
|
||||
assert all(ch not in enc for ch in '{}*') or '%' in enc
|
||||
|
||||
def test_encoding_is_deterministic_and_collision_safe(self):
|
||||
mod = get_valkey_module()
|
||||
enc = mod.ValkeySearchVectorDatabase._encode_file_id
|
||||
# A literal "%7B" must not collide with an encoded "{".
|
||||
assert enc('{') != enc('%7B')
|
||||
assert enc('{') == '%7B'
|
||||
assert enc('%7B') == '%257B'
|
||||
|
||||
def test_filter_encodes_unsafe_chars_in_tag_query(self):
|
||||
backend = make_backend()
|
||||
# The emitted TAG query must contain the encoded form, never raw braces.
|
||||
frag = backend._triples_to_ft({'file_id': 'x}y{z*'})
|
||||
assert '7D' in frag and '7B' in frag and '2A' in frag
|
||||
# No raw '*' from the value, and exactly one opening/closing brace (the
|
||||
# TAG-clause delimiters) — the value's own braces were encoded away.
|
||||
assert '*' not in frag
|
||||
assert frag.count('{') == 1 and frag.count('}') == 1
|
||||
assert frag.startswith('@file_id:{') and frag.endswith('}')
|
||||
|
||||
def test_filter_in_operator_encodes_each_value(self):
|
||||
backend = make_backend()
|
||||
frag = backend._triples_to_ft({'file_id': {'$in': ['a*b', 'c}d']}})
|
||||
assert '2A' in frag and '7D' in frag
|
||||
assert '*' not in frag
|
||||
|
||||
|
||||
class TestFilterToFt:
|
||||
"""Tests for _triples_to_ft filter mapping (all 8 operators)."""
|
||||
|
||||
def test_empty_filter_returns_empty_string(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft(None) == ''
|
||||
assert backend._triples_to_ft({}) == ''
|
||||
|
||||
def test_eq_tag(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft({'file_id': 'abc'}) == '@file_id:{abc}'
|
||||
|
||||
def test_explicit_eq_tag(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft({'file_id': {'$eq': 'abc'}}) == '@file_id:{abc}'
|
||||
|
||||
def test_ne_tag(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft({'file_id': {'$ne': 'abc'}}) == '-@file_id:{abc}'
|
||||
|
||||
def test_in_tag(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft({'file_id': {'$in': ['a', 'b']}}) == '@file_id:{a|b}'
|
||||
|
||||
def test_nin_tag(self):
|
||||
backend = make_backend()
|
||||
assert backend._triples_to_ft({'file_id': {'$nin': ['a', 'b']}}) == '-@file_id:{a|b}'
|
||||
|
||||
def test_numeric_range_operators(self):
|
||||
backend = make_backend()
|
||||
# file_id is the only indexed field; numeric ops still render via the
|
||||
# generic range fragment, so use file_id to keep the field supported.
|
||||
# Values are cast to float (defensive against non-numeric input and a
|
||||
# future NUMERIC field becoming an injection surface).
|
||||
assert backend._triples_to_ft({'file_id': {'$gt': 5}}) == '@file_id:[(5.0 +inf]'
|
||||
assert backend._triples_to_ft({'file_id': {'$gte': 5}}) == '@file_id:[5.0 +inf]'
|
||||
assert backend._triples_to_ft({'file_id': {'$lt': 5}}) == '@file_id:[-inf (5.0]'
|
||||
assert backend._triples_to_ft({'file_id': {'$lte': 5}}) == '@file_id:[-inf 5.0]'
|
||||
|
||||
def test_numeric_range_rejects_non_numeric(self):
|
||||
backend = make_backend()
|
||||
# A non-numeric range value fails closed rather than interpolating raw.
|
||||
with pytest.raises((ValueError, TypeError)):
|
||||
backend._triples_to_ft({'file_id': {'$gt': 'not-a-number'}})
|
||||
|
||||
def test_unsupported_field_dropped(self):
|
||||
backend = make_backend()
|
||||
# Non-indexed fields are dropped (returns empty expression).
|
||||
assert backend._triples_to_ft({'some_other_field': 'x'}) == ''
|
||||
|
||||
def test_multiple_supported_keys_anded(self):
|
||||
backend = make_backend()
|
||||
# Two conditions on the same indexed field are joined with a space (AND).
|
||||
result = backend._triples_to_ft({'file_id': {'$in': ['a', 'b']}})
|
||||
assert result == '@file_id:{a|b}'
|
||||
|
||||
|
||||
class TestTextEscaping:
|
||||
"""Tests for _escape_text full-text escaping."""
|
||||
|
||||
def test_escapes_ft_special_chars(self):
|
||||
mod = get_valkey_module()
|
||||
escaped = mod.ValkeySearchVectorDatabase._escape_text('hello@world|test')
|
||||
assert '\\@' in escaped
|
||||
assert '\\|' in escaped
|
||||
|
||||
|
||||
class TestReplyToChroma:
|
||||
"""Tests for _reply_to_chroma FT.SEARCH reply parsing."""
|
||||
|
||||
def test_parses_knn_reply(self):
|
||||
backend = make_backend()
|
||||
# glide returns [total, {key: {field: value}}]
|
||||
reply = [
|
||||
2,
|
||||
{
|
||||
b'kb:col1:id1': {
|
||||
b'distance': b'0.10',
|
||||
b'document': b'hello',
|
||||
b'metadata_json': b'{"file_id": "f1"}',
|
||||
},
|
||||
b'kb:col1:id2': {
|
||||
b'distance': b'0.25',
|
||||
b'document': b'world',
|
||||
b'metadata_json': b'{"file_id": "f2"}',
|
||||
},
|
||||
},
|
||||
]
|
||||
result = backend._reply_to_chroma('idx:col1', reply, has_distance=True)
|
||||
assert result['ids'][0] == ['id1', 'id2']
|
||||
assert result['distances'][0] == [pytest.approx(0.10), pytest.approx(0.25)]
|
||||
assert result['metadatas'][0][0] == {'file_id': 'f1'}
|
||||
assert result['metadatas'][0][1] == {'file_id': 'f2'}
|
||||
|
||||
def test_empty_reply(self):
|
||||
backend = make_backend()
|
||||
result = backend._reply_to_chroma('idx:col1', [0, {}], has_distance=True)
|
||||
assert result == {'ids': [[]], 'metadatas': [[]], 'distances': [[]]}
|
||||
|
||||
def test_malformed_reply(self):
|
||||
backend = make_backend()
|
||||
result = backend._reply_to_chroma('idx:col1', [], has_distance=True)
|
||||
assert result == {'ids': [[]], 'metadatas': [[]], 'distances': [[]]}
|
||||
|
||||
def test_text_search_reply_no_distance(self):
|
||||
backend = make_backend()
|
||||
reply = [
|
||||
1,
|
||||
{
|
||||
b'kb:col1:id1': {
|
||||
b'document': b'hello',
|
||||
b'metadata_json': b'{"file_id": "f1"}',
|
||||
},
|
||||
},
|
||||
]
|
||||
result = backend._reply_to_chroma('idx:col1', reply, has_distance=False)
|
||||
assert result['ids'][0] == ['id1']
|
||||
assert result['distances'][0] == [0.0]
|
||||
|
||||
|
||||
class TestImportGuard:
|
||||
"""Tests for the ImportError guard when glide is unavailable."""
|
||||
|
||||
def test_constructor_raises_when_unavailable(self, monkeypatch):
|
||||
mod = get_valkey_module()
|
||||
monkeypatch.setattr(mod, 'VALKEY_SEARCH_AVAILABLE', False)
|
||||
with pytest.raises(ImportError, match='valkey-glide'):
|
||||
mod.ValkeySearchVectorDatabase(ap=None)
|
||||
|
||||
|
||||
class TestSupportedSearchTypes:
|
||||
"""Tests for supported_search_types."""
|
||||
|
||||
def test_supports_vector_full_text_hybrid(self):
|
||||
mod = get_valkey_module()
|
||||
from langbot.pkg.vector.vdb import SearchType
|
||||
|
||||
types = mod.ValkeySearchVectorDatabase.supported_search_types()
|
||||
assert SearchType.VECTOR in types
|
||||
assert SearchType.FULL_TEXT in types
|
||||
assert SearchType.HYBRID in types
|
||||
|
||||
|
||||
class TestDeleteByFilterGuard:
|
||||
"""Regression tests for the delete_by_filter mass-deletion guard.
|
||||
|
||||
A non-empty filter referencing only non-indexed fields must NOT fall back
|
||||
to match-all and wipe the whole collection: it must skip and return 0.
|
||||
"""
|
||||
|
||||
async def test_unsupported_only_filter_skips_and_returns_zero(self):
|
||||
backend = make_backend()
|
||||
# Make the client/index lookups succeed without a real server.
|
||||
backend._client = AsyncMock()
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensure_client = AsyncMock(return_value=backend._client)
|
||||
backend._index_exists = AsyncMock(return_value=True)
|
||||
# _search_keys must never be reached for an unusable filter.
|
||||
backend._search_keys = AsyncMock(
|
||||
side_effect=AssertionError('_search_keys must not be called for an unusable filter')
|
||||
)
|
||||
|
||||
# Filter references only a non-indexed field -> maps to no FT conditions.
|
||||
deleted = await backend.delete_by_filter('col1', {'some_other_field': 'x'})
|
||||
|
||||
assert deleted == 0
|
||||
backend._client.delete.assert_not_called()
|
||||
|
||||
async def test_supported_filter_deletes_matching_keys(self):
|
||||
backend = make_backend()
|
||||
backend._client = AsyncMock()
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensure_client = AsyncMock(return_value=backend._client)
|
||||
backend._index_exists = AsyncMock(return_value=True)
|
||||
backend._search_keys = AsyncMock(return_value=['kb:col1:id1', 'kb:col1:id2'])
|
||||
|
||||
deleted = await backend.delete_by_filter('col1', {'file_id': 'f1'})
|
||||
|
||||
assert deleted == 2
|
||||
backend._client.delete.assert_awaited_once_with(['kb:col1:id1', 'kb:col1:id2'])
|
||||
|
||||
|
||||
class TestClose:
|
||||
"""Tests for the close() teardown."""
|
||||
|
||||
async def test_close_resets_client_and_indexes(self):
|
||||
backend = make_backend()
|
||||
client = AsyncMock()
|
||||
backend._client = client
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensured_indexes = {'idx:col1'}
|
||||
|
||||
await backend.close()
|
||||
|
||||
client.close.assert_awaited_once()
|
||||
assert backend._client is None
|
||||
assert backend._ensured_indexes == set()
|
||||
|
||||
async def test_close_is_noop_when_no_client(self):
|
||||
backend = make_backend()
|
||||
backend._client = None
|
||||
backend.ap = type('Ap', (), {'logger': AsyncMock()})()
|
||||
backend._ensured_indexes = set()
|
||||
# Should not raise.
|
||||
await backend.close()
|
||||
assert backend._client is None
|
||||
|
||||
|
||||
class TestCredentialsBuild:
|
||||
"""Tests for the auth-credential construction in _ensure_client."""
|
||||
|
||||
def _prep_backend(self, mod, monkeypatch, *, username, password):
|
||||
backend = make_backend()
|
||||
backend._client = None
|
||||
backend._host = 'localhost'
|
||||
backend._port = 6379
|
||||
backend._db = 0
|
||||
backend._tls = False
|
||||
backend._username = username
|
||||
backend._password = password
|
||||
backend._request_timeout = 5000
|
||||
backend._ensured_indexes = set()
|
||||
warnings: list[str] = []
|
||||
backend.ap = type(
|
||||
'Ap',
|
||||
(),
|
||||
{
|
||||
'logger': type(
|
||||
'L', (), {'info': lambda self, *a, **k: None, 'warning': lambda s, m, *a, **k: warnings.append(m)}
|
||||
)()
|
||||
},
|
||||
)()
|
||||
|
||||
created = {}
|
||||
|
||||
class _FakeClient:
|
||||
@staticmethod
|
||||
async def create(conf):
|
||||
created['conf'] = conf
|
||||
return AsyncMock()
|
||||
|
||||
cred_calls: list[dict] = []
|
||||
|
||||
def _fake_credentials(**kwargs):
|
||||
cred_calls.append(kwargs)
|
||||
return ('CRED', kwargs)
|
||||
|
||||
monkeypatch.setattr(mod, 'GlideClient', _FakeClient)
|
||||
monkeypatch.setattr(mod, 'ServerCredentials', _fake_credentials)
|
||||
monkeypatch.setattr(mod, 'GlideClientConfiguration', lambda **kw: kw)
|
||||
monkeypatch.setattr(mod, 'NodeAddress', lambda *a, **k: ('node', a, k))
|
||||
return backend, created, cred_calls, warnings
|
||||
|
||||
async def test_username_without_password_fails_closed(self, monkeypatch):
|
||||
mod = get_valkey_module()
|
||||
backend, created, cred_calls, warnings = self._prep_backend(mod, monkeypatch, username='acluser', password=None)
|
||||
|
||||
# A username without a password must fail closed rather than silently
|
||||
# connecting unauthenticated to a (potentially shared) Valkey instance.
|
||||
with pytest.raises(ValueError, match='without a password'):
|
||||
await backend._ensure_client()
|
||||
|
||||
assert cred_calls == [] # ServerCredentials NOT constructed
|
||||
assert 'conf' not in created # client never created
|
||||
|
||||
async def test_password_builds_credentials(self, monkeypatch):
|
||||
mod = get_valkey_module()
|
||||
backend, created, cred_calls, warnings = self._prep_backend(
|
||||
mod, monkeypatch, username='acluser', password='secret'
|
||||
)
|
||||
|
||||
await backend._ensure_client()
|
||||
|
||||
assert len(cred_calls) == 1
|
||||
assert cred_calls[0] == {'password': 'secret', 'username': 'acluser'}
|
||||
assert created['conf']['credentials'] == ('CRED', {'password': 'secret', 'username': 'acluser'})
|
||||
Reference in New Issue
Block a user