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:
chen-alan
2026-07-09 19:34:00 +08:00
committed by GitHub
2 changed files with 12 additions and 3 deletions

View File

@@ -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"),

View File

@@ -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()