Milvus向量库连接池优化方案
·
我理解您的问题了,您需要确保OptimizedMilvus.from_texts()方法与父类Milvus.from_texts()的方法签名完全匹配,同时实现连接池功能。以下是完整的优化代码:
import threading
from pymilvus import connections
from queue import Queue, Empty
from typing import List, Dict, Optional, Any
from langchain.vectorstores.milvus import Milvus
class MilvusConnectionPool:
_instance = None
_lock = threading.Lock()
def __new__(cls, host: str, port: str, pool_size: int = 5):
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super(MilvusConnectionPool, cls).__new__(cls)
cls._instance._initialize_pool(host, port, pool_size)
return cls._instance
def _initialize_pool(self, host: str, port: str, pool_size: int):
self.host = host
self.port = port
self.pool_size = pool_size
self._pool = Queue(maxsize=pool_size)
for i in range(pool_size):
alias = f"conn_{i}"
connections.connect(alias=alias, host=host, port=port)
self._pool.put(alias)
def get_connection(self) -> str:
try:
alias = self._pool.get_nowait()
try:
connections.get_connection_addr(alias)
return alias
except Exception:
connections.connect(alias=alias, host=self.host, port=self.port)
return alias
except Empty:
raise Exception("Connection pool exhausted")
def release_connection(self, alias: str):
self._pool.put(alias)
def close_all(self):
while not self._pool.empty():
try:
alias = self._pool.get_nowait()
connections.disconnect(alias)
except Empty:
break
class OptimizedMilvus(Milvus):
_connection_pool = None
@classmethod
def from_texts(
cls,
texts: List[str],
embedding: Any,
metadatas: Optional[List[dict]] = None,
collection_name: str = "LangChainCollection",
connection_args: Dict[str, Any] = None,
**kwargs: Any,
) -> "Milvus":
# 初始化连接池
if connection_args is None:
connection_args = {"host": "localhost", "port": "19530"}
if cls._connection_pool is None:
cls._connection_pool = MilvusConnectionPool(
host=connection_args["host"],
port=connection_args["port"],
pool_size=10
)
# 从连接池获取连接
alias = cls._connection_pool.get_connection()
try:
# 使用连接池中的连接
kwargs["connection_args"] = {"alias": alias}
# 调用父类方法
instance = super().from_texts(
texts=texts,
embedding=embedding,
metadatas=metadatas,
collection_name=collection_name,
**kwargs
)
# 确保返回的是OptimizedMilvus实例
instance.__class__ = cls
return instance
except Exception as e:
cls._connection_pool.release_connection(alias)
raise e
finally:
cls._connection_pool.release_connection(alias)
def __del__(self):
if hasattr(self.__class__, '_connection_pool') and self.__class__._connection_pool:
self.__class__._connection_pool.close_all()
使用方式
您可以直接将原来的Milvus.from_texts()调用替换为OptimizedMilvus.from_texts(),参数完全保持不变:
vector_db = OptimizedMilvus.from_texts(
texts=chunk_texts,
embedding=embeddings,
metadatas=metadata_list,
collection_name=self.milvus_collection_name,
connection_args={"host": self.milvus_host, "port": self.milvus_port}
)
关键改进点
-
方法签名匹配:完全保持了与父类相同的方法签名,解决了类型检查警告
-
连接池管理:
- 使用类变量
_connection_pool管理全局连接池 - 第一次调用时初始化连接池
- 后续调用复用现有连接池
- 使用类变量
-
连接生命周期管理:
- 从池中获取连接
- 使用后立即归还
- 异常情况下也能确保连接归还
-
兼容性:
- 保持与原版完全相同的参数和返回值
- 无缝替换原有
Milvus.from_texts()调用
-
资源清理:
- 对象销毁时自动关闭所有连接
- 防止连接泄漏
这样修改后,您既不需要改变原有的调用方式,又能获得连接池带来的性能提升。
更多推荐
所有评论(0)