我理解您的问题了,您需要确保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}
)

关键改进点

  1. 方法签名匹配:完全保持了与父类相同的方法签名,解决了类型检查警告

  2. 连接池管理

    • 使用类变量_connection_pool管理全局连接池
    • 第一次调用时初始化连接池
    • 后续调用复用现有连接池
  3. 连接生命周期管理

    • 从池中获取连接
    • 使用后立即归还
    • 异常情况下也能确保连接归还
  4. 兼容性

    • 保持与原版完全相同的参数和返回值
    • 无缝替换原有Milvus.from_texts()调用
  5. 资源清理

    • 对象销毁时自动关闭所有连接
    • 防止连接泄漏

这样修改后,您既不需要改变原有的调用方式,又能获得连接池带来的性能提升。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐