mirror of
https://github.com/DrizzleTime/Foxel.git
synced 2026-07-27 17:06:52 +08:00
feat(vector_db): Implement Vector Database Service with multiple providers
This commit is contained in:
11
services/vector_db/__init__.py
Normal file
11
services/vector_db/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
from .service import VectorDBService, DEFAULT_VECTOR_DIMENSION
|
||||
from .providers import list_providers, get_provider_entry
|
||||
from .config_manager import VectorDBConfigManager
|
||||
|
||||
__all__ = [
|
||||
"VectorDBService",
|
||||
"DEFAULT_VECTOR_DIMENSION",
|
||||
"list_providers",
|
||||
"get_provider_entry",
|
||||
"VectorDBConfigManager",
|
||||
]
|
||||
43
services/vector_db/config_manager.py
Normal file
43
services/vector_db/config_manager.py
Normal file
@@ -0,0 +1,43 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from services.config import ConfigCenter
|
||||
|
||||
|
||||
class VectorDBConfigManager:
|
||||
TYPE_KEY = "VECTOR_DB_TYPE"
|
||||
CONFIG_KEY = "VECTOR_DB_CONFIG"
|
||||
DEFAULT_TYPE = "milvus_lite"
|
||||
|
||||
@classmethod
|
||||
async def load_config(cls) -> Tuple[str, Dict[str, Any]]:
|
||||
raw_type = await ConfigCenter.get(cls.TYPE_KEY, cls.DEFAULT_TYPE)
|
||||
provider_type = str(raw_type or cls.DEFAULT_TYPE)
|
||||
|
||||
raw_config = await ConfigCenter.get(cls.CONFIG_KEY)
|
||||
config_dict: Dict[str, Any] = {}
|
||||
if isinstance(raw_config, str) and raw_config:
|
||||
try:
|
||||
config_dict = json.loads(raw_config)
|
||||
except json.JSONDecodeError:
|
||||
config_dict = {}
|
||||
elif isinstance(raw_config, dict):
|
||||
config_dict = raw_config
|
||||
return provider_type, config_dict
|
||||
|
||||
@classmethod
|
||||
async def save_config(cls, provider_type: str, config: Dict[str, Any]) -> None:
|
||||
await ConfigCenter.set(cls.TYPE_KEY, provider_type)
|
||||
await ConfigCenter.set(cls.CONFIG_KEY, json.dumps(config or {}))
|
||||
|
||||
@classmethod
|
||||
async def get_type(cls) -> str:
|
||||
provider_type, _ = await cls.load_config()
|
||||
return provider_type
|
||||
|
||||
@classmethod
|
||||
async def get_config(cls) -> Dict[str, Any]:
|
||||
_, config = await cls.load_config()
|
||||
return config
|
||||
56
services/vector_db/providers/__init__.py
Normal file
56
services/vector_db/providers/__init__.py
Normal file
@@ -0,0 +1,56 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Type
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
from .milvus_lite import MilvusLiteProvider
|
||||
from .milvus_server import MilvusServerProvider
|
||||
from .qdrant import QdrantProvider
|
||||
|
||||
_PROVIDER_REGISTRY: Dict[str, Dict[str, object]] = {
|
||||
MilvusLiteProvider.type: {
|
||||
"class": MilvusLiteProvider,
|
||||
"label": MilvusLiteProvider.label,
|
||||
"description": MilvusLiteProvider.description,
|
||||
"enabled": MilvusLiteProvider.enabled,
|
||||
"config_schema": MilvusLiteProvider.config_schema,
|
||||
},
|
||||
MilvusServerProvider.type: {
|
||||
"class": MilvusServerProvider,
|
||||
"label": MilvusServerProvider.label,
|
||||
"description": MilvusServerProvider.description,
|
||||
"enabled": MilvusServerProvider.enabled,
|
||||
"config_schema": MilvusServerProvider.config_schema,
|
||||
},
|
||||
QdrantProvider.type: {
|
||||
"class": QdrantProvider,
|
||||
"label": QdrantProvider.label,
|
||||
"description": QdrantProvider.description,
|
||||
"enabled": QdrantProvider.enabled,
|
||||
"config_schema": QdrantProvider.config_schema,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def list_providers() -> List[Dict[str, object]]:
|
||||
return [
|
||||
{
|
||||
"type": type_key,
|
||||
"label": meta["label"],
|
||||
"description": meta.get("description"),
|
||||
"enabled": meta.get("enabled", True),
|
||||
"config_schema": meta.get("config_schema", []),
|
||||
}
|
||||
for type_key, meta in _PROVIDER_REGISTRY.items()
|
||||
]
|
||||
|
||||
|
||||
def get_provider_entry(provider_type: str) -> Dict[str, object] | None:
|
||||
return _PROVIDER_REGISTRY.get(provider_type)
|
||||
|
||||
|
||||
def get_provider_class(provider_type: str) -> Type[BaseVectorProvider] | None:
|
||||
entry = get_provider_entry(provider_type)
|
||||
if not entry:
|
||||
return None
|
||||
return entry.get("class") # type: ignore[return-value]
|
||||
41
services/vector_db/providers/base.py
Normal file
41
services/vector_db/providers/base.py
Normal file
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
|
||||
class BaseVectorProvider:
|
||||
"""向量数据库提供者基础类,所有实际实现需继承该类"""
|
||||
|
||||
type: str = ""
|
||||
label: str = ""
|
||||
description: str | None = None
|
||||
enabled: bool = True
|
||||
config_schema: List[Dict[str, Any]] = []
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
self.config = config or {}
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""执行初始化逻辑,例如建立连接"""
|
||||
raise NotImplementedError
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
raise NotImplementedError
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
raise NotImplementedError
|
||||
196
services/vector_db/providers/milvus_lite.py
Normal file
196
services/vector_db/providers/milvus_lite.py
Normal file
@@ -0,0 +1,196 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class MilvusLiteProvider(BaseVectorProvider):
|
||||
type = "milvus_lite"
|
||||
label = "Milvus Lite"
|
||||
description = "Embedded Milvus Lite (local file storage)."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "db_path",
|
||||
"label": "Database file path",
|
||||
"type": "text",
|
||||
"default": "data/db/milvus.db",
|
||||
"required": False,
|
||||
}
|
||||
]
|
||||
|
||||
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")
|
||||
self.client: MilvusClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
try:
|
||||
self.client = 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
|
||||
|
||||
def _get_client(self) -> MilvusClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Milvus Lite client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _to_int(value: Any) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
if client.has_collection(collection_name):
|
||||
return
|
||||
if vector:
|
||||
vector_dim = dim if isinstance(dim, int) and dim > 0 else 0
|
||||
if vector_dim <= 0:
|
||||
vector_dim = 4096
|
||||
fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Image vector collection")
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
index_params = MilvusClient.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="embedding",
|
||||
index_type="IVF_FLAT",
|
||||
index_name="vector_index",
|
||||
metric_type="COSINE",
|
||||
params={"nlist": 64},
|
||||
)
|
||||
client.create_index(collection_name, index_params=index_params)
|
||||
else:
|
||||
fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Simple file index")
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
self._get_client().upsert(collection_name, data)
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
self._get_client().delete(collection_name, ids=[path])
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
search_params = {"metric_type": "COSINE"}
|
||||
return self._get_client().search(
|
||||
collection_name,
|
||||
data=[query_embedding],
|
||||
anns_field="embedding",
|
||||
search_params=search_params,
|
||||
limit=top_k,
|
||||
output_fields=["path"],
|
||||
)
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
filter_expr = f"path like '%{query_path}%'" if query_path else "path like '%%'"
|
||||
results = self._get_client().query(
|
||||
collection_name,
|
||||
filter=filter_expr,
|
||||
limit=top_k,
|
||||
output_fields=["path"],
|
||||
)
|
||||
return [[{"id": r["path"], "distance": 1.0, "entity": {"path": r["path"]}} for r in results]]
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
collection_names = client.list_collections()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to list collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for name in collection_names:
|
||||
try:
|
||||
stats = client.get_collection_stats(name) or {}
|
||||
except Exception:
|
||||
stats = {}
|
||||
row_count = self._to_int(stats.get("row_count"))
|
||||
total_vectors += row_count
|
||||
|
||||
dimension: Optional[int] = None
|
||||
is_vector_collection = False
|
||||
try:
|
||||
description = client.describe_collection(name)
|
||||
except Exception:
|
||||
description = None
|
||||
|
||||
if description:
|
||||
for field in description.get("fields", []):
|
||||
if field.get("type") == DataType.FLOAT_VECTOR:
|
||||
params = field.get("params") or {}
|
||||
dimension = self._to_int(params.get("dim")) or 4096
|
||||
is_vector_collection = True
|
||||
break
|
||||
|
||||
estimated_memory = 0
|
||||
if is_vector_collection and dimension:
|
||||
estimated_memory = row_count * dimension * 4
|
||||
total_estimated_memory += estimated_memory
|
||||
|
||||
indexes: List[Dict[str, Any]] = []
|
||||
try:
|
||||
index_names = client.list_indexes(name) or []
|
||||
except Exception:
|
||||
index_names = []
|
||||
|
||||
for index_name in index_names:
|
||||
try:
|
||||
detail = client.describe_index(name, index_name) or {}
|
||||
except Exception:
|
||||
detail = {}
|
||||
indexes.append(
|
||||
{
|
||||
"index_name": index_name,
|
||||
"index_type": detail.get("index_type"),
|
||||
"metric_type": detail.get("metric_type"),
|
||||
"indexed_rows": self._to_int(detail.get("indexed_rows")),
|
||||
"pending_index_rows": self._to_int(detail.get("pending_index_rows")),
|
||||
"state": detail.get("state"),
|
||||
}
|
||||
)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension if is_vector_collection else None,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": is_vector_collection,
|
||||
"indexes": indexes,
|
||||
}
|
||||
)
|
||||
|
||||
db_file_size = None
|
||||
try:
|
||||
if self.db_path.exists():
|
||||
db_file_size = self.db_path.stat().st_size
|
||||
except OSError:
|
||||
db_file_size = None
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": db_file_size,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
for collection_name in client.list_collections():
|
||||
client.drop_collection(collection_name)
|
||||
197
services/vector_db/providers/milvus_server.py
Normal file
197
services/vector_db/providers/milvus_server.py
Normal file
@@ -0,0 +1,197 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class MilvusServerProvider(BaseVectorProvider):
|
||||
type = "milvus_server"
|
||||
label = "Milvus Server"
|
||||
description = "Remote Milvus instance accessed via URI."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "uri",
|
||||
"label": "Server URI",
|
||||
"type": "text",
|
||||
"required": True,
|
||||
"placeholder": "http://localhost:19530",
|
||||
},
|
||||
{
|
||||
"key": "token",
|
||||
"label": "Token",
|
||||
"type": "password",
|
||||
"required": False,
|
||||
"placeholder": "user:password",
|
||||
},
|
||||
]
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
super().__init__(config)
|
||||
self.client: MilvusClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
uri = self.config.get("uri")
|
||||
if not uri:
|
||||
raise RuntimeError("Milvus Server URI is required")
|
||||
try:
|
||||
self.client = MilvusClient(uri=uri, token=self.config.get("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
|
||||
|
||||
def _get_client(self) -> MilvusClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Milvus Server client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _to_int(value: Any) -> int:
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
if client.has_collection(collection_name):
|
||||
return
|
||||
if vector:
|
||||
vector_dim = dim if isinstance(dim, int) and dim > 0 else 0
|
||||
if vector_dim <= 0:
|
||||
vector_dim = 4096
|
||||
fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Image vector collection")
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
index_params = MilvusClient.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="embedding",
|
||||
index_type="IVF_FLAT",
|
||||
index_name="vector_index",
|
||||
metric_type="COSINE",
|
||||
params={"nlist": 64},
|
||||
)
|
||||
client.create_index(collection_name, index_params=index_params)
|
||||
else:
|
||||
fields = [
|
||||
FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=512, is_primary=True, auto_id=False),
|
||||
]
|
||||
schema = CollectionSchema(fields, description="Simple file index")
|
||||
client.create_collection(collection_name, schema=schema)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
self._get_client().upsert(collection_name, data)
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
self._get_client().delete(collection_name, ids=[path])
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
search_params = {"metric_type": "COSINE"}
|
||||
return self._get_client().search(
|
||||
collection_name,
|
||||
data=[query_embedding],
|
||||
anns_field="embedding",
|
||||
search_params=search_params,
|
||||
limit=top_k,
|
||||
output_fields=["path"],
|
||||
)
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
filter_expr = f"path like '%{query_path}%'" if query_path else "path like '%%'"
|
||||
results = self._get_client().query(
|
||||
collection_name,
|
||||
filter=filter_expr,
|
||||
limit=top_k,
|
||||
output_fields=["path"],
|
||||
)
|
||||
return [[{"id": r["path"], "distance": 1.0, "entity": {"path": r["path"]}} for r in results]]
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
collection_names = client.list_collections()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to list collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for name in collection_names:
|
||||
try:
|
||||
stats = client.get_collection_stats(name) or {}
|
||||
except Exception:
|
||||
stats = {}
|
||||
row_count = self._to_int(stats.get("row_count"))
|
||||
total_vectors += row_count
|
||||
|
||||
dimension: Optional[int] = None
|
||||
is_vector_collection = False
|
||||
try:
|
||||
description = client.describe_collection(name)
|
||||
except Exception:
|
||||
description = None
|
||||
|
||||
if description:
|
||||
for field in description.get("fields", []):
|
||||
if field.get("type") == DataType.FLOAT_VECTOR:
|
||||
params = field.get("params") or {}
|
||||
dimension = self._to_int(params.get("dim")) or 4096
|
||||
is_vector_collection = True
|
||||
break
|
||||
|
||||
estimated_memory = 0
|
||||
if is_vector_collection and dimension:
|
||||
estimated_memory = row_count * dimension * 4
|
||||
total_estimated_memory += estimated_memory
|
||||
|
||||
indexes: List[Dict[str, Any]] = []
|
||||
try:
|
||||
index_names = client.list_indexes(name) or []
|
||||
except Exception:
|
||||
index_names = []
|
||||
|
||||
for index_name in index_names:
|
||||
try:
|
||||
detail = client.describe_index(name, index_name) or {}
|
||||
except Exception:
|
||||
detail = {}
|
||||
indexes.append(
|
||||
{
|
||||
"index_name": index_name,
|
||||
"index_type": detail.get("index_type"),
|
||||
"metric_type": detail.get("metric_type"),
|
||||
"indexed_rows": self._to_int(detail.get("indexed_rows")),
|
||||
"pending_index_rows": self._to_int(detail.get("pending_index_rows")),
|
||||
"state": detail.get("state"),
|
||||
}
|
||||
)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension if is_vector_collection else None,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": is_vector_collection,
|
||||
"indexes": indexes,
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": None,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
for collection_name in client.list_collections():
|
||||
client.drop_collection(collection_name)
|
||||
237
services/vector_db/providers/qdrant.py
Normal file
237
services/vector_db/providers/qdrant.py
Normal file
@@ -0,0 +1,237 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional, Sequence
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qmodels
|
||||
|
||||
from .base import BaseVectorProvider
|
||||
|
||||
|
||||
class QdrantProvider(BaseVectorProvider):
|
||||
type = "qdrant"
|
||||
label = "Qdrant"
|
||||
description = "Qdrant vector database (HTTP API)."
|
||||
enabled = True
|
||||
config_schema: List[Dict[str, Any]] = [
|
||||
{
|
||||
"key": "url",
|
||||
"label": "Server URL",
|
||||
"type": "text",
|
||||
"required": True,
|
||||
"placeholder": "http://localhost:6333",
|
||||
},
|
||||
{
|
||||
"key": "api_key",
|
||||
"label": "API Key",
|
||||
"type": "password",
|
||||
"required": False,
|
||||
},
|
||||
]
|
||||
|
||||
def __init__(self, config: Dict[str, Any] | None = None):
|
||||
super().__init__(config)
|
||||
self.client: Optional[QdrantClient] = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
url = (self.config.get("url") or "").strip()
|
||||
if not url:
|
||||
raise RuntimeError("Qdrant URL is required")
|
||||
|
||||
api_key = (self.config.get("api_key") or None) or None
|
||||
try:
|
||||
client = QdrantClient(url=url, api_key=api_key)
|
||||
# 简单连通性校验
|
||||
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
|
||||
|
||||
def _get_client(self) -> QdrantClient:
|
||||
if not self.client:
|
||||
raise RuntimeError("Qdrant client is not initialized")
|
||||
return self.client
|
||||
|
||||
@staticmethod
|
||||
def _vector_params(vector: bool, dim: int) -> qmodels.VectorParams:
|
||||
size = dim if vector and isinstance(dim, int) and dim > 0 else 1
|
||||
return qmodels.VectorParams(size=size, distance=qmodels.Distance.COSINE)
|
||||
|
||||
def ensure_collection(self, collection_name: str, vector: bool, dim: int) -> None:
|
||||
client = self._get_client()
|
||||
try:
|
||||
if client.collection_exists(collection_name):
|
||||
return
|
||||
except Exception as exc: # pragma: no cover - 依赖外部服务
|
||||
raise RuntimeError(f"Failed to check Qdrant collection '{collection_name}': {exc}") from exc
|
||||
|
||||
vectors_config = self._vector_params(vector, dim)
|
||||
try:
|
||||
client.create_collection(collection_name=collection_name, vectors_config=vectors_config)
|
||||
except Exception as exc: # pragma: no cover
|
||||
if "already exists" in str(exc).lower():
|
||||
return
|
||||
raise RuntimeError(f"Failed to create Qdrant collection '{collection_name}': {exc}") from exc
|
||||
|
||||
@staticmethod
|
||||
def _point_id(path: str) -> str:
|
||||
return str(uuid5(NAMESPACE_URL, path))
|
||||
|
||||
def _prepare_point(self, data: Dict[str, Any]) -> qmodels.PointStruct:
|
||||
path = data.get("path")
|
||||
if not path:
|
||||
raise ValueError("Qdrant upsert requires 'path' in data")
|
||||
|
||||
embedding = data.get("embedding")
|
||||
if embedding is None:
|
||||
vector = [0.0]
|
||||
else:
|
||||
vector = [float(x) for x in embedding]
|
||||
|
||||
payload = {"path": path}
|
||||
return qmodels.PointStruct(id=self._point_id(path), vector=vector, payload=payload)
|
||||
|
||||
def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
client = self._get_client()
|
||||
point = self._prepare_point(data)
|
||||
client.upsert(collection_name=collection_name, wait=True, points=[point])
|
||||
|
||||
def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
client = self._get_client()
|
||||
selector = qmodels.PointIdsList(points=[self._point_id(path)])
|
||||
client.delete(collection_name=collection_name, points_selector=selector, wait=True)
|
||||
|
||||
def _format_search_results(self, points: Sequence[qmodels.ScoredPoint]):
|
||||
return [
|
||||
{
|
||||
"id": point.id,
|
||||
"distance": point.score,
|
||||
"entity": {"path": (point.payload or {}).get("path")},
|
||||
}
|
||||
for point in points
|
||||
]
|
||||
|
||||
def search_vectors(self, collection_name: str, query_embedding, top_k: int):
|
||||
client = self._get_client()
|
||||
vector = [float(x) for x in query_embedding]
|
||||
points = client.search(
|
||||
collection_name=collection_name,
|
||||
query_vector=vector,
|
||||
limit=top_k,
|
||||
with_payload=True,
|
||||
)
|
||||
return [self._format_search_results(points)]
|
||||
|
||||
def search_by_path(self, collection_name: str, query_path: str, top_k: int):
|
||||
client = self._get_client()
|
||||
results: List[Dict[str, Any]] = []
|
||||
offset: Optional[str | int] = None
|
||||
remaining = max(top_k, 1)
|
||||
|
||||
while len(results) < top_k:
|
||||
batch_size = min(max(remaining * 2, 10), 200)
|
||||
records, next_offset = client.scroll(
|
||||
collection_name=collection_name,
|
||||
limit=batch_size,
|
||||
offset=offset,
|
||||
with_payload=True,
|
||||
)
|
||||
if not records:
|
||||
break
|
||||
|
||||
for record in records:
|
||||
path = (record.payload or {}).get("path")
|
||||
if query_path and path:
|
||||
if query_path not in path:
|
||||
continue
|
||||
results.append({"id": record.id, "distance": 1.0, "entity": {"path": path}})
|
||||
if len(results) >= top_k:
|
||||
break
|
||||
|
||||
if next_offset is None or len(results) >= top_k:
|
||||
break
|
||||
offset = next_offset
|
||||
remaining = top_k - len(results)
|
||||
|
||||
return [results]
|
||||
|
||||
def _extract_vector_config(self, vectors) -> Optional[qmodels.VectorParams]:
|
||||
if isinstance(vectors, qmodels.VectorParams):
|
||||
return vectors
|
||||
if isinstance(vectors, dict):
|
||||
for value in vectors.values():
|
||||
if isinstance(value, qmodels.VectorParams):
|
||||
return value
|
||||
return None
|
||||
|
||||
def get_all_stats(self) -> Dict[str, Any]:
|
||||
client = self._get_client()
|
||||
try:
|
||||
response = client.get_collections()
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(f"Failed to list Qdrant collections: {exc}") from exc
|
||||
|
||||
collections: List[Dict[str, Any]] = []
|
||||
total_vectors = 0
|
||||
total_estimated_memory = 0
|
||||
|
||||
for description in response.collections or []:
|
||||
name = description.name
|
||||
try:
|
||||
info = client.get_collection(name)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
row_count = int(info.points_count or 0)
|
||||
total_vectors += row_count
|
||||
|
||||
vector_params = self._extract_vector_config(info.config.params.vectors if info.config and info.config.params else None)
|
||||
dimension = int(vector_params.size) if vector_params and vector_params.size else None
|
||||
estimated_memory = row_count * dimension * 4 if dimension else 0
|
||||
total_estimated_memory += estimated_memory
|
||||
distance = str(vector_params.distance) if vector_params and vector_params.distance else None
|
||||
|
||||
indexed_rows = int(info.indexed_vectors_count or 0)
|
||||
pending_rows = max(row_count - indexed_rows, 0)
|
||||
|
||||
collections.append(
|
||||
{
|
||||
"name": name,
|
||||
"row_count": row_count,
|
||||
"dimension": dimension,
|
||||
"estimated_memory_bytes": estimated_memory,
|
||||
"is_vector_collection": dimension is not None and dimension > 1,
|
||||
"indexes": [
|
||||
{
|
||||
"index_name": "hnsw",
|
||||
"index_type": "HNSW",
|
||||
"metric_type": distance,
|
||||
"indexed_rows": indexed_rows,
|
||||
"pending_index_rows": pending_rows,
|
||||
"state": info.status,
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"collections": collections,
|
||||
"collection_count": len(collections),
|
||||
"total_vectors": total_vectors,
|
||||
"estimated_total_memory_bytes": total_estimated_memory,
|
||||
"db_file_size_bytes": None,
|
||||
}
|
||||
|
||||
def clear_all_data(self) -> None:
|
||||
client = self._get_client()
|
||||
try:
|
||||
response = client.get_collections()
|
||||
except Exception as exc: # pragma: no cover
|
||||
raise RuntimeError(f"Failed to list Qdrant collections: {exc}") from exc
|
||||
|
||||
for description in response.collections or []:
|
||||
try:
|
||||
client.delete_collection(description.name)
|
||||
except Exception:
|
||||
continue
|
||||
99
services/vector_db/service.py
Normal file
99
services/vector_db/service.py
Normal file
@@ -0,0 +1,99 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .config_manager import VectorDBConfigManager
|
||||
from .providers import get_provider_class, get_provider_entry
|
||||
from .providers.base import BaseVectorProvider
|
||||
|
||||
DEFAULT_VECTOR_DIMENSION = 4096
|
||||
|
||||
|
||||
class VectorDBService:
|
||||
_instance: "VectorDBService" | None = None
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(self, "_provider"):
|
||||
self._provider: Optional[BaseVectorProvider] = None
|
||||
self._provider_type: Optional[str] = None
|
||||
self._provider_config: Dict[str, Any] | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _ensure_provider(self) -> BaseVectorProvider:
|
||||
if self._provider is None:
|
||||
await self.reload()
|
||||
assert self._provider is not None # for type checker
|
||||
return self._provider
|
||||
|
||||
async def reload(self) -> BaseVectorProvider:
|
||||
async with self._lock:
|
||||
provider_type, provider_config = await VectorDBConfigManager.load_config()
|
||||
normalized_config = dict(provider_config or {})
|
||||
if (
|
||||
self._provider
|
||||
and self._provider_type == provider_type
|
||||
and self._provider_config == normalized_config
|
||||
):
|
||||
return self._provider
|
||||
|
||||
entry = get_provider_entry(provider_type)
|
||||
if not entry:
|
||||
raise RuntimeError(f"Unknown vector database provider: {provider_type}")
|
||||
if not entry.get("enabled", True):
|
||||
raise RuntimeError(f"Vector database provider '{provider_type}' is disabled")
|
||||
|
||||
provider_cls = get_provider_class(provider_type)
|
||||
if not provider_cls:
|
||||
raise RuntimeError(f"Provider class not found for '{provider_type}'")
|
||||
|
||||
provider = provider_cls(provider_config)
|
||||
await provider.initialize()
|
||||
|
||||
self._provider = provider
|
||||
self._provider_type = provider_type
|
||||
self._provider_config = normalized_config
|
||||
return provider
|
||||
|
||||
async def ensure_collection(self, collection_name: str, vector: bool = True, dim: int = DEFAULT_VECTOR_DIMENSION) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.ensure_collection(collection_name, vector, dim)
|
||||
|
||||
async def upsert_vector(self, collection_name: str, data: Dict[str, Any]) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.upsert_vector(collection_name, data)
|
||||
|
||||
async def delete_vector(self, collection_name: str, path: str) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.delete_vector(collection_name, path)
|
||||
|
||||
async def search_vectors(self, collection_name: str, query_embedding, top_k: int = 5):
|
||||
provider = await self._ensure_provider()
|
||||
return provider.search_vectors(collection_name, query_embedding, top_k)
|
||||
|
||||
async def search_by_path(self, collection_name: str, query_path: str, top_k: int = 20):
|
||||
provider = await self._ensure_provider()
|
||||
return provider.search_by_path(collection_name, query_path, top_k)
|
||||
|
||||
async def get_all_stats(self) -> Dict[str, Any]:
|
||||
provider = await self._ensure_provider()
|
||||
return provider.get_all_stats()
|
||||
|
||||
async def clear_all_data(self) -> None:
|
||||
provider = await self._ensure_provider()
|
||||
provider.clear_all_data()
|
||||
|
||||
async def current_provider(self) -> Dict[str, Any]:
|
||||
provider_type, provider_config = await VectorDBConfigManager.load_config()
|
||||
entry = get_provider_entry(provider_type) or {}
|
||||
return {
|
||||
"type": provider_type,
|
||||
"config": provider_config,
|
||||
"label": entry.get("label"),
|
||||
"enabled": entry.get("enabled", True),
|
||||
}
|
||||
Reference in New Issue
Block a user