mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-07-16 17:15:22 +00:00
feat(ext): add Valkey vector store integration (#3051)
Signed-off-by: Daria Korenieva <daric2612@gmail.com>
This commit is contained in:
Daria Korenieva
GitHub
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.",
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
---
|
||||
title: "Valkey Config Configuration"
|
||||
description: "Valkey vector store config."
|
||||
---
|
||||
|
||||
import { ConfigDetail } from "@site/src/components/mdx/ConfigDetail";
|
||||
|
||||
<ConfigDetail config={{
|
||||
"name": "ValkeyVectorConfig",
|
||||
"description": "Valkey vector store config.",
|
||||
"documentationUrl": null,
|
||||
"parameters": [
|
||||
{
|
||||
"name": "user",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The user of vector store, if not set, will use the default user."
|
||||
},
|
||||
{
|
||||
"name": "password",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The password for Valkey store."
|
||||
},
|
||||
{
|
||||
"name": "max_chunks_once_load",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "The max chunks once load in vector store, if not set, will use the default value 10."
|
||||
},
|
||||
{
|
||||
"name": "max_threads",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "The max threads in vector store, if not set, will use the default value 1."
|
||||
},
|
||||
{
|
||||
"name": "host",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The host of Valkey store.",
|
||||
"defaultValue": "localhost"
|
||||
},
|
||||
{
|
||||
"name": "port",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "The port of Valkey store.",
|
||||
"defaultValue": "6379"
|
||||
},
|
||||
{
|
||||
"name": "use_ssl",
|
||||
"type": "boolean",
|
||||
"required": false,
|
||||
"description": "Whether to use SSL for the Valkey connection.",
|
||||
"defaultValue": "False"
|
||||
},
|
||||
{
|
||||
"name": "index_type",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The vector index type: 'HNSW' (approximate, fast) or 'FLAT' (exact, slower).",
|
||||
"defaultValue": "HNSW"
|
||||
},
|
||||
{
|
||||
"name": "distance_metric",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The distance metric: 'COSINE', 'L2' (Euclidean), or 'IP' (Inner Product).",
|
||||
"defaultValue": "COSINE"
|
||||
},
|
||||
{
|
||||
"name": "key_prefix",
|
||||
"type": "string",
|
||||
"required": false,
|
||||
"description": "The key prefix for all vector store keys.",
|
||||
"defaultValue": "dbgpt_vec:"
|
||||
},
|
||||
{
|
||||
"name": "hnsw_m",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "HNSW: number of connections per node.",
|
||||
"defaultValue": "16"
|
||||
},
|
||||
{
|
||||
"name": "hnsw_ef_construction",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "HNSW: construction time quality factor.",
|
||||
"defaultValue": "200"
|
||||
},
|
||||
{
|
||||
"name": "hnsw_ef_runtime",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "HNSW: runtime search quality factor.",
|
||||
"defaultValue": "10"
|
||||
},
|
||||
{
|
||||
"name": "request_timeout",
|
||||
"type": "integer",
|
||||
"required": false,
|
||||
"description": "Request timeout in milliseconds for Valkey operations. Prevents indefinite hangs on network issues.",
|
||||
"defaultValue": "5000"
|
||||
},
|
||||
{
|
||||
"name": "metadata_schema",
|
||||
"type": "object",
|
||||
"required": false,
|
||||
"description": "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."
|
||||
}
|
||||
]
|
||||
}} />
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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` |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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` |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for vector store implementations."""
|
||||
@@ -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("") == ""
|
||||
@@ -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 == "*"
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user