feat: enhance vector database configuration handling and improve provider initialization

This commit is contained in:
shiyu
2026-04-10 19:40:41 +08:00
parent 0609cf6971
commit 398dbcf8ae
5 changed files with 65 additions and 23 deletions
+5 -2
View File
@@ -1,3 +1,4 @@
import asyncio
from pathlib import Path
from typing import Any, Dict, List, Optional
@@ -23,12 +24,14 @@ class MilvusLiteProvider(BaseVectorProvider):
def __init__(self, config: Dict[str, Any] | None = None):
super().__init__(config)
self.db_path = Path(self.config.get("db_path") or "data/db/milvus.db")
raw_db_path = self.config.get("db_path")
db_path = str(raw_db_path).strip() if raw_db_path is not None else ""
self.db_path = Path(db_path or "data/db/milvus.db")
self.client: MilvusClient | None = None
async def initialize(self) -> None:
try:
self.client = MilvusClient(str(self.db_path))
self.client = await asyncio.to_thread(MilvusClient, str(self.db_path))
except Exception as exc: # pragma: no cover - depends on local environment
raise RuntimeError(f"Failed to open Milvus Lite at {self.db_path}: {exc}") from exc
+6 -2
View File
@@ -1,3 +1,4 @@
import asyncio
from typing import Any, Dict, List, Optional
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
@@ -32,11 +33,14 @@ class MilvusServerProvider(BaseVectorProvider):
self.client: MilvusClient | None = None
async def initialize(self) -> None:
uri = self.config.get("uri")
uri = str(self.config.get("uri") or "").strip()
if not uri:
raise RuntimeError("Milvus Server URI is required")
token = self.config.get("token")
if isinstance(token, str):
token = token.strip() or None
try:
self.client = MilvusClient(uri=uri, token=self.config.get("token"))
self.client = await asyncio.to_thread(MilvusClient, uri=uri, token=token)
except Exception as exc: # pragma: no cover - depends on remote availability
raise RuntimeError(f"Failed to connect to Milvus Server {uri}: {exc}") from exc
+2 -1
View File
@@ -1,3 +1,4 @@
import asyncio
from typing import Any, Dict, List, Optional, Sequence
from uuid import NAMESPACE_URL, uuid5
@@ -40,7 +41,7 @@ class QdrantProvider(BaseVectorProvider):
api_key = (self.config.get("api_key") or None) or None
try:
client = QdrantClient(url=url, api_key=api_key)
client.get_collections()
await asyncio.to_thread(client.get_collections)
self.client = client
except Exception as exc: # pragma: no cover - 依赖外部服务
raise RuntimeError(f"Failed to connect to Qdrant at {url}: {exc}") from exc