mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-07-17 10:16:49 +00:00
fix(config): allow Milvus vector store type in TOML configuration (#3127)
## Summary Fixes #3041 When users set `type = "milvus"` in their TOML vector storage config, the server fails to start with: ``` ValueError: Invalid data for ChromaVectorConfig: Unknown type value: milvus, known types: ['chroma', 'elasticsearch', 'pgvector', 'weaviate', 'oceanbase'] ``` ### Root Causes 1. **Wrong type annotation** — `StorageConfig.vector` was annotated as `ChromaVectorConfig` instead of the polymorphic base class `VectorStoreConfig`, preventing the config parser from dispatching to non-chroma vector store types. 2. **Top-level `pymilvus` import** — `milvus_store.py` imported `pymilvus` at module level (line 12). When `pymilvus` is not installed, the entire module fails to load and `MilvusVectorConfig` is never registered in the type registry. ### Fix - **`packages/dbgpt-app/src/dbgpt_app/config.py`**: Change the `vector` field type annotation from `ChromaVectorConfig` to `VectorStoreConfig` (keeping `ChromaVectorConfig()` as the default value for backward compatibility). - **`packages/dbgpt-ext/src/dbgpt_ext/storage/vector_store/milvus_store.py`**: Move the top-level `from pymilvus.milvus_client import IndexParams, MilvusClient` into the methods that actually use them (lazy import), matching the pattern used by all other vector store implementations (chroma, elasticsearch, pgvector, weaviate, etc.). ### Verification Tested with `pymilvus` **not installed** (same environment as the bug reporter): **Before fix** — loading a TOML config with `type = "milvus"`: ``` Error scanning module dbgpt_ext.storage.vector_store.milvus_store: No module named 'pymilvus' ValueError: Invalid data for ChromaVectorConfig: Unknown type value: milvus ``` **After fix** — same config loads successfully: ``` ✅ vector type: MilvusVectorConfig (uri=localhost, port=19530) ``` - `MilvusVectorConfig` registers in the type registry even without `pymilvus` installed - `VectorStoreConfig.get_subclass("milvus")` correctly resolves to `MilvusVectorConfig` - When actually creating a `MilvusStore` without `pymilvus`, a clear error message is shown: `"Please install it with pip install pymilvus"` - All 26 existing storage unit tests pass - `ruff` lint and format checks pass Made with [Cursor](https://cursor.com)
This commit is contained in:
@@ -8,6 +8,7 @@ from dbgpt.model.parameter import (
|
||||
ModelServiceConfig,
|
||||
)
|
||||
from dbgpt.storage.cache.manager import ModelCacheParameters
|
||||
from dbgpt.storage.vector_store.base import VectorStoreConfig
|
||||
from dbgpt.util.configure import HookConfig
|
||||
from dbgpt.util.i18n_utils import _
|
||||
from dbgpt.util.parameter_utils import BaseParameters
|
||||
@@ -55,7 +56,7 @@ class SystemParameters:
|
||||
class StorageConfig(BaseParameters):
|
||||
__cfg_type__ = "app"
|
||||
|
||||
vector: Optional[ChromaVectorConfig] = field(
|
||||
vector: Optional[VectorStoreConfig] = field(
|
||||
default_factory=lambda: ChromaVectorConfig(),
|
||||
metadata={
|
||||
"help": _("default vector type"),
|
||||
|
||||
@@ -9,8 +9,6 @@ import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Iterable, List, Optional
|
||||
|
||||
from pymilvus.milvus_client import IndexParams, MilvusClient
|
||||
|
||||
from dbgpt.core import Chunk, Embeddings
|
||||
from dbgpt.core.awel.flow import Parameter, ResourceCategory, register_resource
|
||||
from dbgpt.storage.vector_store.base import (
|
||||
@@ -289,6 +287,14 @@ class MilvusStore(VectorStoreBase):
|
||||
connect_kwargs["user"] = self.username
|
||||
connect_kwargs["password"] = self.password
|
||||
|
||||
try:
|
||||
from pymilvus.milvus_client import MilvusClient
|
||||
except ImportError:
|
||||
raise ValueError(
|
||||
"Could not import pymilvus python package. "
|
||||
"Please install it with `pip install pymilvus`."
|
||||
)
|
||||
|
||||
url = f"http://{self.uri}:{self.port}"
|
||||
self._milvus_client = MilvusClient(
|
||||
uri=url, user=self.username, db_name="default"
|
||||
@@ -374,6 +380,8 @@ class MilvusStore(VectorStoreBase):
|
||||
)
|
||||
schema.add_function(bm25_fn)
|
||||
# Create the collection
|
||||
from pymilvus.milvus_client import IndexParams
|
||||
|
||||
collection = Collection(collection_name, schema)
|
||||
self.col = collection
|
||||
index_params = IndexParams()
|
||||
|
||||
Reference in New Issue
Block a user