diff --git a/docs/VALKEY_SEARCH_INTEGRATION.md b/docs/VALKEY_SEARCH_INTEGRATION.md new file mode 100644 index 000000000..aa95b0e70 --- /dev/null +++ b/docs/VALKEY_SEARCH_INTEGRATION.md @@ -0,0 +1,169 @@ +# Valkey Search Vector Database Integration + +This document describes how to use **Valkey Search** (the search/vector module bundled in +`valkey/valkey-bundle`) as the vector database backend for LangBot's knowledge base (RAG) +feature. + +## What is Valkey Search? + +**Valkey Search** is a module that adds vector similarity search and full-text search to +[Valkey](https://valkey.io/), the open-source, BSD-licensed in-memory data store forked from +Redis OSS. It is distributed in the `valkey/valkey-bundle` image alongside other modules +(JSON, Bloom, LDAP). + +LangBot talks to Valkey through the official [`valkey-glide`](https://pypi.org/project/valkey-glide/) +client (Rust core + async Python wrapper), using its native `ft` (search) command namespace. + +### Key Features + +- **Vector search**: ANN via HNSW or exact via FLAT, with COSINE / L2 / IP distance metrics +- **Full-text search**: term, prefix and phrase matching over indexed text fields +- **Hybrid search**: a metadata/text filter pre-selects candidates, then KNN ranks them +- **In-memory speed**: vectors and documents are stored as Valkey HASH keys +- **Auth + TLS**: optional username/password and TLS for production (toB / SaaS) deployments + +### Licensing + +- Valkey core and the Search module are **BSD-3-Clause**. +- The `valkey-glide` client is **Apache-2.0**. + +Both are compatible with LangBot. + +## Installation + +Valkey Search support is included when you install LangBot — the `valkey-glide` dependency is +declared in `pyproject.toml`. To install manually: + +```bash +pip install 'valkey-glide>=2.4.1,<3.0.0' +``` + +You also need a running Valkey server with the Search module loaded. The simplest way is the +bundled image: + +```bash +# Run valkey-bundle (includes the Search module) on host port 6380 +podman run -d --name valkey-test-langbot -p 6380:6379 valkey/valkey-bundle:9.1.0 +# (docker run ... works identically) +``` + +`valkey-bundle` ships multi-arch images (linux/amd64 + linux/arm64), so it runs on both CI +(x86_64) and Apple-silicon dev machines. + +## Configuration + +Valkey Search is **opt-in and disabled by default** — the default `vdb.use` stays `chroma`, +so existing single-process deployments are unaffected. To enable it, edit your `config.yaml`: + +```yaml +vdb: + use: valkey_search + valkey_search: + host: 'localhost' + port: 6379 # use 6380 if you started the container as shown above + db: 0 + password: '' # optional (ACL / requirepass) — never logged + username: '' # optional (ACL user) + tls: false # optional (toB / SaaS) + index_algorithm: 'HNSW' # HNSW | FLAT + distance_metric: 'COSINE' # COSINE | L2 | IP + request_timeout: 5000 # per-request timeout in ms +``` + +| Option | Default | Description | +|--------|---------|-------------| +| `host` | `localhost` | Valkey host | +| `port` | `6379` | Valkey port | +| `db` | `0` | Logical database id | +| `password` | `''` | Optional auth password (empty = no auth). Never logged. | +| `username` | `''` | Optional ACL username. Configuring a username without a password fails closed (raises) rather than connecting unauthenticated. | +| `tls` | `false` | Enable TLS for the connection | +| `index_algorithm` | `HNSW` | `HNSW` (approximate) or `FLAT` (exact) | +| `distance_metric` | `COSINE` | `COSINE`, `L2`, or `IP` | +| `request_timeout` | `5000` | Per-request timeout in milliseconds. The valkey-glide default (250ms) is too low for vector KNN under load; raise it further for remote/cross-AZ Valkey. | + +### Connection behavior + +The backend uses a **lazy** connection (`lazy_connect=True`): the client is created on first +use and the connection is deferred to the first command. A misconfigured or unreachable Valkey +server therefore does **not** block LangBot from booting — knowledge-base operations will error +at call time instead, and you can recover by switching `vdb.use` back to another backend. + +The connection sets a fixed `client_name` of `langbot_vector_client` so it is identifiable in +`CLIENT LIST` and monitoring dashboards. + +## Supported search types + +| Type | Behavior | +|------|----------| +| `vector` | Pure KNN over the embedding field | +| `full_text` | Term/phrase match over the indexed `document` text field | +| `hybrid` | Metadata/text filter **pre-selects** candidates, then KNN ranks them | + +### ⚠️ Important: `vector_weight` is NOT honored + +Valkey Search hybrid queries follow a **filter-then-KNN** model: the filter (and/or full-text +clause) narrows the candidate set, and the KNN stage ranks the survivors by vector distance. +There is **no native weighted score fusion** (unlike, e.g., SeekDB's RRF boost). + +For interface compatibility the backend still accepts a `vector_weight` argument, but it is +**ignored** — passing different weights does not change result ordering. The first time a +non-default weight is supplied, the backend logs a one-time warning. + +If weighted hybrid ranking is needed in the future, it can be added **application-side** (run +vector KNN and full-text search separately and blend the scores). That is intentionally out of +scope for this integration. + +## Metadata & filtering + +Documents are stored as Valkey HASH keys under the prefix `kb:{collection}:{id}` with fields: + +- `vector` — the embedding, packed as little-endian FLOAT32 +- `document` — the raw text (indexed as TEXT for full-text/hybrid search) +- `file_id` — promoted to an indexed TAG field so it is filterable +- `metadata_json` — the full metadata dict, preserved verbatim as JSON + +Only **indexed** fields are filterable. Currently that is `file_id`. Filters referencing +non-indexed metadata keys are dropped with a warning (the same pragmatism used by the Milvus +and pgvector backends). All other metadata still round-trips intact via `metadata_json`. + +Supported filter operators (canonical Chroma-style `where` syntax): `$eq`, `$ne`, `$gt`, +`$gte`, `$lt`, `$lte`, `$in`, `$nin`. Multiple top-level keys are AND-ed. + +## Testing + +Unit tests (filter mapping, float32 packing, reply parsing, import guard) run in the fast lane +with no server: + +```bash +uv run pytest tests/unit_tests/vector/test_valkey_search_filter.py -q +``` + +Integration tests are **slow-gated** on `TEST_VALKEY_URL` and require a running server: + +```bash +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, matching the existing +PostgreSQL migration-test precedent. + +## Troubleshooting + +| Symptom | Cause / fix | +|---------|-------------| +| Tests skip with "Valkey Search module not available" | The server is plain Valkey without the Search module. Use the `valkey/valkey-bundle` image. | +| `ConnectionError` at call time | Check `host`/`port`/auth; remember `lazy_connect` defers errors to first use. | +| Empty search results right after insert | The Search indexer is asynchronous; results become visible within a short delay. The integration tests poll/retry to account for this. | +| Hybrid ranking ignores `vector_weight` | Expected — see the caveat above. | + +## Production considerations + +- **Cluster mode**: Valkey Search in cluster mode uses an additional coordination port. This + integration targets standalone mode; cluster support is a future consideration. +- **Persistence**: configure Valkey RDB/AOF persistence if the knowledge base must survive + restarts; otherwise an in-memory store is ephemeral. +- **Security**: set `password`/`username` and `tls: true` for any non-local deployment. + Credentials are never written to logs. diff --git a/pyproject.toml b/pyproject.toml index bece64d8d..3204bd35d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,7 @@ dependencies = [ "pgvector>=0.4.1", "botocore>=1.42.39", "litellm>=1.0.0", + "valkey-glide>=2.4.1,<3.0.0", ] keywords = [ "bot", diff --git a/src/langbot/pkg/vector/mgr.py b/src/langbot/pkg/vector/mgr.py index 73ed5d747..765c259f9 100644 --- a/src/langbot/pkg/vector/mgr.py +++ b/src/langbot/pkg/vector/mgr.py @@ -33,6 +33,12 @@ class VectorDBManager: self.vector_db = SeekDBVectorDatabase(self.ap) self.ap.logger.info('Initialized SeekDB vector database backend.') + elif vdb_type == 'valkey_search': + from .vdbs.valkey_search import ValkeySearchVectorDatabase + + self.vector_db = ValkeySearchVectorDatabase(self.ap) + self.ap.logger.info('Initialized Valkey Search vector database backend.') + elif vdb_type == 'milvus': from .vdbs.milvus import MilvusVectorDatabase diff --git a/src/langbot/pkg/vector/vdbs/valkey_search.py b/src/langbot/pkg/vector/vdbs/valkey_search.py new file mode 100644 index 000000000..9f92acc97 --- /dev/null +++ b/src/langbot/pkg/vector/vdbs/valkey_search.py @@ -0,0 +1,828 @@ +from __future__ import annotations + +import asyncio +import json +import struct +from typing import Any + +from langbot.pkg.core import app +from langbot.pkg.vector.vdb import VectorDatabase, SearchType +from langbot.pkg.vector.filter_utils import normalize_filter, strip_unsupported_fields + +try: + from glide import ( + Batch, + GlideClient, + GlideClientConfiguration, + NodeAddress, + RequestError, + ServerCredentials, + ft, + VectorField, + VectorFieldAttributesHnsw, + VectorFieldAttributesFlat, + VectorAlgorithm, + VectorType, + DistanceMetricType, + TagField, + TextField, + FtCreateOptions, + DataType, + FtSearchOptions, + FtSearchLimit, + ReturnField, + ) + + VALKEY_SEARCH_AVAILABLE = True +except ImportError: + VALKEY_SEARCH_AVAILABLE = False + +# Default per-request timeout (ms) for the glide client. The glide library +# default is 250ms, which is too low for vector KNN (``FT.SEARCH ... =>[KNN]``) +# under moderate load or with large indexes and yields spurious TimeoutErrors. +# Overridable via the ``vdb.valkey_search.request_timeout`` config option. +_DEFAULT_REQUEST_TIMEOUT_MS = 5000 + +# Safety cap on the number of SCAN rounds when purging a collection's keys, so +# a cursor-handling bug or pathological keyspace can never spin forever. +_MAX_SCAN_ROUNDS = 100000 + + +# Mandatory client name for production observability (CLIENT LIST / dashboards). +VALKEY_CLIENT_NAME = 'langbot_vector_client' + +# Fixed, indexed metadata schema. LangBot's RAG layer stores ``file_id`` on +# every chunk; it is the only metadata field we promote to a first-class +# (filterable) index field. All other metadata is preserved verbatim inside +# the ``metadata_json`` field so it survives a round-trip, but is NOT +# filterable (the established Milvus / pgvector pragmatism). +_INDEXED_TAG_FIELDS = {'file_id'} +_SUPPORTED_FILTER_FIELDS = set(_INDEXED_TAG_FIELDS) + +# Hash field names used for stored documents. +_FIELD_VECTOR = 'vector' +_FIELD_DOCUMENT = 'document' +_FIELD_FILE_ID = 'file_id' +_FIELD_METADATA = 'metadata_json' +_VEC_SCORE_ALIAS = '__vec_score' + +# Valkey Search has no bare "match everything" token for non-vector queries +# (a standalone ``*`` is a syntax error). A negated match on a sentinel tag +# value that can never exist matches every key, which is the canonical +# match-all idiom for FT.SEARCH. +_MATCH_ALL = '-@file_id:{__langbot_match_all_sentinel__}' + +# Page size used when enumerating matching keys for deletion. Deletes +# paginate through the full result set in batches of this size so that +# files/filters matching more than one page of chunks are fully removed +# (no silent truncation / orphaned vectors). +_DELETE_SCAN_BATCH = 10000 + +# Characters Valkey Search's TAG query parser cannot handle even when +# backslash-escaped (the brace delimiters and the wildcard). file_id TAG +# values are percent-encoded over this set (plus '%' itself, so the encoding +# is reversible/unambiguous) before being stored or queried, so an arbitrary +# file_id round-trips instead of producing an unparseable query. For normal +# UUID/hash file_ids none of these characters occur, so the encoding is a +# no-op and the stored value is unchanged. The original file_id is always +# preserved verbatim inside ``metadata_json``. +_FT_UNSAFE_TAG_CHARS = frozenset('{}*%') + + +class ValkeySearchVectorDatabase(VectorDatabase): + """Valkey Search (valkey-bundle) vector database adapter for LangBot. + + Backed by the Valkey Search module shipped in ``valkey/valkey-bundle``, + accessed through the official ``valkey-glide`` client's native ``ft`` + (search) command namespace. Documents are stored as Valkey HASH keys + under a per-collection prefix and indexed by one ``FT.CREATE`` index per + collection. + + Supported search types: ``VECTOR``, ``FULL_TEXT`` and ``HYBRID``. + + Hybrid search semantics (IMPORTANT) + ----------------------------------- + Valkey Search hybrid queries follow a *filter-then-KNN* model: the text / + metadata filter pre-selects candidate keys and the KNN stage ranks them by + vector distance. This backend does **NOT** implement application-side + weighted score fusion. The ``vector_weight`` argument is therefore + accepted for interface compatibility but is **not honored** — passing + different weights does not change result ordering. A one-time warning is + emitted the first time a non-default weight is supplied. App-side score + fusion can be layered on later if weighted hybrid ranking is required. + """ + + @classmethod + def supported_search_types(cls) -> list[SearchType]: + return [SearchType.VECTOR, SearchType.FULL_TEXT, SearchType.HYBRID] + + def __init__(self, ap: app.Application): + if not VALKEY_SEARCH_AVAILABLE: + raise ImportError( + "valkey-glide is not installed. Install it with: pip install 'valkey-glide>=2.4.1,<3.0.0'" + ) + + self.ap = ap + config = self.ap.instance_config.data['vdb']['valkey_search'] + + self._host = config.get('host', 'localhost') + self._port = int(config.get('port', 6379)) + self._db = int(config.get('db', 0)) + # Auth / TLS are optional (toB / SaaS). Never logged. + self._password = config.get('password', '') or None + self._username = config.get('username', '') or None + self._tls = bool(config.get('tls', False)) + self._request_timeout = int(config.get('request_timeout', _DEFAULT_REQUEST_TIMEOUT_MS)) + + algorithm = str(config.get('index_algorithm', 'HNSW')).upper() + self._algorithm = VectorAlgorithm.FLAT if algorithm == 'FLAT' else VectorAlgorithm.HNSW + + metric = str(config.get('distance_metric', 'COSINE')).upper() + self._distance_metric = { + 'COSINE': DistanceMetricType.COSINE, + 'L2': DistanceMetricType.L2, + 'IP': DistanceMetricType.IP, + }.get(metric, DistanceMetricType.COSINE) + + # Lazily-created client (created on first use so a down Valkey does not + # block LangBot boot). + self._client: GlideClient | None = None + # Serializes lazy client creation so concurrent first-use callers do not + # each construct (and leak) a separate GlideClient. + self._client_lock = asyncio.Lock() + # Index names we have already ensured this process lifetime. + self._ensured_indexes: set[str] = set() + # Whether we have already warned about the non-honored vector_weight. + self._vector_weight_warned = False + + # ------------------------------------------------------------------ # + # Client lifecycle + # ------------------------------------------------------------------ # + async def _ensure_client(self) -> GlideClient: + """Create the glide client on first use (lazy, non-blocking boot).""" + if self._client is not None: + return self._client + # Double-checked locking: serialize creation so two concurrent + # first-use callers don't both build a client and leak one. + async with self._client_lock: + if self._client is not None: + return self._client + + credentials = None + if self._password is not None: + # username is optional alongside a password (ACL "user" vs default user). + credentials = ServerCredentials(password=self._password, username=self._username) + elif self._username is not None: + # A username without a password is not a valid credential pair, and silently + # connecting unauthenticated to a potentially shared Valkey instance is a + # security footgun (e.g. an env var that failed to resolve). Fail closed. + raise ValueError( + 'Valkey Search: a username was configured without a password. ' + 'Set both username and password to use ACL authentication, or remove both.' + ) + + conf = GlideClientConfiguration( + addresses=[NodeAddress(self._host, self._port)], + client_name=VALKEY_CLIENT_NAME, + database_id=self._db, + use_tls=self._tls, + lazy_connect=True, + credentials=credentials, + request_timeout=self._request_timeout, + ) + self._client = await GlideClient.create(conf) + self.ap.logger.info( + f'Initialized Valkey Search client to {self._host}:{self._port} (db={self._db}, tls={self._tls})' + ) + return self._client + + async def close(self) -> None: + """Close the glide client and reset state. + + Safe to call when no client was created. After ``close`` the next + operation transparently re-creates the client (``_ensure_client`` + guards on ``self._client is None``). + """ + if self._client is not None: + try: + await self._client.close() + except Exception: + self.ap.logger.warning('Valkey Search: error while closing client (ignored)') + finally: + self._client = None + self._ensured_indexes.clear() + + # ------------------------------------------------------------------ # + # Naming helpers + # ------------------------------------------------------------------ # + @staticmethod + def _index_name(collection: str) -> str: + return f'idx:{collection}' + + @staticmethod + def _key_prefix(collection: str) -> str: + return f'kb:{collection}:' + + @staticmethod + def _pack_vector(vec: list[float]) -> bytes: + """Pack a float vector into little-endian float32 bytes. + + Valkey Search stores and queries vectors as FLOAT32 little-endian + blobs (per the search query-language spec). + """ + return struct.pack(f'<{len(vec)}f', *[float(x) for x in vec]) + + @staticmethod + def _escape_tag(value: str) -> str: + """Escape characters that are special inside a TAG ``{...}`` clause. + + The backslash is escaped first so it cannot consume a following + escape. This neutralises injection-style values (quotes, parens, + ``|``, ``@``, ``:``, spaces, dashes) so a crafted ``file_id`` cannot + break out of the clause. + + Note: Valkey Search's TAG query parser cannot handle a literal brace + (``{`` / ``}``) or ``*`` even when backslash-escaped. Callers that pass + a ``file_id`` route it through ``_encode_and_escape_tag`` / + ``_encode_file_id`` first, which percent-encodes exactly those + characters, so an arbitrary ``file_id`` round-trips safely. This raw + escaper is still correct for all other special characters. + """ + out = [] + for ch in str(value): + if ch in '\\,.<>{}[]"\':;!@#$%^&*()-+=~| ': + out.append('\\') + out.append(ch) + return ''.join(out) + + @staticmethod + def _encode_file_id(value: str) -> str: + """Make a ``file_id`` safe to use as an FT TAG token AND query value. + + Percent-encodes the characters Valkey Search's TAG parser cannot handle + even when backslash-escaped (``{``, ``}``, ``*``) plus ``%`` itself for + reversibility. Applied identically at write time (the stored TAG field) + and query time (filters / ``delete_by_file_id``) so any value matches + itself. For normal UUID/hash ids none of these characters occur, so + this is a no-op. The original value is always kept verbatim in + ``metadata_json``; this encoded form is only ever used for the indexed + TAG. + """ + out = [] + for ch in str(value): + if ch in _FT_UNSAFE_TAG_CHARS: + out.append('%{:02X}'.format(ord(ch))) + else: + out.append(ch) + return ''.join(out) + + def _encode_and_escape_tag(self, value: str) -> str: + """Encode an FT-unsafe ``file_id`` then escape TAG special chars.""" + return self._escape_tag(self._encode_file_id(value)) + + # ------------------------------------------------------------------ # + # Filter mapping (canonical triples -> FT query fragment) + # ------------------------------------------------------------------ # + def _triples_to_ft(self, filter: dict[str, Any] | None) -> str: + """Translate a canonical filter dict into an FT filter expression. + + Only indexed fields (``file_id``) are filterable; unsupported fields + are dropped with a warning (matching the Milvus / pgvector pattern). + Returns an empty string when there is no usable filter. + """ + triples = normalize_filter(filter) + if not triples: + return '' + triples = strip_unsupported_fields(triples, _SUPPORTED_FILTER_FIELDS) + + fragments: list[str] = [] + for field, op, value in triples: + # All currently-indexed fields are TAG fields; file_id values are + # encoded (FT-unsafe chars) then escaped so any value round-trips. + if op == '$eq': + fragments.append(f'@{field}:{{{self._encode_and_escape_tag(value)}}}') + elif op == '$ne': + fragments.append(f'-@{field}:{{{self._encode_and_escape_tag(value)}}}') + elif op == '$in': + joined = '|'.join(self._encode_and_escape_tag(v) for v in value) + fragments.append(f'@{field}:{{{joined}}}') + elif op == '$nin': + joined = '|'.join(self._encode_and_escape_tag(v) for v in value) + fragments.append(f'-@{field}:{{{joined}}}') + elif op == '$gt': + fragments.append(f'@{field}:[({float(value)} +inf]') + elif op == '$gte': + fragments.append(f'@{field}:[{float(value)} +inf]') + elif op == '$lt': + fragments.append(f'@{field}:[-inf ({float(value)}]') + elif op == '$lte': + fragments.append(f'@{field}:[-inf {float(value)}]') + else: + # normalize_filter() already rejects unknown operators, so this + # only triggers if SUPPORTED_OPS grows without this chain being + # updated. Fail closed (rather than silently dropping the + # condition, which would widen delete_by_filter's match set). + raise ValueError(f'Valkey Search: unhandled filter operator {op!r} on field {field!r}') + + return ' '.join(fragments) + + @staticmethod + def _build_text_clause(text: str) -> str: + """Build a field-scoped full-text clause for the ``document`` field. + + Each whitespace-delimited word becomes a ``@document:`` term and + the terms are AND-ed (space separated). FT special characters in each + term are escaped. Returns an empty string when *text* has no words. + """ + words = [w for w in str(text).split() if w] + if not words: + return '' + terms = [f'@{_FIELD_DOCUMENT}:{ValkeySearchVectorDatabase._escape_text(w)}' for w in words] + return ' '.join(terms) + + @staticmethod + def _escape_text(text: str) -> str: + """Escape FT full-text special characters in a single term.""" + out = [] + for ch in str(text): + if ch in '@!{}[]()|-"~*:\\': + out.append('\\') + out.append(ch) + return ''.join(out) + + # ------------------------------------------------------------------ # + # Index management + # ------------------------------------------------------------------ # + async def _ensure_index(self, client: GlideClient, collection: str, dim: int) -> None: + index = self._index_name(collection) + if index in self._ensured_indexes: + return + + # ft.info is O(1) and raises RequestError when the index is absent — + # cheaper than ft.list (O(n) over all indexes) and it closes the + # check-then-create TOCTOU window. + try: + await ft.info(client, index) + self._ensured_indexes.add(index) + return + except RequestError: + pass + + if self._algorithm == VectorAlgorithm.FLAT: + vector_attrs = VectorFieldAttributesFlat( + dimensions=dim, + distance_metric=self._distance_metric, + type=VectorType.FLOAT32, + ) + else: + vector_attrs = VectorFieldAttributesHnsw( + dimensions=dim, + distance_metric=self._distance_metric, + type=VectorType.FLOAT32, + ) + + schema = [ + VectorField(name=_FIELD_VECTOR, algorithm=self._algorithm, attributes=vector_attrs), + TagField(name=_FIELD_FILE_ID), + TextField(name=_FIELD_DOCUMENT), + ] + options = FtCreateOptions(data_type=DataType.HASH, prefixes=[self._key_prefix(collection)]) + await ft.create(client, index, schema, options) + self._ensured_indexes.add(index) + self.ap.logger.info( + f"Valkey Search index '{index}' created (dim={dim}, algo={self._algorithm.value}, " + f'metric={self._distance_metric.value})' + ) + + @staticmethod + def _decode(value: Any) -> str: + if isinstance(value, (bytes, bytearray, memoryview)): + return bytes(value).decode('utf-8', errors='replace') + return str(value) + + # ------------------------------------------------------------------ # + # VectorDatabase ABC implementation + # ------------------------------------------------------------------ # + async def get_or_create_collection(self, collection: str): + """Ensure a client exists. + + The index itself requires the vector dimension, which is only known at + first ``add_embeddings`` (same constraint as Qdrant / SeekDB), so this + is a best-effort no-op when the index does not yet exist. + """ + await self._ensure_client() + + async def add_embeddings( + self, + collection: str, + ids: list[str], + embeddings_list: list[list[float]], + metadatas: list[dict[str, Any]], + documents: list[str] | None = None, + ) -> None: + if not embeddings_list: + return + + client = await self._ensure_client() + dim = len(embeddings_list[0]) + # The index schema is fixed to the first embedding's dimension. A later + # embedding of a different length would be packed into a wrong-sized + # blob that Valkey stores silently but that yields garbage KNN + # distances, so reject mixed dimensions up-front. + if any(len(e) != dim for e in embeddings_list[1:]): + raise ValueError(f'All embeddings must have dimension {dim}; got mixed lengths') + await self._ensure_index(client, collection, dim) + + prefix = self._key_prefix(collection) + + batch = Batch(is_atomic=False) + for i, _id in enumerate(ids): + key = prefix + str(_id) + metadata = metadatas[i] if i < len(metadatas) else {} + mapping: dict[str, Any] = { + _FIELD_VECTOR: self._pack_vector(embeddings_list[i]), + _FIELD_METADATA: json.dumps(metadata, ensure_ascii=False), + } + file_id = metadata.get('file_id') + if file_id is not None: + mapping[_FIELD_FILE_ID] = self._encode_file_id(str(file_id)) + if documents is not None and i < len(documents) and documents[i] is not None: + mapping[_FIELD_DOCUMENT] = documents[i] + + batch.hset(key, mapping) + + # Pipeline all HSETs into a single round-trip (non-atomic) instead of + # one await per embedding, which is N sequential round-trips for N + # chunks. + await client.exec(batch, raise_on_error=True) + + self.ap.logger.info(f"Added {len(ids)} embeddings to Valkey Search collection '{collection}'") + + async def search( + self, + collection: str, + query_embedding: list[float], + k: int = 5, + search_type: str = 'vector', + query_text: str = '', + filter: dict[str, Any] | None = None, + vector_weight: float | None = None, + ) -> dict[str, Any]: + client = await self._ensure_client() + index = self._index_name(collection) + + if not await self._index_exists(client, index): + return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]} + + # vector_weight is accepted for interface parity but NOT honored by this + # backend (filter-then-KNN, no weighted fusion). Warn once. + if vector_weight is not None and not self._vector_weight_warned: + self.ap.logger.warning( + 'Valkey Search backend does not honor vector_weight: hybrid search uses ' + 'filter-then-KNN without weighted score fusion. The vector_weight value ' + 'is ignored. See docs/VALKEY_SEARCH_INTEGRATION.md.' + ) + self._vector_weight_warned = True + + filter_expr = self._triples_to_ft(filter) + + if search_type == SearchType.FULL_TEXT: + if not query_text: + return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]} + text_clause = self._build_text_clause(query_text) + if not text_clause: + return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]} + query = f'{filter_expr} {text_clause}'.strip() if filter_expr else text_clause + return await self._run_text_search(client, index, query, k) + + if search_type == SearchType.HYBRID: + # Filter / text pre-selects candidates; KNN ranks. No fusion. + pre = filter_expr + if query_text: + text_clause = self._build_text_clause(query_text) + if text_clause: + pre = f'{pre} {text_clause}'.strip() if pre else text_clause + pre = pre or '*' + query = f'{self._wrap_pre(pre)}=>[KNN {k} @{_FIELD_VECTOR} $BLOB AS {_VEC_SCORE_ALIAS}]' + return await self._run_knn_search(client, index, query, query_embedding, k) + + # Default: pure VECTOR search. + pre = filter_expr or '*' + query = f'{self._wrap_pre(pre)}=>[KNN {k} @{_FIELD_VECTOR} $BLOB AS {_VEC_SCORE_ALIAS}]' + return await self._run_knn_search(client, index, query, query_embedding, k) + + @staticmethod + def _wrap_pre(pre: str) -> str: + """Parenthesize a multi-condition pre-filter before the ``=>`` KNN clause. + + When ``pre`` combines several terms (e.g. ``@file_id:{x} @document:term``) + the Valkey Search parser can otherwise mis-associate only the last term + with the KNN clause. Wrapping the whole expression forces correct + grouping. A bare ``*`` (match-all) and single-term expressions are left + untouched. + """ + if pre and pre != '*' and ' ' in pre.strip(): + return f'({pre})' + return pre + + async def _run_knn_search( + self, + client: GlideClient, + index: str, + query: str, + query_embedding: list[float], + k: int, + ) -> dict[str, Any]: + options = FtSearchOptions( + params={'BLOB': self._pack_vector(list(query_embedding))}, + return_fields=[ + ReturnField(field_identifier=_VEC_SCORE_ALIAS, alias='distance'), + ReturnField(field_identifier=_FIELD_DOCUMENT), + ReturnField(field_identifier=_FIELD_METADATA), + ], + limit=FtSearchLimit(0, k), + dialect=2, + ) + try: + reply = await ft.search(client, index, query, options) + except Exception as exc: + if self._is_missing_index_error(exc): + return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]} + raise + return self._reply_to_chroma(index, reply, has_distance=True) + + async def _run_text_search( + self, + client: GlideClient, + index: str, + query: str, + k: int, + ) -> dict[str, Any]: + options = FtSearchOptions( + return_fields=[ + ReturnField(field_identifier=_FIELD_DOCUMENT), + ReturnField(field_identifier=_FIELD_METADATA), + ], + limit=FtSearchLimit(0, k), + dialect=2, + ) + try: + reply = await ft.search(client, index, query, options) + except Exception as exc: + if self._is_missing_index_error(exc): + return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]} + raise + return self._reply_to_chroma(index, reply, has_distance=False) + + @staticmethod + def _is_missing_index_error(exc: Exception) -> bool: + """Return True if *exc* indicates the FT index does not exist. + + ``FT.DROPINDEX`` is applied eventually, so an index can briefly still + appear in ``FT._LIST`` after being dropped; a follow-up search then + fails with a "not found" error which we treat as an empty result. + """ + message = str(exc).lower() + return 'not found' in message and 'index' in message + + def _iter_reply_docs(self, reply: Any, prefix: str): + """Yield ``(doc_id, decoded_fields)`` pairs from an FT.SEARCH reply. + + glide returns ``[total, {key: {field: value}, ...}]``. This shared + iterator decodes each key, strips the per-collection prefix to recover + the original document id, and decodes the field map — the logic both + ``_reply_to_chroma`` and ``list_by_filter`` need. + """ + docs = reply[1] if reply and len(reply) >= 2 and isinstance(reply[1], dict) else {} + for key, fields in docs.items(): + key_str = self._decode(key) + doc_id = key_str[len(prefix) :] if prefix and key_str.startswith(prefix) else key_str + decoded_fields = {self._decode(fk): fv for fk, fv in fields.items()} if isinstance(fields, dict) else {} + yield doc_id, decoded_fields + + def _reply_to_chroma(self, index: str, reply: Any, has_distance: bool) -> dict[str, Any]: + """Convert an FT.SEARCH reply into Chroma-style nested lists. + + The KNN score field (aliased ``distance``) is a COSINE/L2 distance + directly, so no inversion is needed (unlike Qdrant). + """ + ids: list[str] = [] + distances: list[float] = [] + metadatas: list[dict[str, Any]] = [] + + if not reply or len(reply) < 2: + return {'ids': [ids], 'metadatas': [metadatas], 'distances': [distances]} + + prefix = self._key_prefix(index[len('idx:') :]) if index.startswith('idx:') else '' + + for doc_id, decoded_fields in self._iter_reply_docs(reply, prefix): + ids.append(doc_id) + + if has_distance and 'distance' in decoded_fields: + try: + distances.append(float(self._decode(decoded_fields['distance']))) + except (TypeError, ValueError): + distances.append(0.0) + else: + distances.append(0.0) + + metadata: dict[str, Any] = {} + raw_meta = decoded_fields.get(_FIELD_METADATA) + if raw_meta is not None: + try: + metadata = json.loads(self._decode(raw_meta)) + except (TypeError, ValueError): + metadata = {} + metadatas.append(metadata) + + return {'ids': [ids], 'metadatas': [metadatas], 'distances': [distances]} + + async def delete_by_file_id(self, collection: str, file_id: str) -> None: + client = await self._ensure_client() + index = self._index_name(collection) + if not await self._index_exists(client, index): + self.ap.logger.warning(f"Valkey Search collection '{collection}' not found for deletion") + return + + query = f'@{_FIELD_FILE_ID}:{{{self._encode_and_escape_tag(file_id)}}}' + keys = await self._search_keys(client, index, query) + if keys: + await client.delete(keys) + self.ap.logger.info( + f"Deleted {len(keys)} embeddings from Valkey Search collection '{collection}' with file_id: {file_id}" + ) + + async def delete_by_filter(self, collection: str, filter: dict[str, Any]) -> int: + client = await self._ensure_client() + index = self._index_name(collection) + if not await self._index_exists(client, index): + self.ap.logger.warning(f"Valkey Search collection '{collection}' not found for deletion") + return 0 + + # Guard against accidental mass deletion: a non-empty filter that maps + # to no usable (indexed) conditions must NOT fall back to match-all and + # wipe the whole collection. Skip instead (matching Milvus / pgvector). + query = self._triples_to_ft(filter) + if not query: + self.ap.logger.warning( + "Valkey Search delete_by_filter on '%s': filter produced no usable conditions, skipping", + collection, + ) + return 0 + keys = await self._search_keys(client, index, query) + if keys: + await client.delete(keys) + self.ap.logger.info(f"Deleted {len(keys)} embeddings from Valkey Search collection '{collection}' by filter") + return len(keys) + + async def list_by_filter( + self, + collection: str, + filter: dict[str, Any] | None = None, + limit: int = 20, + offset: int = 0, + ) -> tuple[list[dict[str, Any]], int]: + client = await self._ensure_client() + index = self._index_name(collection) + if not await self._index_exists(client, index): + return [], 0 + + query = self._triples_to_ft(filter) or _MATCH_ALL + options = FtSearchOptions( + return_fields=[ + ReturnField(field_identifier=_FIELD_DOCUMENT), + ReturnField(field_identifier=_FIELD_METADATA), + ], + limit=FtSearchLimit(offset, limit), + dialect=2, + ) + try: + reply = await ft.search(client, index, query, options) + except Exception as exc: + if self._is_missing_index_error(exc): + return [], 0 + raise + + total = 0 + if reply: + try: + total = int(reply[0]) + except (TypeError, ValueError): + total = 0 + + prefix = self._key_prefix(collection) + items: list[dict[str, Any]] = [] + for doc_id, decoded_fields in self._iter_reply_docs(reply, prefix): + document = decoded_fields.get(_FIELD_DOCUMENT) + metadata: dict[str, Any] = {} + raw_meta = decoded_fields.get(_FIELD_METADATA) + if raw_meta is not None: + try: + metadata = json.loads(self._decode(raw_meta)) + except (TypeError, ValueError): + metadata = {} + + items.append( + { + 'id': doc_id, + 'document': self._decode(document) if document is not None else None, + 'metadata': metadata, + } + ) + + return items, total + + async def delete_collection(self, collection: str): + client = await self._ensure_client() + index = self._index_name(collection) + self._ensured_indexes.discard(index) + + if await self._index_exists(client, index): + try: + await ft.dropindex(client, index) + except RequestError: + # The index was already dropped (e.g. by a concurrent process) + # between the existence check and this call — benign. Other + # errors (connection / auth) must propagate so the caller knows + # the operation failed rather than silently SCAN-deleting next. + pass + + # DROPINDEX does not remove the underlying hashes; delete them too. + prefix = self._key_prefix(collection) + cursor = b'0' + deleted = 0 + for _ in range(_MAX_SCAN_ROUNDS): + cursor, keys = await client.scan(cursor, match=f'{prefix}*', count=500) + if keys: + await client.delete(keys) + deleted += len(keys) + if cursor in (b'0', '0', 0): + break + self.ap.logger.info(f"Valkey Search collection '{collection}' deleted ({deleted} keys removed)") + + # ------------------------------------------------------------------ # + # Internal search helpers + # ------------------------------------------------------------------ # + async def _index_exists(self, client: GlideClient, index: str) -> bool: + if index in self._ensured_indexes: + return True + # ft.info is O(1) and raises RequestError when the index does not + # exist, vs ft.list which is O(n) over every index on the server and + # was being paid on the first query to each collection. + try: + await ft.info(client, index) + self._ensured_indexes.add(index) + return True + except RequestError: + return False + + async def _search_keys(self, client: GlideClient, index: str, query: str) -> list[str]: + """Return all matching document keys for a query (NOCONTENT). + + Paginates through the full result set in pages of ``_DELETE_SCAN_BATCH`` + so that queries matching more than one page of chunks are fully + enumerated (avoids silently truncating deletes and leaving orphaned + vectors). + """ + keys: list[str] = [] + offset = 0 + while True: + options = FtSearchOptions( + nocontent=True, + limit=FtSearchLimit(offset, _DELETE_SCAN_BATCH), + dialect=2, + ) + try: + reply = await ft.search(client, index, query, options) + except Exception as exc: + if self._is_missing_index_error(exc): + return keys + raise + + if not reply or len(reply) < 2: + break + + # reply[0] is the total match count; reply[1] holds this page. + total = 0 + try: + total = int(reply[0]) + except (TypeError, ValueError): + total = 0 + + docs = reply[1] + if isinstance(docs, dict): + page = [self._decode(k) for k in docs.keys()] + elif isinstance(docs, (list, tuple)): + page = [self._decode(k) for k in docs] + else: + page = [] + + if not page: + break + keys.extend(page) + + offset += len(page) + if offset >= total or len(page) < _DELETE_SCAN_BATCH: + break + + return keys diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index fea6db6f3..f4ae79bf8 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -87,6 +87,16 @@ vdb: database: 'langbot' user: 'postgres' password: 'postgres' + valkey_search: + host: 'localhost' + port: 6379 # integration tests use 6380 -> valkey/valkey-bundle:9.1.0 + db: 0 + password: '' # optional (toB auth) + username: '' # optional (ACL user, toB) + tls: false # optional (toB/SaaS) + index_algorithm: 'HNSW' # HNSW | FLAT + distance_metric: 'COSINE' # COSINE | L2 | IP + request_timeout: 5000 # per-request timeout in ms (glide default 250ms is too low for KNN) storage: use: local cleanup: diff --git a/tests/e2e/utils/config_factory.py b/tests/e2e/utils/config_factory.py index b2bc2f7d4..2df35b903 100644 --- a/tests/e2e/utils/config_factory.py +++ b/tests/e2e/utils/config_factory.py @@ -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', diff --git a/tests/integration/vector/__init__.py b/tests/integration/vector/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/vector/test_valkey_search.py b/tests/integration/vector/test_valkey_search.py new file mode 100644 index 000000000..64349987b --- /dev/null +++ b/tests/integration/vector/test_valkey_search.py @@ -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 diff --git a/tests/unit_tests/COVERAGE_EXCLUSIONS.md b/tests/unit_tests/COVERAGE_EXCLUSIONS.md index 1e3f28cec..f0e161158 100644 --- a/tests/unit_tests/COVERAGE_EXCLUSIONS.md +++ b/tests/unit_tests/COVERAGE_EXCLUSIONS.md @@ -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 测试 diff --git a/tests/unit_tests/vector/test_mgr.py b/tests/unit_tests/vector/test_mgr.py index 997861383..5c8927b54 100644 --- a/tests/unit_tests/vector/test_mgr.py +++ b/tests/unit_tests/vector/test_mgr.py @@ -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 = { diff --git a/tests/unit_tests/vector/test_valkey_search_filter.py b/tests/unit_tests/vector/test_valkey_search_filter.py new file mode 100644 index 000000000..7439712dd --- /dev/null +++ b/tests/unit_tests/vector/test_valkey_search_filter.py @@ -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(' 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'}) diff --git a/uv.lock b/uv.lock index 5fd22f65d..3e4b2a124 100644 --- a/uv.lock +++ b/uv.lock @@ -2084,6 +2084,7 @@ dependencies = [ { name = "tiktoken" }, { name = "urllib3" }, { name = "uv" }, + { name = "valkey-glide" }, { name = "websockets" }, ] @@ -2172,6 +2173,7 @@ requires-dist = [ { name = "tiktoken", specifier = ">=0.9.0" }, { name = "urllib3", specifier = ">=2.7.0" }, { name = "uv", specifier = ">=0.11.15" }, + { name = "valkey-glide", specifier = ">=2.4.1,<3.0.0" }, { name = "websockets", specifier = ">=15.0.1" }, ] @@ -5984,6 +5986,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e4/16/c1fd27e9549f3c4baf1dc9c20c456cd2f822dbf8de9f463824b0c0357e06/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6cde23eeda1a25c75b2e07d39970f3374105d5eafbaab2a4482be82f272d5a5e", size = 4296730, upload-time = "2025-10-16T22:17:00.744Z" }, ] +[[package]] +name = "valkey-glide" +version = "2.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "protobuf" }, + { name = "sniffio" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/a2/582b34c6acc8dc857c537f6007459cba48dfa0dc404789a657e5c1a998c0/valkey_glide-2.4.1.tar.gz", hash = "sha256:f1155d84156d11b90488aa67e90102f0bf98a45314f5b99308ac9074c05f7241", size = 898030, upload-time = "2026-05-28T21:41:55.881Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/60/60/961ce40492a56ef831a905dfe03df4a81c0705152f6a8e49c541c634f49e/valkey_glide-2.4.1-cp311-cp311-macosx_10_7_x86_64.whl", hash = "sha256:d7285d03c2df040f26874b7f4ae96f040da2daecc9a34fa99da6f4e6ce5149c8", size = 7482152, upload-time = "2026-05-28T21:41:02.205Z" }, + { url = "https://files.pythonhosted.org/packages/a4/b2/5a05567f0fc385dcbbbf6ab1061f0bc00443d51c2996e95eed45feaedda9/valkey_glide-2.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:5d2e82b74127897ccb7a957ad455787816a75fdc8c60a5e8004aef65ea93e99c", size = 6928601, upload-time = "2026-05-28T21:41:04.543Z" }, + { url = "https://files.pythonhosted.org/packages/c5/d9/7ea2b47cff0a2f99921eb0db404215f828ced7814bd09ede9c93b65d20bc/valkey_glide-2.4.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4094128cb07e06e87013b7afab1e9388f8f5aeebe48ea6cbd54de15bd772e644", size = 7236977, upload-time = "2026-05-28T21:41:06.055Z" }, + { url = "https://files.pythonhosted.org/packages/00/7a/6cda6b42156ed260e765e4ad2d6ab831607775e218a00fbb0d93411c4e8f/valkey_glide-2.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9f8dc0f3a36adb1cbe4e167972ca4758acdfed6baf58a4db94bbb713df56c8f5", size = 7691446, upload-time = "2026-05-28T21:41:07.833Z" }, + { url = "https://files.pythonhosted.org/packages/c5/b4/da8c058baaee414a6bb2450742359f3b3b6993b23281bf227c5089f0099c/valkey_glide-2.4.1-cp312-cp312-macosx_10_7_x86_64.whl", hash = "sha256:5f8df64f6a4f0fd7203113103101fdf0aaa7ff0e7557312611de11ab89c6db75", size = 7472646, upload-time = "2026-05-28T21:41:09.451Z" }, + { url = "https://files.pythonhosted.org/packages/f5/94/e1e311cb56597272b9cb69afb3fe8e2e7dd3371f88c92836015deddc6f49/valkey_glide-2.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b45e35f44c17e88f8cd8082f8d8061a9763238c44ef20b11b615f6d87235864a", size = 6943375, upload-time = "2026-05-28T21:41:11.079Z" }, + { url = "https://files.pythonhosted.org/packages/76/00/0e42e2f6866ebf0de552e076dc585a487b488b5b818c52460d28b50de65b/valkey_glide-2.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cf812b498925a30abab6e1a9f82f5eb821e967904fe7724729b2c82c47e29edf", size = 7237469, upload-time = "2026-05-28T21:41:12.733Z" }, + { url = "https://files.pythonhosted.org/packages/f5/4c/c5dd9a1ed995453b0d9ca75a5af87e881c14e6eebdbf5a5fa78c3bae23fc/valkey_glide-2.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:214e2faca98966eea3eaf9e09de616862423815a5059843a9884125e2427a344", size = 7678744, upload-time = "2026-05-28T21:41:14.634Z" }, + { url = "https://files.pythonhosted.org/packages/6a/2f/3df5702fc68684cef3e09f9cb6ed85578ddb08dc43593b1694c977f396fa/valkey_glide-2.4.1-cp313-cp313-macosx_10_7_x86_64.whl", hash = "sha256:c18976553ba663c03f7cc18c7e6075f4cbd2236c18b051e3d55bb213c6c44cb4", size = 7472972, upload-time = "2026-05-28T21:41:16.063Z" }, + { url = "https://files.pythonhosted.org/packages/54/a3/6a74c6f996fa9e411e66b6f0e645fead2e0a341f1371e4cf3212efa54412/valkey_glide-2.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:43006e19cd63d66051263fa34a8ad47ba7d08a199585689b3f12f56ed6c9a005", size = 6943012, upload-time = "2026-05-28T21:41:17.492Z" }, + { url = "https://files.pythonhosted.org/packages/fc/e7/d10ec41dca703f8c5dcbcba2b905e660c1cf56be53c4d5e368d7aa23d220/valkey_glide-2.4.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b652a2a62aad87738e8f0e0aa5bf660ba91449c9fdb88550ccbc42e5fec08fe7", size = 7237842, upload-time = "2026-05-28T21:41:18.995Z" }, + { url = "https://files.pythonhosted.org/packages/0a/a3/8916a9ed9e871686db444c86e601773245852ba1ad451ce1bb06f7aed91d/valkey_glide-2.4.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fbd27d26947fd9f1b6e9eaf0abce4bccfde779c1e618b310c4d725424b609793", size = 7678919, upload-time = "2026-05-28T21:41:20.502Z" }, + { url = "https://files.pythonhosted.org/packages/05/35/6d39ec3cbd24d85ad8e1051e29e6509c0999f760aff5af7851c1a1981471/valkey_glide-2.4.1-cp314-cp314-macosx_10_7_x86_64.whl", hash = "sha256:91fb7ff97acdabc8f641255b548a48627bb731e65037b1126745bf8a0022e87d", size = 7471906, upload-time = "2026-05-28T21:41:22.135Z" }, + { url = "https://files.pythonhosted.org/packages/ab/fc/3c28f794b7d35e13101598669c1d249c0a9f0408c545c87212e364c6ee4e/valkey_glide-2.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d49a2537c2de44b0fc57691b1ae6c3d6f481e6f7f7eb879c0d28921d0aaec67d", size = 6943495, upload-time = "2026-05-28T21:41:23.783Z" }, + { url = "https://files.pythonhosted.org/packages/2e/15/fb884631f5df78dc538c56bca9391165e40906b9b63ca65633d1be5bf980/valkey_glide-2.4.1-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cded9f14e448da5a96f61c066395f2c7e2846f2afe74cacc8634da0ae0c3425f", size = 7257720, upload-time = "2026-05-28T21:41:25.361Z" }, + { url = "https://files.pythonhosted.org/packages/73/79/0b881017194386d21812b929a81dd8afd51d6b8d92280895b45913854785/valkey_glide-2.4.1-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5f249ab5bd0d69befe35897cf51a8fc9e01e9c8c9fe03087a68e6fe6d3e31d0d", size = 7682318, upload-time = "2026-05-28T21:41:26.996Z" }, + { url = "https://files.pythonhosted.org/packages/7a/4d/f2b4e508692fcd21e76c7cbdc4f988bec7f4675e60f4f35ef482a826f6ae/valkey_glide-2.4.1-pp311-pypy311_pp73-macosx_10_7_x86_64.whl", hash = "sha256:775df9c7421a187c41caf003e4af5f073ed7e4b8abe50f8b9bec712cb03e12bf", size = 7479155, upload-time = "2026-05-28T21:41:42.399Z" }, + { url = "https://files.pythonhosted.org/packages/52/d8/8a3495f5582dccb4c8e7faf6a73baf3dbc4580701923f06d8abf210ff22d/valkey_glide-2.4.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:0d87f21c77004240189cc3c5aab156966487afd81ffdee04225a52c7bd7132e4", size = 6938571, upload-time = "2026-05-28T21:41:44.078Z" }, + { url = "https://files.pythonhosted.org/packages/f3/5a/a70077f76c2f18e94ec4309857b248beb7a8c7a3a50e30242abde2c3827d/valkey_glide-2.4.1-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:44376ef5fe7a25287095b073d8abde510a50b1ead0143662394b3da9717863ef", size = 7260021, upload-time = "2026-05-28T21:41:45.837Z" }, + { url = "https://files.pythonhosted.org/packages/aa/12/72d31522e06fcc9b391118c1f69a09002224e78114b1db0d01b96008dc59/valkey_glide-2.4.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a59cc0a21d7a8b1b3caeb299f23817429b5fe6579bd4cb016382e6b7a10de984", size = 7693093, upload-time = "2026-05-28T21:41:47.617Z" }, +] + [[package]] name = "virtualenv" version = "20.36.1"