From ca126edfab46e4d0a9e48f0766cdbb4cfe13a1b5 Mon Sep 17 00:00:00 2001 From: Daria Korenieva Date: Sat, 9 May 2026 04:13:52 -0700 Subject: [PATCH] feat(ext): add Valkey vector store integration (#3051) Signed-off-by: Daria Korenieva --- .../app/config_storageconfig_028579.mdx | 51 +- docs/docs/config-reference/overview.mdx | 57 +- .../config-reference/vector_store/index.mdx | 15 +- ...valkey_store_valkeyvectorconfig_a0fffb.mdx | 115 +++ docs/docs/faq/kbqa.md | 2 +- docs/docs/getting-started/cli-quickstart.md | 1 + docs/docs/installation/build_image.md | 1 + docs/docs/modules/rag.md | 1 + packages/dbgpt-ext/pyproject.toml | 1 + .../src/dbgpt_ext/storage/__init__.py | 12 + .../storage/vector_store/tests/__init__.py | 1 + .../vector_store/tests/test_valkey_store.py | 533 +++++++++++ .../tests/test_valkey_store_integration.py | 664 ++++++++++++++ .../storage/vector_store/valkey_store.py | 868 ++++++++++++++++++ 14 files changed, 2277 insertions(+), 45 deletions(-) create mode 100644 docs/docs/config-reference/vector_store/valkey_store_valkeyvectorconfig_a0fffb.mdx create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/__init__.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store_integration.py create mode 100644 packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/valkey_store.py diff --git a/docs/docs/config-reference/app/config_storageconfig_028579.mdx b/docs/docs/config-reference/app/config_storageconfig_028579.mdx index d9969dfea..648d92b9d 100644 --- a/docs/docs/config-reference/app/config_storageconfig_028579.mdx +++ b/docs/docs/config-reference/app/config_storageconfig_028579.mdx @@ -28,8 +28,18 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; }, { "type": "link", - "text": "pgvector configuration", - "url": "../vector_store/pgvector_store_pgvectorconfig_3ef448" + "text": "qdrant configuration", + "url": "../vector_store/qdrant_store_qdrantvectorconfig_3d0339" + }, + { + "type": "link", + "text": "valkey configuration", + "url": "../vector_store/valkey_store_valkeyvectorconfig_a0fffb" + }, + { + "type": "link", + "text": "oceanbase configuration", + "url": "../vector_store/oceanbase_store_oceanbaseconfig_220e36" }, { "type": "link", @@ -38,13 +48,8 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; }, { "type": "link", - "text": "milvus configuration", - "url": "../vector_store/milvus_store_milvusvectorconfig_20af52" - }, - { - "type": "link", - "text": "oceanbase configuration", - "url": "../vector_store/oceanbase_store_oceanbaseconfig_220e36" + "text": "pgvector configuration", + "url": "../vector_store/pgvector_store_pgvectorconfig_3ef448" } ], "defaultValue": "ChromaVectorConfig" @@ -59,6 +64,11 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; "type": "link", "text": "tugraph configuration", "url": "../graph_store/tugraph_store_tugraphstoreconfig_7ca8a8" + }, + { + "type": "link", + "text": "neo4j configuration", + "url": "../graph_store/neo4j_store_neo4jstoreconfig_a4db5d" } ] }, @@ -80,8 +90,18 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; }, { "type": "link", - "text": "pgvector configuration", - "url": "../vector_store/pgvector_store_pgvectorconfig_3ef448" + "text": "qdrant configuration", + "url": "../vector_store/qdrant_store_qdrantvectorconfig_3d0339" + }, + { + "type": "link", + "text": "valkey configuration", + "url": "../vector_store/valkey_store_valkeyvectorconfig_a0fffb" + }, + { + "type": "link", + "text": "oceanbase configuration", + "url": "../vector_store/oceanbase_store_oceanbaseconfig_220e36" }, { "type": "link", @@ -90,13 +110,8 @@ import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; }, { "type": "link", - "text": "milvus configuration", - "url": "../vector_store/milvus_store_milvusvectorconfig_20af52" - }, - { - "type": "link", - "text": "oceanbase configuration", - "url": "../vector_store/oceanbase_store_oceanbaseconfig_220e36" + "text": "pgvector configuration", + "url": "../vector_store/pgvector_store_pgvectorconfig_3ef448" } ] } diff --git a/docs/docs/config-reference/overview.mdx b/docs/docs/config-reference/overview.mdx index 352ab836d..14008bc2c 100644 --- a/docs/docs/config-reference/overview.mdx +++ b/docs/docs/config-reference/overview.mdx @@ -9,17 +9,17 @@ This document provides an overview of all configuration classes organized by typ ## Configuration Types - [app](#type-app) (6 classes) -- [datasource](#type-datasource) (14 classes) -- [embedding](#type-embedding) (6 classes) -- [graph_store](#type-graph_store) (4 classes) -- [llm](#type-llm) (22 classes) +- [datasource](#type-datasource) (18 classes) +- [embedding](#type-embedding) (8 classes) +- [graph_store](#type-graph_store) (3 classes) +- [llm](#type-llm) (28 classes) - [memory](#type-memory) (2 classes) - [other](#type-other) (1 classes) -- [reranker](#type-reranker) (3 classes) +- [reranker](#type-reranker) (6 classes) - [serve](#type-serve) (13 classes) - [service](#type-service) (1 classes) - [utils](#type-utils) (2 classes) -- [vector_store](#type-vector_store) (6 classes) +- [vector_store](#type-vector_store) (7 classes) ## Type Details @@ -42,7 +42,7 @@ This type contains 6 configuration classes. ### datasource {#type-datasource} -This type contains 14 configuration classes. +This type contains 18 configuration classes. #### Configuration Classes @@ -51,10 +51,13 @@ This type contains 14 configuration classes. | [ClickhouseParameters](datasource/conn_clickhouse_clickhouseparameters_4a1237) | | | [DorisParameters](datasource/conn_doris_dorisparameters_e33c53) | | | [DuckDbConnectorParameters](datasource/conn_duckdb_duckdbconnectorparameters_c672c7) | | +| [GaussDBParameters](datasource/conn_gaussdb_gaussdbparameters_9c9811) | | | [HiveParameters](datasource/conn_hive_hiveparameters_ec3601) | | | [MSSQLParameters](datasource/conn_mssql_mssqlparameters_d79d1c) | | | [MySQLParameters](datasource/conn_mysql_mysqlparameters_4393c4) | | +| [Neo4jParameters](datasource/conn_neo4j_neo4jparameters_041b12) | | | [OceanBaseParameters](datasource/conn_oceanbase_oceanbaseparameters_260d2d) | | +| [OracleParameters](datasource/conn_oracle_oracleparameters_c011ed) | | | [PostgreSQLParameters](datasource/conn_postgresql_postgresqlparameters_22efa5) | | | [RDBMSDatasourceParameters](datasource/base_rdbmsdatasourceparameters_4f774f) | | | [SQLiteConnectorParameters](datasource/conn_sqlite_sqliteconnectorparameters_82c8b5) | | @@ -62,35 +65,37 @@ This type contains 14 configuration classes. | [StarRocksParameters](datasource/conn_starrocks_starrocksparameters_e511f7) | | | [TuGraphParameters](datasource/conn_tugraph_tugraphparameters_0c844e) | | | [VerticaParameters](datasource/conn_vertica_verticaparameters_c712b8) | | +| [openGaussParameters](datasource/conn_opengauss_opengaussparameters_85409c) | | --- ### embedding {#type-embedding} -This type contains 6 configuration classes. +This type contains 8 configuration classes. #### Configuration Classes | Class | Description | |-------|-------------| +| [AimlapiEmbeddingDeployModelParameters](embedding/aimlapi_aimlapiembeddingdeploymodelparameters_97d1c4) | | | [HFEmbeddingDeployModelParameters](embedding/embeddings_hfembeddingdeploymodelparameters_f588e1) | | | [JinaEmbeddingsDeployModelParameters](embedding/jina_jinaembeddingsdeploymodelparameters_40b0f2) | | | [OllamaEmbeddingDeployModelParameters](embedding/ollama_ollamaembeddingdeploymodelparameters_b511e0) | | | [OpenAPIEmbeddingDeployModelParameters](embedding/embeddings_openapiembeddingdeploymodelparameters_f9ba47) | | | [QianfanEmbeddingDeployModelParameters](embedding/qianfan_qianfanembeddingdeploymodelparameters_257d2a) | | +| [SiliconFlowEmbeddingDeployModelParameters](embedding/siliconflow_siliconflowembeddingdeploymodelparameters_113c43) | | | [TongyiEmbeddingDeployModelParameters](embedding/tongyi_tongyiembeddingdeploymodelparameters_a7cbb4) | | --- ### graph_store {#type-graph_store} -This type contains 4 configuration classes. +This type contains 3 configuration classes. #### Configuration Classes | Class | Description | |-------|-------------| -| [BuiltinKnowledgeGraphConfig](graph_store/knowledge_graph_builtinknowledgegraphconfig_f26e05) | | | [Neo4jStoreConfig](graph_store/neo4j_store_neo4jstoreconfig_a4db5d) | | | [OpenSPGConfig](graph_store/open_spg_openspgconfig_a744fd) | | | [TuGraphStoreConfig](graph_store/tugraph_store_tugraphstoreconfig_7ca8a8) | | @@ -99,7 +104,7 @@ This type contains 4 configuration classes. ### llm {#type-llm} -This type contains 22 configuration classes. +This type contains 28 configuration classes. #### Relationships @@ -114,17 +119,23 @@ graph TD | Class | Description | |-------|-------------| +| [AimlapiDeployModelParameters](llm/aimlapi_aimlapideploymodelparameters_7c1b54) | | | [BaichuanDeployModelParameters](llm/baichuan_baichuandeploymodelparameters_0bf9cc) | | | [BitsandbytesQuantization](llm/parameter_bitsandbytesquantization_d40e3b) | | | [BitsandbytesQuantization4bits](llm/parameter_bitsandbytesquantization4bits_52b778) | | | [BitsandbytesQuantization8bits](llm/parameter_bitsandbytesquantization8bits_909aed) | | +| [BurnCloudDeployModelParameters](llm/burncloud_burnclouddeploymodelparameters_e59cbb) | | | [ClaudeDeployModelParameters](llm/claude_claudedeploymodelparameters_1f0c45) | | | [DeepSeekDeployModelParameters](llm/deepseek_deepseekdeploymodelparameters_194cbd) | | | [GeminiDeployModelParameters](llm/gemini_geminideploymodelparameters_5113b9) | | | [GiteeDeployModelParameters](llm/gitee_giteedeploymodelparameters_d1bdb3) | | | [HFLLMDeployModelParameters](llm/hf_adapter_hfllmdeploymodelparameters_103e81) | | +| [InfiniAIDeployModelParameters](llm/infiniai_infiniaideploymodelparameters_682301) | | +| [LiteLLMDeployModelParameters](llm/litellm_litellmdeploymodelparameters_d053f9) | | | [LlamaCppModelParameters](llm/llama_cpp_py_adapter_llamacppmodelparameters_e88874) | | | [LlamaServerParameters](llm/llama_cpp_adapter_llamaserverparameters_421f40) | | +| [MLXDeployModelParameters](llm/mlx_adapter_mlxdeploymodelparameters_f26f45) | | +| [MiniMaxDeployModelParameters](llm/minimax_minimaxdeploymodelparameters_97ce61) | | | [MoonshotDeployModelParameters](llm/moonshot_moonshotdeploymodelparameters_aa2f6b) | | | [OllamaDeployModelParameters](llm/ollama_ollamadeploymodelparameters_d55be6) | | | [OpenAICompatibleDeployModelParameters](llm/chatgpt_openaicompatibledeploymodelparameters_c3d426) | | @@ -166,15 +177,18 @@ This type contains 1 configuration classes. ### reranker {#type-reranker} -This type contains 3 configuration classes. +This type contains 6 configuration classes. #### Configuration Classes | Class | Description | |-------|-------------| | [CrossEncoderRerankEmbeddingsParameters](reranker/rerank_crossencoderrerankembeddingsparameters_63ec13) | | +| [InfiniAIRerankEmbeddingsParameters](reranker/rerank_infiniairerankembeddingsparameters_879440) | | | [OpenAPIRerankerDeployModelParameters](reranker/rerank_openapirerankerdeploymodelparameters_778108) | | +| [QwenRerankEmbeddingsParameters](reranker/rerank_qwenrerankembeddingsparameters_18e7d2) | | | [SiliconFlowRerankEmbeddingsParameters](reranker/rerank_siliconflowrerankembeddingsparameters_af0257) | | +| [TeiEmbeddingsParameters](reranker/rerank_teiembeddingsparameters_e38e28) | | --- @@ -229,7 +243,7 @@ This type contains 2 configuration classes. ### vector_store {#type-vector_store} -This type contains 6 configuration classes. +This type contains 7 configuration classes. #### Configuration Classes @@ -237,9 +251,10 @@ This type contains 6 configuration classes. |-------|-------------| | [ChromaVectorConfig](vector_store/chroma_store_chromavectorconfig_16224f) | | | [ElasticsearchStoreConfig](vector_store/elastic_store_elasticsearchstoreconfig_15bdb6) | | -| [MilvusVectorConfig](vector_store/milvus_store_milvusvectorconfig_20af52) | | | [OceanBaseConfig](vector_store/oceanbase_store_oceanbaseconfig_220e36) | | | [PGVectorConfig](vector_store/pgvector_store_pgvectorconfig_3ef448) | | +| [QdrantVectorConfig](vector_store/qdrant_store_qdrantvectorconfig_3d0339) | | +| [ValkeyVectorConfig](vector_store/valkey_store_valkeyvectorconfig_a0fffb) | | | [WeaviateVectorConfig](vector_store/weaviate_store_weaviatevectorconfig_093ce3) | | --- @@ -251,18 +266,18 @@ The following diagram shows relationships between different configuration types: ```mermaid graph TD other[other - 1 classes] - datasource[datasource - 14 classes] - llm[llm - 22 classes] - embedding[embedding - 6 classes] - reranker[reranker - 3 classes] + datasource[datasource - 18 classes] + llm[llm - 28 classes] + embedding[embedding - 8 classes] + reranker[reranker - 6 classes] service[service - 1 classes] - graph_store[graph_store - 4 classes] - vector_store[vector_store - 6 classes] + graph_store[graph_store - 3 classes] + vector_store[vector_store - 7 classes] serve[serve - 13 classes] memory[memory - 2 classes] app[app - 6 classes] utils[utils - 2 classes] - service -->|14 connections| datasource + service -->|18 connections| datasource other -->|13 connections| serve serve -->|2 connections| utils app -->|14 connections| memory diff --git a/docs/docs/config-reference/vector_store/index.mdx b/docs/docs/config-reference/vector_store/index.mdx index c29492f87..708615fcc 100644 --- a/docs/docs/config-reference/vector_store/index.mdx +++ b/docs/docs/config-reference/vector_store/index.mdx @@ -22,11 +22,6 @@ import { ConfigClassTable } from '@site/src/components/mdx/ConfigClassTable'; "description": "Elasticsearch vector config.", "link": "./elastic_store_elasticsearchstoreconfig_15bdb6" }, - { - "name": "MilvusVectorConfig", - "description": "Milvus vector config.", - "link": "./milvus_store_milvusvectorconfig_20af52" - }, { "name": "OceanBaseConfig", "description": "OceanBase vector store config.", @@ -37,6 +32,16 @@ import { ConfigClassTable } from '@site/src/components/mdx/ConfigClassTable'; "description": "PG vector config.", "link": "./pgvector_store_pgvectorconfig_3ef448" }, + { + "name": "QdrantVectorConfig", + "description": "Qdrant vector store config.", + "link": "./qdrant_store_qdrantvectorconfig_3d0339" + }, + { + "name": "ValkeyVectorConfig", + "description": "Valkey vector store config.", + "link": "./valkey_store_valkeyvectorconfig_a0fffb" + }, { "name": "WeaviateVectorConfig", "description": "Weaviate vector config.", diff --git a/docs/docs/config-reference/vector_store/valkey_store_valkeyvectorconfig_a0fffb.mdx b/docs/docs/config-reference/vector_store/valkey_store_valkeyvectorconfig_a0fffb.mdx new file mode 100644 index 000000000..aee8509a4 --- /dev/null +++ b/docs/docs/config-reference/vector_store/valkey_store_valkeyvectorconfig_a0fffb.mdx @@ -0,0 +1,115 @@ +--- +title: "Valkey Config Configuration" +description: "Valkey vector store config." +--- + +import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail"; + + + diff --git a/docs/docs/faq/kbqa.md b/docs/docs/faq/kbqa.md index 0ed3afee4..61c919e7b 100644 --- a/docs/docs/faq/kbqa.md +++ b/docs/docs/faq/kbqa.md @@ -18,7 +18,7 @@ git lfs clone https://huggingface.co/GanymedeNil/text2vec-large-chinese Update .env file and set VECTOR_STORE_TYPE. -DB-GPT currently support Chroma(Default), Milvus(>2.1), Weaviate, OceanBase vector database. +DB-GPT currently support Chroma(Default), Milvus(>2.1), Weaviate, Valkey, OceanBase vector database. If you want to change vector db, Update your .env, set your vector store type, VECTOR_STORE_TYPE=Chroma (now only support Chroma and Milvus(>2.1), if you set Milvus, please set MILVUS_URL and MILVUS_PORT). If you want to use OceanBase, please first start a docker container via the following command: diff --git a/docs/docs/getting-started/cli-quickstart.md b/docs/docs/getting-started/cli-quickstart.md index 849ad7158..337f37fa0 100644 --- a/docs/docs/getting-started/cli-quickstart.md +++ b/docs/docs/getting-started/cli-quickstart.md @@ -369,6 +369,7 @@ The core framework is included by default when you `pip install dbgpt-app`. Use |-------|---------|-------------| | `storage_chromadb` | ChromaDB | `chromadb`, `onnxruntime` | | `storage_milvus` | Milvus | `pymilvus` | +| `storage_valkey` | Valkey | `valkey-glide` | | `storage_weaviate` | Weaviate | `weaviate-client` | | `storage_elasticsearch` | Elasticsearch | `elasticsearch` | | `storage_obvector` | OBVector | `pyobvector` | diff --git a/docs/docs/installation/build_image.md b/docs/docs/installation/build_image.md index 46975ed23..6c76e3b99 100644 --- a/docs/docs/installation/build_image.md +++ b/docs/docs/installation/build_image.md @@ -182,6 +182,7 @@ Here are some useful extras you can add: | Extra Package | Description | |--------------|-------------| | `storage_milvus` | Vector store integration with Milvus | +| `storage_valkey` | Vector store integration with Valkey | | `storage_elasticsearch` | Vector store integration with Elasticsearch | | `datasource_postgres` | Database connector for PostgreSQL | | `vllm` | VLLM integration for optimized inference | diff --git a/docs/docs/modules/rag.md b/docs/docs/modules/rag.md index eea3e3bcf..1a1a68b14 100644 --- a/docs/docs/modules/rag.md +++ b/docs/docs/modules/rag.md @@ -95,6 +95,7 @@ When creating a knowledge base, you can choose from three storage types: | **ChromaDB** | Default embedded vector database, zero setup | `storage_chromadb` | | **Milvus** | Distributed vector database for production scale | `storage_milvus` | | **PGVector** | PostgreSQL extension for vector operations | `storage_pgvector` | +| **Valkey** | High-performance in-memory vector store with HNSW/FLAT indexing | `storage_valkey` | | **Weaviate** | Cloud-native vector search engine | `storage_weaviate` | | **Elasticsearch** | Full-text + vector hybrid search | `storage_elasticsearch` | | **OceanBase** | Cloud-native distributed database | `storage_oceanbase` | diff --git a/packages/dbgpt-ext/pyproject.toml b/packages/dbgpt-ext/pyproject.toml index d7c388f6e..6db7403eb 100644 --- a/packages/dbgpt-ext/pyproject.toml +++ b/packages/dbgpt-ext/pyproject.toml @@ -80,6 +80,7 @@ storage_chromadb = [ storage_elasticsearch = ["elasticsearch"] storage_qdrant = ["qdrant-client>=1.17.1; python_version < '3.13'"] storage_obvector = ["pyobvector"] +storage_valkey = ["valkey-glide>=2.3.0"] file_oss = [ "oss2" # Aliyun OSS diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/__init__.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/__init__.py index d6b897c30..d750f1f9a 100644 --- a/packages/dbgpt-ext/src/dbgpt_ext/storage/__init__.py +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/__init__.py @@ -66,6 +66,15 @@ def _import_qdrant() -> Tuple[Type, Type]: return QdrantStore, QdrantVectorConfig +def _import_valkey() -> Tuple[Type, Type]: + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + return ValkeyStore, ValkeyVectorConfig + + def _import_builtin_knowledge_graph() -> Tuple[Type, Type]: from dbgpt_ext.storage.knowledge_graph.knowledge_graph import ( BuiltinKnowledgeGraph, @@ -112,6 +121,8 @@ def _select_rag_storage(name: str) -> Tuple[Type, Type]: return _import_elastic() elif name == "Qdrant": return _import_qdrant() + elif name == "Valkey": + return _import_valkey() elif name == "KnowledgeGraph": return _import_builtin_knowledge_graph() elif name == "CommunitySummaryKnowledgeGraph": @@ -161,6 +172,7 @@ __vector_store__ = [ "PGVector", "ElasticSearch", "Qdrant", + "Valkey", ] __knowledge_graph__ = ["KnowledgeGraph", "CommunitySummaryKnowledgeGraph", "OpenSPG"] diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/__init__.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/__init__.py new file mode 100644 index 000000000..dc12c7e07 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for vector store implementations.""" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store.py new file mode 100644 index 000000000..a5864e561 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store.py @@ -0,0 +1,533 @@ +"""Unit tests for ValkeyStore. + +These tests cover pure logic that does not require a running Valkey server +or heavy mocking. Integration tests cover actual Valkey operations. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from dbgpt.storage.vector_store.filters import ( + FilterCondition, + FilterOperator, + MetadataFilter, + MetadataFilters, +) +from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + _escape_tag_value, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def mock_embedding_fn(): + """Create a mock embedding function.""" + embedding = MagicMock() + embedding.embed_query.return_value = [0.1, 0.2, 0.3, 0.4] + embedding.embed_documents.return_value = [ + [0.1, 0.2, 0.3, 0.4], + [0.5, 0.6, 0.7, 0.8], + ] + return embedding + + +@pytest.fixture +def valkey_config(): + """Create a ValkeyVectorConfig for testing.""" + return ValkeyVectorConfig( + host="localhost", + port=6379, + password=None, + use_ssl=False, + index_type="HNSW", + distance_metric="COSINE", + key_prefix="test_vec:", + hnsw_m=16, + hnsw_ef_construction=200, + hnsw_ef_runtime=10, + metadata_schema={ + "source": "tag", + "page": "numeric", + "score": "numeric", + "category": "tag", + "status": "tag", + }, + ) + + +@pytest.fixture +def valkey_store(valkey_config, mock_embedding_fn): + """Create a ValkeyStore with mocked client for pure logic tests.""" + mock_client = MagicMock() + with ( + patch.object(ValkeyStore, "_create_client", return_value=mock_client), + patch.object(ValkeyStore, "_index_exists", return_value=True), + ): + store = ValkeyStore( + vector_store_config=valkey_config, + name="test_collection", + embedding_fn=mock_embedding_fn, + ) + store._client = mock_client + return store + + +# --------------------------------------------------------------------------- +# ValkeyVectorConfig tests (pure dataclass logic) +# --------------------------------------------------------------------------- + + +class TestValkeyVectorConfig: + """Tests for ValkeyVectorConfig.""" + + def test_default_config(self): + """Test default configuration values.""" + config = ValkeyVectorConfig() + assert config.host == "localhost" + assert config.port == 6379 + assert config.index_type == "HNSW" + assert config.distance_metric == "COSINE" + assert config.hnsw_m == 16 + assert config.hnsw_ef_construction == 200 + assert config.hnsw_ef_runtime == 10 + assert config.metadata_schema is None + + def test_custom_config(self, valkey_config): + """Test custom configuration values.""" + assert valkey_config.host == "localhost" + assert valkey_config.port == 6379 + assert valkey_config.key_prefix == "test_vec:" + assert valkey_config.index_type == "HNSW" + assert valkey_config.distance_metric == "COSINE" + assert valkey_config.metadata_schema == { + "source": "tag", + "page": "numeric", + "score": "numeric", + "category": "tag", + "status": "tag", + } + + def test_config_type(self, valkey_config): + """Test that __type__ is set correctly.""" + assert valkey_config.__type__ == "valkey" + + def test_flat_index_config(self): + """Test FLAT index type configuration.""" + config = ValkeyVectorConfig(index_type="FLAT") + assert config.index_type == "FLAT" + + def test_distance_metrics(self): + """Test different distance metric configurations.""" + for metric in ["COSINE", "L2", "IP"]: + config = ValkeyVectorConfig(distance_metric=metric) + assert config.distance_metric == metric + + +# --------------------------------------------------------------------------- +# Store initialization tests (pure validation logic) +# --------------------------------------------------------------------------- + + +class TestValkeyStoreInit: + """Tests for ValkeyStore initialization logic.""" + + def test_init_without_embedding_raises(self, valkey_config): + """Test that init without embedding_fn raises ValueError.""" + with pytest.raises(ValueError, match="embedding_fn is required"): + ValkeyStore(vector_store_config=valkey_config, name="test") + + def test_init_sets_collection_name(self, valkey_store): + """Test that collection name is set correctly.""" + assert valkey_store._collection_name == "test_collection" + + def test_init_sets_index_name(self, valkey_store): + """Test that index name is derived from collection name.""" + assert valkey_store._index_name == "idx:test_collection" + + def test_init_sets_key_prefix(self, valkey_store): + """Test that key prefix includes collection name.""" + assert valkey_store._key_prefix == "test_vec:test_collection:" + + def test_init_default_collection_name(self, valkey_config, mock_embedding_fn): + """Test default collection name when none is provided.""" + mock_client = MagicMock() + with ( + patch.object(ValkeyStore, "_create_client", return_value=mock_client), + patch.object(ValkeyStore, "_index_exists", return_value=True), + ): + store = ValkeyStore( + vector_store_config=valkey_config, + embedding_fn=mock_embedding_fn, + ) + assert store._collection_name == "dbgpt_collection" + + def test_get_config(self, valkey_store, valkey_config): + """Test get_config returns the config.""" + assert valkey_store.get_config() == valkey_config + + +# --------------------------------------------------------------------------- +# Filter expression building (pure string logic) +# --------------------------------------------------------------------------- + + +class TestValkeyStoreFilters: + """Tests for metadata filter conversion — pure logic, no I/O.""" + + def test_no_filters(self, valkey_store): + """Test that no filters returns wildcard.""" + assert valkey_store._build_filter_expression(None) == "*" + + def test_empty_filters(self, valkey_store): + """Test that empty filter list returns wildcard.""" + filters = MetadataFilters(filters=[]) + assert valkey_store._build_filter_expression(filters) == "*" + + def test_eq_string_filter(self, valkey_store): + """Test equality filter for string values.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="web") + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_source:{web}" in result + + def test_eq_numeric_filter(self, valkey_store): + """Test equality filter for numeric values.""" + filters = MetadataFilters( + filters=[MetadataFilter(key="page", operator=FilterOperator.EQ, value=5)] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_page:[5 5]" in result + + def test_gt_filter(self, valkey_store): + """Test greater-than filter.""" + filters = MetadataFilters( + filters=[MetadataFilter(key="score", operator=FilterOperator.GT, value=0.5)] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_score:[(0.5 +inf]" in result + + def test_gte_filter(self, valkey_store): + """Test greater-than-or-equal filter.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="score", operator=FilterOperator.GTE, value=0.5) + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_score:[0.5 +inf]" in result + + def test_lt_filter(self, valkey_store): + """Test less-than filter.""" + filters = MetadataFilters( + filters=[MetadataFilter(key="score", operator=FilterOperator.LT, value=0.8)] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_score:[-inf (0.8]" in result + + def test_lte_filter(self, valkey_store): + """Test less-than-or-equal filter.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="score", operator=FilterOperator.LTE, value=0.8) + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_score:[-inf 0.8]" in result + + def test_in_filter(self, valkey_store): + """Test IN filter.""" + filters = MetadataFilters( + filters=[ + MetadataFilter( + key="category", operator=FilterOperator.IN, value=["a", "b", "c"] + ) + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_category:{a|b|c}" in result + + def test_nin_filter(self, valkey_store): + """Test NOT IN filter.""" + filters = MetadataFilters( + filters=[ + MetadataFilter( + key="category", operator=FilterOperator.NIN, value=["x", "y"] + ) + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "-@meta_category:{x|y}" in result + + def test_ne_string_filter(self, valkey_store): + """Test not-equal filter for strings.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="status", operator=FilterOperator.NE, value="draft") + ] + ) + result = valkey_store._build_filter_expression(filters) + assert "-@meta_status:{draft}" in result + + def test_ne_numeric_filter(self, valkey_store): + """Test not-equal filter for numeric values.""" + filters = MetadataFilters( + filters=[MetadataFilter(key="page", operator=FilterOperator.NE, value=3)] + ) + result = valkey_store._build_filter_expression(filters) + assert "-@meta_page:[3 3]" in result + + def test_and_condition(self, valkey_store): + """Test AND condition joins with space.""" + filters = MetadataFilters( + condition=FilterCondition.AND, + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="web"), + MetadataFilter(key="page", operator=FilterOperator.GT, value=1), + ], + ) + result = valkey_store._build_filter_expression(filters) + assert "@meta_source:{web}" in result + assert "@meta_page:[(1 +inf]" in result + # AND uses space separator + assert " | " not in result + + def test_or_condition(self, valkey_store): + """Test OR condition joins with pipe.""" + filters = MetadataFilters( + condition=FilterCondition.OR, + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="web"), + MetadataFilter(key="source", operator=FilterOperator.EQ, value="pdf"), + ], + ) + result = valkey_store._build_filter_expression(filters) + assert " | " in result + + def test_filters_without_metadata_schema(self, mock_embedding_fn): + """Test that filters raise ValueError when no metadata_schema is set.""" + config = ValkeyVectorConfig(key_prefix="test_vec:", metadata_schema=None) + mock_client = MagicMock() + with ( + patch.object(ValkeyStore, "_create_client", return_value=mock_client), + patch.object(ValkeyStore, "_index_exists", return_value=True), + ): + store = ValkeyStore( + vector_store_config=config, + name="test_no_schema", + embedding_fn=mock_embedding_fn, + ) + store._client = mock_client + + filters = MetadataFilters( + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="web") + ] + ) + with pytest.raises(ValueError, match="no metadata_schema configured"): + store._build_filter_expression(filters) + + +# --------------------------------------------------------------------------- +# Result parsing (pure logic) +# --------------------------------------------------------------------------- + + +class TestValkeyStoreParseResults: + """Tests for search result parsing — pure logic, no I/O.""" + + def test_parse_empty_results(self, valkey_store): + """Test parsing None and empty list.""" + assert valkey_store._parse_search_results(None) == [] + assert valkey_store._parse_search_results([]) == [] + + def test_parse_structured_results(self, valkey_store): + """Test parsing results with .results attribute (newer glide).""" + mock_doc = MagicMock(spec=[]) + mock_doc.fields = { + "content": "test content", + "metadata": '{"source": "test"}', + "chunk_id": "c1", + "score": "0.1", + } + + mock_result = MagicMock(spec=[]) + mock_result.results = [mock_doc] + + chunks = valkey_store._parse_search_results(mock_result) + assert len(chunks) == 1 + assert chunks[0].content == "test content" + assert chunks[0].chunk_id == "c1" + assert chunks[0].score == pytest.approx(0.9) + + def test_parse_list_results_with_dict(self, valkey_store): + """Test parsing list-based results with dict format.""" + result = [ + 1, # total count + { + b"key1": { + b"content": b"hello", + b"metadata": b'{"k": "v"}', + b"chunk_id": b"id1", + b"score": b"0.3", + } + }, + ] + chunks = valkey_store._parse_search_results(result) + assert len(chunks) == 1 + assert chunks[0].content == "hello" + assert chunks[0].chunk_id == "id1" + assert chunks[0].score == pytest.approx(0.7) + + def test_doc_to_chunk_with_bytes(self, valkey_store): + """Test parsing document with bytes values.""" + doc = { + "content": b"byte content", + "metadata": b'{"key": "val"}', + "chunk_id": b"c1", + "score": b"0.2", + } + chunk = valkey_store._doc_to_chunk(doc) + assert chunk is not None + assert chunk.content == "byte content" + assert chunk.metadata == {"key": "val"} + assert chunk.chunk_id == "c1" + assert chunk.score == pytest.approx(0.8) + + def test_doc_to_chunk_with_strings(self, valkey_store): + """Test parsing document with string values.""" + doc = { + "content": "string content", + "metadata": '{"key": "val"}', + "chunk_id": "c2", + "score": "0.0", + } + chunk = valkey_store._doc_to_chunk(doc) + assert chunk is not None + assert chunk.content == "string content" + assert chunk.score == pytest.approx(1.0) + + def test_doc_to_chunk_with_flat_list(self, valkey_store): + """Test parsing flat list format [field, value, field, value, ...].""" + doc = [ + "content", + "list content", + "metadata", + "{}", + "chunk_id", + "c3", + "score", + "0.5", + ] + chunk = valkey_store._doc_to_chunk(doc) + assert chunk is not None + assert chunk.content == "list content" + assert chunk.chunk_id == "c3" + assert chunk.score == pytest.approx(0.5) + + def test_doc_to_chunk_invalid_returns_none(self, valkey_store): + """Test that invalid doc types return None.""" + assert valkey_store._doc_to_chunk(42) is None + assert valkey_store._doc_to_chunk("invalid") is None + + +# --------------------------------------------------------------------------- +# Score conversion for different metrics (pure logic) +# --------------------------------------------------------------------------- + + +class TestScoreConversion: + """Tests for distance-to-similarity score conversion.""" + + def _make_store_with_metric(self, metric, mock_embedding_fn): + """Helper to create a store with a specific distance metric.""" + config = ValkeyVectorConfig( + key_prefix="test_vec:", + distance_metric=metric, + ) + mock_client = MagicMock() + with ( + patch.object(ValkeyStore, "_create_client", return_value=mock_client), + patch.object(ValkeyStore, "_index_exists", return_value=True), + ): + store = ValkeyStore( + vector_store_config=config, + name="test", + embedding_fn=mock_embedding_fn, + ) + store._client = mock_client + return store + + def test_cosine_score(self, mock_embedding_fn): + """COSINE: score = 1.0 - distance.""" + store = self._make_store_with_metric("COSINE", mock_embedding_fn) + doc = {"content": "x", "metadata": "{}", "chunk_id": "c1", "score": "0.3"} + chunk = store._doc_to_chunk(doc) + assert chunk.score == pytest.approx(0.7) + + def test_l2_score(self, mock_embedding_fn): + """L2: score = 1/(1+distance).""" + store = self._make_store_with_metric("L2", mock_embedding_fn) + doc = {"content": "x", "metadata": "{}", "chunk_id": "c1", "score": "4.0"} + chunk = store._doc_to_chunk(doc) + assert chunk.score == pytest.approx(0.2) # 1/(1+4) + + def test_l2_score_zero_distance(self, mock_embedding_fn): + """L2: identical vectors have distance 0, score = 1.0.""" + store = self._make_store_with_metric("L2", mock_embedding_fn) + doc = {"content": "x", "metadata": "{}", "chunk_id": "c1", "score": "0.0"} + chunk = store._doc_to_chunk(doc) + assert chunk.score == pytest.approx(1.0) + + def test_ip_score(self, mock_embedding_fn): + """IP: score = 1.0 + distance (distance is negative inner product).""" + store = self._make_store_with_metric("IP", mock_embedding_fn) + doc = {"content": "x", "metadata": "{}", "chunk_id": "c1", "score": "-0.8"} + chunk = store._doc_to_chunk(doc) + assert chunk.score == pytest.approx(0.2) # 1.0 + (-0.8) + + +# --------------------------------------------------------------------------- +# Tag value escaping (pure function) +# --------------------------------------------------------------------------- + + +class TestEscapeTagValue: + """Tests for tag value escaping.""" + + def test_simple_value(self): + """Test that simple values pass through.""" + assert _escape_tag_value("hello") == "hello" + + def test_space(self): + """Test that spaces are escaped.""" + assert _escape_tag_value("hello world") == "hello\\ world" + + def test_dot_and_comma(self): + """Test escaping dots and commas.""" + assert _escape_tag_value("a.b,c") == "a\\.b\\,c" + + def test_special_chars_comprehensive(self): + """Test all special chars that need escaping.""" + # Braces are special in Valkey tag queries + assert _escape_tag_value("a{b}c") == "a\\{b\\}c" + assert _escape_tag_value("x@y") == "x\\@y" + assert _escape_tag_value("path/to") == "path\\/to" + + def test_pipe_escaped(self): + """Test that pipe character is escaped to prevent OR injection.""" + assert _escape_tag_value("foo|bar") == "foo\\|bar" + + def test_empty_string(self): + """Test empty string passes through.""" + assert _escape_tag_value("") == "" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store_integration.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store_integration.py new file mode 100644 index 000000000..3b3a1f846 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/tests/test_valkey_store_integration.py @@ -0,0 +1,664 @@ +"""Integration tests for ValkeyStore. + +These tests require a running Valkey server with the valkey-search module loaded. +Skip if Valkey is not available. + +To run locally:: + + docker run -d --name valkey -p 6379:6379 valkey/valkey:latest \\ + --loadmodule /usr/lib/valkey/modules/valkey-search.so + + pytest -v -m integration +""" + +from __future__ import annotations + +import os +import uuid +from typing import List + +import pytest + +from dbgpt.core import Chunk +from dbgpt.storage.vector_store.filters import ( + FilterOperator, + MetadataFilter, + MetadataFilters, +) + +pytestmark = pytest.mark.integration + +VALKEY_HOST = os.environ.get("VALKEY_HOST", "localhost") +VALKEY_PORT = int(os.environ.get("VALKEY_PORT", "6379")) +VALKEY_PASSWORD = os.environ.get("VALKEY_PASSWORD", None) + + +def _valkey_available() -> bool: + """Check if Valkey is available and has the search module.""" + try: + import asyncio + + from glide import GlideClient, GlideClientConfiguration, NodeAddress + + async def _check(): + config = GlideClientConfiguration( + addresses=[NodeAddress(host=VALKEY_HOST, port=VALKEY_PORT)] + ) + client = await GlideClient.create(config) + result = await client.custom_command(["MODULE", "LIST"]) + await client.close() + if result: + for module in result: + if isinstance(module, dict): + name = module.get(b"name", b"") + if name == b"search": + return True + return False + + return asyncio.run(_check()) + except Exception: + return False + + +if not _valkey_available(): + pytest.skip( + "Valkey server with search module not available", allow_module_level=True + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class MockEmbeddings: + """Deterministic embeddings for integration testing.""" + + def __init__(self, dim: int = 128): + self.dim = dim + + def embed_query(self, text: str) -> List[float]: + """Generate a deterministic embedding from text.""" + import hashlib + + h = hashlib.sha256(text.encode()).digest() + vector = [] + for i in range(self.dim): + byte_idx = i % len(h) + vector.append((h[byte_idx] - 128) / 128.0) + return vector + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + """Generate embeddings for multiple documents.""" + return [self.embed_query(text) for text in texts] + + +def _unique_name(prefix: str = "test") -> str: + return f"{prefix}_{uuid.uuid4().hex[:8]}" + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def collection_name(): + """Generate a unique collection name for test isolation.""" + return _unique_name() + + +@pytest.fixture +def valkey_store(collection_name): + """Create a ValkeyStore for integration testing (COSINE, HNSW).""" + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + config = ValkeyVectorConfig( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + index_type="HNSW", + distance_metric="COSINE", + key_prefix="inttest_vec:", + metadata_schema={"source": "tag", "page": "numeric"}, + ) + + store = ValkeyStore( + vector_store_config=config, + name=collection_name, + embedding_fn=MockEmbeddings(dim=128), + ) + + yield store + + try: + store.delete_vector_name(collection_name) + except Exception: + pass + store.close() + + +# --------------------------------------------------------------------------- +# Tests: create_collection +# --------------------------------------------------------------------------- + + +class TestCreateCollection: + """Tests for index creation.""" + + def test_create_collection(self, valkey_store): + """Test creating a vector index.""" + valkey_store.create_collection(valkey_store._collection_name) + assert valkey_store._index_exists() + + def test_create_collection_idempotent(self, valkey_store): + """Test that calling create_collection twice does not error.""" + valkey_store.create_collection(valkey_store._collection_name) + valkey_store.create_collection(valkey_store._collection_name) + assert valkey_store._index_exists() + + +# --------------------------------------------------------------------------- +# Tests: load_document +# --------------------------------------------------------------------------- + + +class TestLoadDocument: + """Tests for document loading.""" + + def test_load_document_returns_ids(self, valkey_store): + """Test that load_document returns chunk IDs.""" + chunks = [ + Chunk(content="doc one", metadata={"source": "test"}, chunk_id="d1"), + Chunk(content="doc two", metadata={"source": "test"}, chunk_id="d2"), + ] + ids = valkey_store.load_document(chunks) + assert ids == ["d1", "d2"] + + def test_load_document_empty_list(self, valkey_store): + """Test loading an empty list.""" + ids = valkey_store.load_document([]) + assert ids == [] + + def test_load_document_with_metadata_fields(self, valkey_store): + """Test that metadata fields are stored for indexing.""" + chunks = [ + Chunk( + content="Python guide", + metadata={"source": "wiki", "page": 5}, + chunk_id="m1", + ), + ] + valkey_store.load_document(chunks) + # Verify the data is searchable + assert valkey_store.vector_name_exists() is True + + +# --------------------------------------------------------------------------- +# Tests: similar_search +# --------------------------------------------------------------------------- + + +class TestSimilarSearch: + """Tests for vector similarity search.""" + + def test_similar_search_returns_results(self, valkey_store): + """Test basic search returns results.""" + chunks = [ + Chunk( + content="Python is a programming language", + metadata={"source": "wiki"}, + chunk_id="s1", + ), + Chunk( + content="Java is also a programming language", + metadata={"source": "wiki"}, + chunk_id="s2", + ), + Chunk( + content="The weather is sunny today", + metadata={"source": "news"}, + chunk_id="s3", + ), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search("programming language", topk=2) + assert len(results) <= 2 + assert len(results) > 0 + # Programming-related chunks should rank higher + contents = [r.content.lower() for r in results] + assert any("programming" in c for c in contents) + + def test_similar_search_topk_limit(self, valkey_store): + """Test that topk limits results.""" + chunks = [ + Chunk(content=f"document number {i}", metadata={}, chunk_id=f"t{i}") + for i in range(5) + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search("document", topk=2) + assert len(results) <= 2 + + def test_similar_search_has_scores(self, valkey_store): + """Test that results have valid scores.""" + chunks = [ + Chunk(content="exact query text", metadata={}, chunk_id="sc1"), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search("exact query text", topk=1) + assert len(results) == 1 + # Score should be high for exact match (COSINE similarity) + assert results[0].score > 0.5 + + def test_similar_search_preserves_metadata(self, valkey_store): + """Test that search results include metadata.""" + chunks = [ + Chunk( + content="metadata test", + metadata={"source": "unit_test", "page": 42}, + chunk_id="meta1", + ), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search("metadata test", topk=1) + assert len(results) == 1 + assert results[0].metadata["source"] == "unit_test" + assert results[0].metadata["page"] == 42 + + def test_similar_search_preserves_chunk_id(self, valkey_store): + """Test that search results include chunk_id.""" + chunks = [ + Chunk(content="id test", metadata={}, chunk_id="myid123"), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search("id test", topk=1) + assert len(results) == 1 + assert results[0].chunk_id == "myid123" + + +# --------------------------------------------------------------------------- +# Tests: similar_search_with_scores +# --------------------------------------------------------------------------- + + +class TestSimilarSearchWithScores: + """Tests for score-threshold filtering.""" + + def test_high_threshold_filters_low_scores(self, valkey_store): + """Test that high threshold filters out dissimilar results.""" + chunks = [ + Chunk(content="exact match phrase", metadata={}, chunk_id="h1"), + Chunk(content="completely unrelated xyz abc", metadata={}, chunk_id="h2"), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search_with_scores( + "exact match phrase", topk=5, score_threshold=0.8 + ) + if results: + assert all(r.score >= 0.8 for r in results) + + def test_zero_threshold_returns_all(self, valkey_store): + """Test that threshold=0 does not filter positive-score results.""" + chunks = [ + Chunk(content="search query text", metadata={}, chunk_id="z1"), + Chunk(content="search query text again", metadata={}, chunk_id="z2"), + ] + valkey_store.load_document(chunks) + + results = valkey_store.similar_search_with_scores( + "search query text", topk=5, score_threshold=0.0 + ) + # Both chunks are very similar to query, scores should be >= 0 + assert len(results) >= 1 + + +# --------------------------------------------------------------------------- +# Tests: vector_name_exists +# --------------------------------------------------------------------------- + + +class TestVectorNameExists: + """Tests for checking index existence.""" + + def test_exists_false_before_load(self, valkey_store): + """Test that exists is False before any data is loaded.""" + valkey_store.create_collection(valkey_store._collection_name) + assert valkey_store.vector_name_exists() is False + + def test_exists_true_after_load(self, valkey_store): + """Test that exists is True after loading data.""" + chunks = [Chunk(content="exists test", metadata={}, chunk_id="e1")] + valkey_store.load_document(chunks) + assert valkey_store.vector_name_exists() is True + + def test_exists_false_after_delete(self, valkey_store): + """Test that exists is False after deleting the index.""" + chunks = [Chunk(content="delete test", metadata={}, chunk_id="ed1")] + valkey_store.load_document(chunks) + valkey_store.delete_vector_name(valkey_store._collection_name) + assert valkey_store.vector_name_exists() is False + + +# --------------------------------------------------------------------------- +# Tests: delete_by_ids +# --------------------------------------------------------------------------- + + +class TestDeleteByIds: + """Tests for deleting specific documents.""" + + def test_delete_single_id(self, valkey_store): + """Test deleting a single document.""" + chunks = [ + Chunk(content="keep this", metadata={}, chunk_id="keep1"), + Chunk(content="delete this", metadata={}, chunk_id="del1"), + ] + valkey_store.load_document(chunks) + + deleted = valkey_store.delete_by_ids("del1") + assert deleted == ["del1"] + + def test_delete_multiple_ids(self, valkey_store): + """Test deleting multiple documents.""" + chunks = [ + Chunk(content=f"doc {i}", metadata={}, chunk_id=f"multi{i}") + for i in range(4) + ] + valkey_store.load_document(chunks) + + deleted = valkey_store.delete_by_ids("multi0, multi1, multi2") + assert deleted == ["multi0", "multi1", "multi2"] + + def test_delete_nonexistent_id(self, valkey_store): + """Test deleting an ID that doesn't exist doesn't error.""" + valkey_store.create_collection(valkey_store._collection_name) + deleted = valkey_store.delete_by_ids("nonexistent_id") + assert deleted == ["nonexistent_id"] + + +# --------------------------------------------------------------------------- +# Tests: delete_vector_name +# --------------------------------------------------------------------------- + + +class TestDeleteVectorName: + """Tests for deleting the entire index.""" + + def test_delete_removes_index(self, valkey_store): + """Test that delete removes the index.""" + chunks = [Chunk(content="to delete", metadata={}, chunk_id="dv1")] + valkey_store.load_document(chunks) + + result = valkey_store.delete_vector_name(valkey_store._collection_name) + assert result is True + assert valkey_store._index_exists() is False + + def test_delete_nonexistent_index(self, valkey_store): + """Test deleting when no index exists doesn't error.""" + result = valkey_store.delete_vector_name("nonexistent") + assert result is True + + +# --------------------------------------------------------------------------- +# Tests: truncate +# --------------------------------------------------------------------------- + + +class TestTruncate: + """Tests for truncating data.""" + + def test_truncate_removes_keys(self, valkey_store): + """Test that truncate removes all data keys.""" + chunks = [ + Chunk(content="data 1", metadata={}, chunk_id="tr1"), + Chunk(content="data 2", metadata={}, chunk_id="tr2"), + ] + valkey_store.load_document(chunks) + + deleted = valkey_store.truncate() + assert len(deleted) >= 2 + + def test_truncate_empty_collection(self, valkey_store): + """Test truncating when no data exists.""" + valkey_store.create_collection(valkey_store._collection_name) + deleted = valkey_store.truncate() + assert deleted == [] + + +# --------------------------------------------------------------------------- +# Tests: metadata filtering +# --------------------------------------------------------------------------- + + +class TestMetadataFiltering: + """Tests for metadata-based filtering during search.""" + + @pytest.fixture + def store_with_data(self, collection_name): + """Create a store with diverse metadata for filter testing.""" + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + config = ValkeyVectorConfig( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + index_type="HNSW", + distance_metric="COSINE", + key_prefix="inttest_filter:", + metadata_schema={"source": "tag", "page": "numeric"}, + ) + + store = ValkeyStore( + vector_store_config=config, + name=collection_name + "_filter", + embedding_fn=MockEmbeddings(dim=128), + ) + + chunks = [ + Chunk( + content="Python programming tutorial", + metadata={"source": "wiki", "page": 1}, + chunk_id="ft1", + ), + Chunk( + content="Java programming tutorial", + metadata={"source": "wiki", "page": 2}, + chunk_id="ft2", + ), + Chunk( + content="Weather forecast report", + metadata={"source": "news", "page": 1}, + chunk_id="ft3", + ), + ] + store.load_document(chunks) + + yield store + + try: + store.delete_vector_name(collection_name + "_filter") + except Exception: + pass + store.close() + + def test_filter_by_tag_eq(self, store_with_data): + """Test filtering by tag equality.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="wiki") + ] + ) + results = store_with_data.similar_search("programming", topk=5, filters=filters) + assert len(results) > 0 + for r in results: + assert r.metadata.get("source") == "wiki" + + def test_filter_by_numeric_gt(self, store_with_data): + """Test filtering by numeric greater-than.""" + filters = MetadataFilters( + filters=[MetadataFilter(key="page", operator=FilterOperator.GT, value=1)] + ) + results = store_with_data.similar_search("programming", topk=5, filters=filters) + if results: + for r in results: + assert r.metadata.get("page", 0) > 1 + + def test_filter_excludes_non_matching(self, store_with_data): + """Test that filter actually excludes non-matching docs.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="news") + ] + ) + results = store_with_data.similar_search("programming", topk=5, filters=filters) + # Only the news doc should match the filter + for r in results: + assert r.metadata.get("source") == "news" + + +# --------------------------------------------------------------------------- +# Tests: FLAT index type +# --------------------------------------------------------------------------- + + +class TestFlatIndex: + """Tests for FLAT index type.""" + + def test_flat_index_load_and_search(self, collection_name): + """Test with FLAT index type.""" + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + config = ValkeyVectorConfig( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + index_type="FLAT", + distance_metric="L2", + key_prefix="inttest_flat:", + ) + + store = ValkeyStore( + vector_store_config=config, + name=collection_name + "_flat", + embedding_fn=MockEmbeddings(dim=128), + ) + + try: + chunks = [ + Chunk(content="flat index test", metadata={}, chunk_id="f1"), + Chunk(content="another flat doc", metadata={}, chunk_id="f2"), + ] + ids = store.load_document(chunks) + assert ids == ["f1", "f2"] + + results = store.similar_search("flat index", topk=1) + assert len(results) >= 1 + # L2 scores should be in (0, 1] range via 1/(1+d) + assert results[0].score > 0 + assert results[0].score <= 1.0 + finally: + try: + store.delete_vector_name(collection_name + "_flat") + except Exception: + pass + store.close() + + +# --------------------------------------------------------------------------- +# Tests: close() +# --------------------------------------------------------------------------- + + +class TestClose: + """Tests for resource cleanup.""" + + def test_close_does_not_error(self, collection_name): + """Test that close() can be called without error.""" + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + config = ValkeyVectorConfig( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + key_prefix="inttest_close:", + ) + + store = ValkeyStore( + vector_store_config=config, + name=collection_name + "_close", + embedding_fn=MockEmbeddings(dim=128), + ) + + # Force client creation + _ = store.client + store.close() + assert store._client is None + assert store._loop.is_closed() + + def test_close_without_client(self, collection_name): + """Test that close() works even if client was never created.""" + from dbgpt_ext.storage.vector_store.valkey_store import ( + ValkeyStore, + ValkeyVectorConfig, + ) + + config = ValkeyVectorConfig( + host=VALKEY_HOST, + port=VALKEY_PORT, + password=VALKEY_PASSWORD, + key_prefix="inttest_close2:", + ) + + store = ValkeyStore( + vector_store_config=config, + name=collection_name + "_close2", + embedding_fn=MockEmbeddings(dim=128), + ) + + # Don't create client, just close + store.close() + assert store._client is None + + +# --------------------------------------------------------------------------- +# Tests: convert_metadata_filters (public method) +# --------------------------------------------------------------------------- + + +class TestConvertMetadataFilters: + """Tests for the public convert_metadata_filters method.""" + + def test_convert_returns_filter_string(self, valkey_store): + """Test that convert_metadata_filters returns a valid string.""" + filters = MetadataFilters( + filters=[ + MetadataFilter(key="source", operator=FilterOperator.EQ, value="wiki") + ] + ) + result = valkey_store.convert_metadata_filters(filters) + assert isinstance(result, str) + assert "@meta_source:{wiki}" in result + + def test_convert_empty_filters(self, valkey_store): + """Test converting empty filters.""" + filters = MetadataFilters(filters=[]) + result = valkey_store.convert_metadata_filters(filters) + assert result == "*" diff --git a/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/valkey_store.py b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/valkey_store.py new file mode 100644 index 000000000..351a4a6f3 --- /dev/null +++ b/packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/valkey_store.py @@ -0,0 +1,868 @@ +"""Valkey vector store. + +Requires a Valkey server with the valkey-search module loaded for vector +similarity search. Uses the valkey-glide client library. + +To run Valkey with the search module:: + + docker run -d --name valkey -p 6379:6379 valkey/valkey:latest \\ + --loadmodule /usr/lib/valkey/modules/valkey-search.so +""" + +from __future__ import annotations + +import json +import logging +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from dbgpt.core import Chunk, Embeddings +from dbgpt.core.awel.flow import Parameter, ResourceCategory, register_resource +from dbgpt.storage.vector_store.base import ( + _COMMON_PARAMETERS, + _VECTOR_STORE_COMMON_PARAMETERS, + VectorStoreBase, + VectorStoreConfig, +) +from dbgpt.storage.vector_store.filters import ( + FilterCondition, + FilterOperator, + MetadataFilters, +) +from dbgpt.util.i18n_utils import _ + +logger = logging.getLogger(__name__) + +_VALKEY_DEFAULT_INDEX_TYPE = "HNSW" +_VALKEY_DEFAULT_DISTANCE_METRIC = "COSINE" +_VALKEY_DEFAULT_KEY_PREFIX = "dbgpt_vec:" +_VALKEY_VECTOR_FIELD = "vector" +_VALKEY_CONTENT_FIELD = "content" +_VALKEY_METADATA_FIELD = "metadata" +_VALKEY_CHUNK_ID_FIELD = "chunk_id" +_VALKEY_METADATA_PREFIX = "meta_" + + +@register_resource( + _("Valkey Config"), + "valkey_vector_config", + category=ResourceCategory.VECTOR_STORE, + description=_("Valkey vector store config."), + parameters=[ + *_COMMON_PARAMETERS, + Parameter.build_from( + _("Host"), + "host", + str, + description=_("The host of the Valkey instance."), + optional=True, + default="localhost", + ), + Parameter.build_from( + _("Port"), + "port", + int, + description=_("The port of the Valkey instance."), + optional=True, + default=6379, + ), + Parameter.build_from( + _("Password"), + "password", + str, + description=_("The password for the Valkey instance."), + optional=True, + default=None, + ), + Parameter.build_from( + _("Use SSL"), + "use_ssl", + bool, + description=_("Whether to use SSL for the Valkey connection."), + optional=True, + default=False, + ), + Parameter.build_from( + _("Index Type"), + "index_type", + str, + description=_( + "The vector index type: 'HNSW' (approximate, fast) or " + "'FLAT' (exact, slower)." + ), + optional=True, + default="HNSW", + ), + Parameter.build_from( + _("Distance Metric"), + "distance_metric", + str, + description=_( + "The distance metric: 'COSINE', 'L2' (Euclidean), or " + "'IP' (Inner Product)." + ), + optional=True, + default="COSINE", + ), + ], +) +@dataclass +class ValkeyVectorConfig(VectorStoreConfig): + """Valkey vector store config.""" + + __type__ = "valkey" + + host: str = field( + default_factory=lambda: os.getenv("VALKEY_HOST", "localhost"), + metadata={"help": _("The host of Valkey store.")}, + ) + port: int = field( + default_factory=lambda: int(os.getenv("VALKEY_PORT", "6379")), + metadata={"help": _("The port of Valkey store.")}, + ) + password: Optional[str] = field( + default_factory=lambda: os.getenv("VALKEY_PASSWORD"), + metadata={"help": _("The password for Valkey store.")}, + ) + use_ssl: bool = field( + default=False, + metadata={"help": _("Whether to use SSL for the Valkey connection.")}, + ) + index_type: str = field( + default_factory=lambda: os.getenv( + "VALKEY_INDEX_TYPE", _VALKEY_DEFAULT_INDEX_TYPE + ), + metadata={ + "help": _( + "The vector index type: 'HNSW' (approximate, fast) or " + "'FLAT' (exact, slower)." + ) + }, + ) + distance_metric: str = field( + default_factory=lambda: os.getenv( + "VALKEY_DISTANCE_METRIC", _VALKEY_DEFAULT_DISTANCE_METRIC + ), + metadata={ + "help": _( + "The distance metric: 'COSINE', 'L2' (Euclidean), or " + "'IP' (Inner Product)." + ) + }, + ) + key_prefix: str = field( + default_factory=lambda: os.getenv( + "VALKEY_KEY_PREFIX", _VALKEY_DEFAULT_KEY_PREFIX + ), + metadata={"help": _("The key prefix for all vector store keys.")}, + ) + hnsw_m: int = field( + default=16, + metadata={"help": _("HNSW: number of connections per node.")}, + ) + hnsw_ef_construction: int = field( + default=200, + metadata={"help": _("HNSW: construction time quality factor.")}, + ) + hnsw_ef_runtime: int = field( + default=10, + metadata={"help": _("HNSW: runtime search quality factor.")}, + ) + request_timeout: int = field( + default_factory=lambda: int(os.getenv("VALKEY_REQUEST_TIMEOUT", "5000")), + metadata={ + "help": _( + "Request timeout in milliseconds for Valkey operations. " + "Prevents indefinite hangs on network issues." + ) + }, + ) + metadata_schema: Optional[Dict[str, str]] = field( + default=None, + metadata={ + "help": _( + "Metadata fields to index for filtering. Dict mapping field name " + "to type: 'tag' (string) or 'numeric'. " + "E.g. {'source': 'tag', 'page': 'numeric'}. " + "Must be defined at store creation time." + ) + }, + ) + + def create_store(self, **kwargs) -> "ValkeyStore": + """Create ValkeyStore.""" + return ValkeyStore(vector_store_config=self, **kwargs) + + +@register_resource( + _("Valkey Vector Store"), + "valkey_vector_store", + category=ResourceCategory.VECTOR_STORE, + description=_("Valkey vector store."), + parameters=[ + Parameter.build_from( + _("Valkey Config"), + "vector_store_config", + ValkeyVectorConfig, + description=_("the valkey config of vector store."), + optional=True, + default=None, + ), + *_VECTOR_STORE_COMMON_PARAMETERS, + ], +) +class ValkeyStore(VectorStoreBase): + """Valkey vector store using valkey-glide client. + + Requires a Valkey server with the valkey-search module loaded. + """ + + def __init__( + self, + vector_store_config: ValkeyVectorConfig, + name: Optional[str] = None, + embedding_fn: Optional[Embeddings] = None, + max_chunks_once_load: Optional[int] = None, + max_threads: Optional[int] = None, + ) -> None: + """Initialize ValkeyStore. + + Args: + vector_store_config: Valkey connection and index configuration. + name: Collection/index name. + embedding_fn: Embedding function for vectorizing text. + max_chunks_once_load: Max chunks per load batch. + max_threads: Max threads for parallel loading. + """ + try: + import glide # noqa: F401 + except ImportError: + raise ImportError( + "Please install valkey-glide: pip install 'valkey-glide>=2.3.0'" + ) + + super().__init__( + max_chunks_once_load=max_chunks_once_load, max_threads=max_threads + ) + + if embedding_fn is None: + raise ValueError("embedding_fn is required for ValkeyStore") + + self._vector_store_config = vector_store_config + self._embedding_fn = embedding_fn + self._collection_name = name or "dbgpt_collection" + self._key_prefix = vector_store_config.key_prefix + self._collection_name + ":" + self._index_name = f"idx:{self._collection_name}" + self._client: Optional[Any] = None + self._dim: Optional[int] = None + + # Dedicated event loop for async glide operations + import asyncio + + self._loop = asyncio.new_event_loop() + + def close(self): + """Close the client connection and event loop.""" + if self._client: + try: + self._loop.run_until_complete(self._client.close()) + except Exception: + pass + self._client = None + if self._loop and not self._loop.is_closed(): + self._loop.close() + + def __del__(self): + """Clean up resources if close() was not called explicitly.""" + if hasattr(self, "_loop"): + self.close() + + @property + def client(self) -> Any: + """Get or create the Valkey client (lazy initialization).""" + if self._client is None: + self._client = self._create_client() + return self._client + + def _create_client(self) -> Any: + """Create a Valkey-glide client.""" + from glide import GlideClient, GlideClientConfiguration, NodeAddress + + config = self._vector_store_config + node = NodeAddress(host=config.host, port=config.port) + + if config.password: + from glide import ServerCredentials + + client_config = GlideClientConfiguration( + addresses=[node], + use_tls=config.use_ssl, + request_timeout=config.request_timeout, + credentials=ServerCredentials(password=config.password), + ) + else: + client_config = GlideClientConfiguration( + addresses=[node], + use_tls=config.use_ssl, + request_timeout=config.request_timeout, + ) + + # GlideClient.create() is async — run it in our dedicated loop + return self._loop.run_until_complete(GlideClient.create(client_config)) + + def _get_dimension(self) -> int: + """Get embedding dimension by running a probe embedding.""" + if self._dim is None: + self._dim = len(self._embedding_fn.embed_query("probe")) + return self._dim + + def _run_async(self, coro): + """Run an async coroutine synchronously using the dedicated event loop.""" + if self._loop.is_closed(): + raise RuntimeError("ValkeyStore has been closed") + return self._loop.run_until_complete(coro) + + def get_config(self) -> ValkeyVectorConfig: + """Get the vector store config.""" + return self._vector_store_config + + def create_collection(self, collection_name: str, **kwargs) -> None: + """Create a vector index in Valkey. + + Args: + collection_name: The name for the collection/index. + """ + from glide import DataType, ft + + # Check if index already exists + if self._index_exists(): + return + + dim = self._get_dimension() + config = self._vector_store_config + + # Build the schema fields + schema = self._build_index_schema(dim, config) + + # Create the index + prefix = self._key_prefix + options = ft.FtCreateOptions( + data_type=DataType.HASH, + prefixes=[prefix], + ) + + self._run_async( + ft.create(self.client, self._index_name, schema=schema, options=options) + ) + logger.info( + f"Created Valkey index '{self._index_name}' with {config.index_type} " + f"algorithm, {config.distance_metric} metric, dim={dim}" + ) + + def _build_index_schema(self, dim: int, config: ValkeyVectorConfig) -> List: + """Build the FT.CREATE schema fields.""" + from glide import ( + DistanceMetricType, + NumericField, + TagField, + TextField, + VectorAlgorithm, + VectorField, + VectorFieldAttributesFlat, + VectorFieldAttributesHnsw, + VectorType, + ) + + # Map string config to enum values + distance_map = { + "COSINE": DistanceMetricType.COSINE, + "L2": DistanceMetricType.L2, + "IP": DistanceMetricType.IP, + } + distance_metric = distance_map.get( + config.distance_metric.upper(), DistanceMetricType.COSINE + ) + + fields = [ + TextField(_VALKEY_CONTENT_FIELD), + TextField(_VALKEY_METADATA_FIELD), + TagField(_VALKEY_CHUNK_ID_FIELD), + ] + + # Add user-defined metadata fields to schema + if config.metadata_schema: + for field_name, field_type in config.metadata_schema.items(): + prefixed = _VALKEY_METADATA_PREFIX + field_name + if field_type.lower() == "numeric": + fields.append(NumericField(prefixed)) + else: + # Default to TAG for string fields + fields.append(TagField(prefixed)) + + # Vector field with algorithm-specific parameters + if config.index_type.upper() == "FLAT": + attributes = VectorFieldAttributesFlat( + dimensions=dim, + distance_metric=distance_metric, + type=VectorType.FLOAT32, + ) + vector_field = VectorField( + _VALKEY_VECTOR_FIELD, + algorithm=VectorAlgorithm.FLAT, + attributes=attributes, + ) + else: + # Default to HNSW + attributes = VectorFieldAttributesHnsw( + dimensions=dim, + distance_metric=distance_metric, + type=VectorType.FLOAT32, + number_of_edges=config.hnsw_m, + vectors_examined_on_construction=config.hnsw_ef_construction, + vectors_examined_on_runtime=config.hnsw_ef_runtime, + ) + vector_field = VectorField( + _VALKEY_VECTOR_FIELD, + algorithm=VectorAlgorithm.HNSW, + attributes=attributes, + ) + + fields.append(vector_field) + return fields + + def _index_exists(self) -> bool: + """Check if the index already exists.""" + from glide import ft + + existing = self._run_async(ft.list(self.client)) + names = { + i.decode() if isinstance(i, bytes) else str(i) for i in (existing or []) + } + return self._index_name in names + + def load_document(self, chunks: List[Chunk]) -> List[str]: + """Load document chunks into Valkey. + + Args: + chunks: Document chunks to store. + + Returns: + List of chunk IDs that were stored. + """ + import struct + + # Ensure index exists + self.create_collection(self._collection_name) + + texts = [chunk.content for chunk in chunks] + vectors = self._embedding_fn.embed_documents(texts) + + for chunk, vector in zip(chunks, vectors): + key = self._key_prefix + chunk.chunk_id + + # Pack vector as binary float32 + vector_bytes = struct.pack(f"{len(vector)}f", *vector) + + # Store as hash fields + field_map = { + _VALKEY_CONTENT_FIELD: chunk.content, + _VALKEY_METADATA_FIELD: json.dumps( + chunk.metadata if chunk.metadata else {} + ), + _VALKEY_CHUNK_ID_FIELD: chunk.chunk_id, + _VALKEY_VECTOR_FIELD: vector_bytes, + } + + # Store indexed metadata fields as top-level hash fields + metadata_schema = self._vector_store_config.metadata_schema + if metadata_schema and chunk.metadata: + for field_name in metadata_schema: + if field_name in chunk.metadata: + prefixed = _VALKEY_METADATA_PREFIX + field_name + field_map[prefixed] = str(chunk.metadata[field_name]) + + self._run_async(self.client.hset(key, field_map)) + + return [chunk.chunk_id for chunk in chunks] + + def similar_search( + self, text: str, topk: int, filters: Optional[MetadataFilters] = None + ) -> List[Chunk]: + """Search for similar documents. + + Args: + text: Query text. + topk: Number of results to return. + filters: Optional metadata filters. + + Returns: + List of similar chunks. + """ + return self._search(text, topk, filters=filters) + + def similar_search_with_scores( + self, + text: str, + topk: int, + score_threshold: float, + filters: Optional[MetadataFilters] = None, + ) -> List[Chunk]: + """Search for similar documents with score filtering. + + Args: + text: Query text. + topk: Number of results to return. + score_threshold: Minimum score threshold. + filters: Optional metadata filters. + + Returns: + List of similar chunks with scores above threshold. + """ + chunks = self._search(text, topk, filters=filters) + return self.filter_by_score_threshold(chunks, score_threshold) + + def _search( + self, + text: str, + topk: int, + filters: Optional[MetadataFilters] = None, + ) -> List[Chunk]: + """Execute a vector similarity search. + + Args: + text: Query text to embed and search. + topk: Number of results. + filters: Optional metadata filters. + + Returns: + List of matching chunks with scores. + """ + import struct + + from glide import FtSearchLimit, FtSearchOptions, ReturnField, ft + + query_vector = self._embedding_fn.embed_query(text) + vector_bytes = struct.pack(f"{len(query_vector)}f", *query_vector) + + # Build the KNN query + filter_expr = self._build_filter_expression(filters) + query_str = f"{filter_expr}=>[KNN {topk} @{_VALKEY_VECTOR_FIELD} $vec AS score]" + + # Execute search + options = FtSearchOptions( + params={"vec": vector_bytes}, + return_fields=[ + ReturnField(_VALKEY_CONTENT_FIELD), + ReturnField(_VALKEY_METADATA_FIELD), + ReturnField(_VALKEY_CHUNK_ID_FIELD), + ReturnField("score"), + ], + limit=FtSearchLimit(0, topk), + ) + + result = self._run_async( + ft.search(self.client, self._index_name, query_str, options) + ) + + return self._parse_search_results(result) + + def _parse_search_results(self, result) -> List[Chunk]: + """Parse FT.SEARCH results into Chunk objects. + + The result format from valkey-glide ft.search is: + [total_count, {key1: {field: value, ...}, key2: {field: value, ...}}] + """ + chunks: List[Chunk] = [] + + if not result: + return chunks + + # Handle structured results (newer glide versions) + if hasattr(result, "results"): + for doc in result.results: + chunk = self._doc_to_chunk(doc) + if chunk: + chunks.append(chunk) + return chunks + + # Handle list-based results: [total_count, {key: {fields}, ...}] + if not isinstance(result, (list, tuple)) or len(result) < 2: + return chunks + + docs = result[1] + if isinstance(docs, dict): + # Format: {b'key1': {b'field': b'value', ...}, ...} + for key, fields in docs.items(): + if isinstance(fields, dict): + chunk = self._doc_to_chunk(fields) + if chunk: + chunks.append(chunk) + elif isinstance(docs, list): + for entry in docs: + chunk = self._doc_to_chunk(entry) + if chunk: + chunks.append(chunk) + + return chunks + + def _doc_to_chunk(self, doc) -> Optional[Chunk]: + """Convert a search result document to a Chunk.""" + try: + if hasattr(doc, "fields"): + fields = doc.fields + elif isinstance(doc, dict): + # Normalize bytes keys to strings + fields = {} + for k, v in doc.items(): + key = k.decode() if isinstance(k, bytes) else k + fields[key] = v + elif isinstance(doc, (list, tuple)): + # Convert flat list [field, value, field, value, ...] to dict + fields = {} + for j in range(0, len(doc), 2): + key = doc[j] if isinstance(doc[j], str) else doc[j].decode() + fields[key] = doc[j + 1] + else: + return None + + content = fields.get(_VALKEY_CONTENT_FIELD, "") + if isinstance(content, bytes): + content = content.decode() + + metadata_raw = fields.get(_VALKEY_METADATA_FIELD, "{}") + if isinstance(metadata_raw, bytes): + metadata_raw = metadata_raw.decode() + metadata = json.loads(metadata_raw) if metadata_raw else {} + + chunk_id = fields.get(_VALKEY_CHUNK_ID_FIELD, "") + if isinstance(chunk_id, bytes): + chunk_id = chunk_id.decode() + + score_raw = fields.get("score", 0.0) + if isinstance(score_raw, bytes): + score_raw = score_raw.decode() + # Convert distance to similarity score based on metric + distance = float(score_raw) + metric = self._vector_store_config.distance_metric.upper() + if metric == "COSINE": + # COSINE distance range: 0 (identical) to 2 (opposite) + score = 1.0 - distance + elif metric == "IP": + # Inner Product: assumes normalized vectors; clamp for safety + score = max(0.0, min(1.0, 1.0 + distance)) + else: + # L2 (Euclidean): distance >= 0, convert via 1/(1+d) + score = 1.0 / (1.0 + distance) + + return Chunk( + content=content, + metadata=metadata, + chunk_id=chunk_id, + score=score, + ) + except Exception as e: + logger.warning(f"Failed to parse search result document: {e}") + return None + + def _build_filter_expression(self, filters: Optional[MetadataFilters]) -> str: + """Build a Valkey search filter expression. + + Args: + filters: Metadata filters to convert. + + Returns: + Filter expression string for FT.SEARCH query. + """ + if not filters or not filters.filters: + return "*" + + if not self._vector_store_config.metadata_schema: + raise ValueError( + "Metadata filters provided but no metadata_schema configured. " + "Configure metadata_schema in ValkeyVectorConfig to enable filtering." + ) + + expressions = [] + for f in filters.filters: + expr = self._single_filter_to_expr(f) + if expr: + expressions.append(expr) + + if not expressions: + return "*" + + if filters.condition == FilterCondition.OR: + return " | ".join(expressions) + else: + # AND condition + return " ".join(expressions) + + def _single_filter_to_expr(self, f) -> Optional[str]: + """Convert a single MetadataFilter to a Valkey search expression.""" + key = _VALKEY_METADATA_PREFIX + f.key + op = f.operator + val = f.value + + if op == FilterOperator.EQ: + if isinstance(val, (int, float)): + return f"@{key}:[{val} {val}]" + else: + # Tag field exact match + return f"@{key}:{{{_escape_tag_value(str(val))}}}" + elif op == FilterOperator.GT: + return f"@{key}:[({val} +inf]" + elif op == FilterOperator.GTE: + return f"@{key}:[{val} +inf]" + elif op == FilterOperator.LT: + return f"@{key}:[-inf ({val}]" + elif op == FilterOperator.LTE: + return f"@{key}:[-inf {val}]" + elif op == FilterOperator.NE: + if isinstance(val, (int, float)): + return f"-@{key}:[{val} {val}]" + else: + return f"-@{key}:{{{_escape_tag_value(str(val))}}}" + elif op == FilterOperator.IN: + if isinstance(val, list): + escaped = "|".join(_escape_tag_value(str(v)) for v in val) + return f"@{key}:{{{escaped}}}" + return None + elif op == FilterOperator.NIN: + if isinstance(val, list): + escaped = "|".join(_escape_tag_value(str(v)) for v in val) + return f"-@{key}:{{{escaped}}}" + return None + else: + logger.warning(f"Unsupported filter operator for Valkey: {op}") + return None + + def vector_name_exists(self) -> bool: + """Check whether the vector index exists and has data.""" + try: + if not self._index_exists(): + return False + # Use FT.INFO to get the number of documents in the index + from glide import ft + + info = self._run_async(ft.info(self.client, self._index_name)) + # info is a dict with index stats + if isinstance(info, dict): + num_docs = info.get(b"num_docs", info.get("num_docs", 0)) + if isinstance(num_docs, bytes): + num_docs = num_docs.decode() + return int(num_docs) > 0 + return False + except Exception as e: + logger.error(f"vector_name_exists error: {e}") + return False + + def delete_vector_name(self, vector_name: str) -> bool: + """Delete the vector index and all associated keys. + + Args: + vector_name: The name of the vector index to delete. + + Returns: + True if deletion was successful. + """ + from glide import ft + + logger.info(f"Deleting Valkey vector index: {self._index_name}") + try: + # Drop the index (does not delete the underlying hash keys) + self._run_async(ft.dropindex(self.client, self._index_name)) + except Exception as e: + logger.warning(f"Error dropping index: {e}") + + # Also delete any remaining keys with the prefix + self._delete_keys_with_prefix(self._key_prefix) + return True + + def delete_by_ids(self, ids: str) -> List[str]: + """Delete vectors by their IDs. + + Args: + ids: Comma-separated string of chunk IDs to delete. + + Returns: + List of deleted chunk IDs. + """ + id_list = [i.strip() for i in ids.split(",") if i.strip()] + for chunk_id in id_list: + key = self._key_prefix + chunk_id + self._run_async(self.client.delete([key])) + return id_list + + def truncate(self) -> List[str]: + """Truncate all data in the collection. + + Returns: + List of deleted key names. + """ + logger.info(f"Truncating Valkey collection: {self._collection_name}") + deleted = self._delete_keys_with_prefix(self._key_prefix) + return deleted + + def _delete_keys_with_prefix(self, prefix: str) -> List[str]: + """Delete all keys matching a prefix using SCAN. + + Args: + prefix: Key prefix to match. + + Returns: + List of deleted key names. + """ + deleted_keys = [] + cursor = "0" + while True: + # Use SCAN to find keys with prefix + result = self._run_async( + self.client.custom_command( + ["SCAN", cursor, "MATCH", f"{prefix}*", "COUNT", "100"] + ) + ) + if isinstance(result, (list, tuple)) and len(result) == 2: + cursor = result[0] + if isinstance(cursor, bytes): + cursor = cursor.decode() + keys = result[1] + if keys: + key_list = [k.decode() if isinstance(k, bytes) else k for k in keys] + self._run_async(self.client.delete(key_list)) + deleted_keys.extend(key_list) + else: + break + if str(cursor) == "0": + break + return deleted_keys + + def convert_metadata_filters(self, filters: MetadataFilters) -> str: + """Convert metadata filters to Valkey search filter expression. + + Args: + filters: Metadata filters. + + Returns: + Filter expression string. + """ + return self._build_filter_expression(filters) + + +def _escape_tag_value(value: str) -> str: + """Escape special characters in tag values for Valkey search. + + Args: + value: The tag value to escape. + + Returns: + Escaped tag value safe for use in FT.SEARCH queries. + """ + special_chars = r".,<>{}[]\"':;!@#$%^&*()-+=~/ |" + escaped = "" + for char in value: + if char in special_chars: + escaped += f"\\{char}" + else: + escaped += char + return escaped