修复代码格式和文件名大小写问题
This commit is contained in:
@@ -1,19 +1,21 @@
|
||||
from .base import VectorDBBase
|
||||
from .chromadb_impl import ChromaDBImpl
|
||||
|
||||
|
||||
def get_vector_db_service() -> VectorDBBase:
|
||||
"""
|
||||
工厂函数,初始化并返回向量数据库服务实例。
|
||||
|
||||
|
||||
目前硬编码为 ChromaDB,未来可以从配置中读取。
|
||||
"""
|
||||
# TODO: 从全局配置中读取数据库类型和路径
|
||||
db_path = "data/chroma_db"
|
||||
|
||||
|
||||
# ChromaDBImpl 是一个单例,所以这里每次调用都会返回同一个实例
|
||||
return ChromaDBImpl(path=db_path)
|
||||
|
||||
|
||||
# 全局向量数据库服务实例
|
||||
vector_db_service: VectorDBBase = get_vector_db_service()
|
||||
|
||||
__all__ = ["vector_db_service", "VectorDBBase"]
|
||||
__all__ = ["vector_db_service", "VectorDBBase"]
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class VectorDBBase(ABC):
|
||||
"""
|
||||
向量数据库的抽象基类 (ABC),定义了所有向量数据库实现必须遵循的接口。
|
||||
@@ -133,7 +134,7 @@ class VectorDBBase(ABC):
|
||||
int: 条目总数。
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@abstractmethod
|
||||
def delete_collection(self, name: str) -> None:
|
||||
"""
|
||||
@@ -142,4 +143,4 @@ class VectorDBBase(ABC):
|
||||
Args:
|
||||
name (str): 要删除的集合的名称。
|
||||
"""
|
||||
pass
|
||||
pass
|
||||
|
||||
@@ -9,11 +9,13 @@ from src.common.logger import get_logger
|
||||
|
||||
logger = get_logger("chromadb_impl")
|
||||
|
||||
|
||||
class ChromaDBImpl(VectorDBBase):
|
||||
"""
|
||||
ChromaDB 的具体实现,遵循 VectorDBBase 接口。
|
||||
采用单例模式,确保全局只有一个 ChromaDB 客户端实例。
|
||||
"""
|
||||
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
@@ -29,13 +31,12 @@ class ChromaDBImpl(VectorDBBase):
|
||||
初始化 ChromaDB 客户端。
|
||||
由于是单例,这个初始化只会执行一次。
|
||||
"""
|
||||
if not hasattr(self, '_initialized'):
|
||||
if not hasattr(self, "_initialized"):
|
||||
with self._lock:
|
||||
if not hasattr(self, '_initialized'):
|
||||
if not hasattr(self, "_initialized"):
|
||||
try:
|
||||
self.client = chromadb.PersistentClient(
|
||||
path=path,
|
||||
settings=Settings(anonymized_telemetry=False)
|
||||
path=path, settings=Settings(anonymized_telemetry=False)
|
||||
)
|
||||
self._collections: Dict[str, Any] = {}
|
||||
self._initialized = True
|
||||
@@ -48,10 +49,10 @@ class ChromaDBImpl(VectorDBBase):
|
||||
def get_or_create_collection(self, name: str, **kwargs: Any) -> Any:
|
||||
if not self.client:
|
||||
raise ConnectionError("ChromaDB 客户端未初始化")
|
||||
|
||||
|
||||
if name in self._collections:
|
||||
return self._collections[name]
|
||||
|
||||
|
||||
try:
|
||||
collection = self.client.get_or_create_collection(name=name, **kwargs)
|
||||
self._collections[name] = collection
|
||||
@@ -151,15 +152,15 @@ class ChromaDBImpl(VectorDBBase):
|
||||
except Exception as e:
|
||||
logger.error(f"获取集合 '{collection_name}' 计数失败: {e}")
|
||||
return 0
|
||||
|
||||
|
||||
def delete_collection(self, name: str) -> None:
|
||||
if not self.client:
|
||||
raise ConnectionError("ChromaDB 客户端未初始化")
|
||||
|
||||
|
||||
try:
|
||||
self.client.delete_collection(name=name)
|
||||
if name in self._collections:
|
||||
del self._collections[name]
|
||||
logger.info(f"集合 '{name}' 已被删除")
|
||||
except Exception as e:
|
||||
logger.error(f"删除集合 '{name}' 失败: {e}")
|
||||
logger.error(f"删除集合 '{name}' 失败: {e}")
|
||||
|
||||
Reference in New Issue
Block a user