crewplus 0.2.52__tar.gz → 0.2.54__tar.gz

{crewplus-0.2.52 → crewplus-0.2.54}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.52
+Version: 0.2.54
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

{crewplus-0.2.52 → crewplus-0.2.54}/crewplus/vectorstores/milvus/vdb_service.py

@@ -95,6 +95,7 @@ class VDBService(object):
     _client: MilvusClient
     _async_client: AsyncMilvusClient
     _instances: Dict[str, Milvus] = {}
+    _async_instances: Dict[str, Milvus] = {}
 
     schema: str
     embedding_function: Embeddings
@@ -153,8 +154,13 @@ class VDBService(object):
             self.logger.error(msg)
             raise ValueError(msg)
 
-        self.alias = f"crewplus-vdb-{uuid.uuid4()}"
-        self.connection_args['alias'] = self.alias
+        # Create separate aliases for sync and async clients to avoid connection handler race conditions.
+        self.sync_alias = f"crewplus-vdb-sync-{uuid.uuid4()}"
+        self.async_alias = f"crewplus-vdb-async-{uuid.uuid4()}"
+
+        # The default alias in connection_args should be the sync one, as langchain_milvus
+        # primarily uses a synchronous client and will pick up this alias.
+        self.connection_args['alias'] = self.sync_alias
 
         self._client = self._initialize_milvus_client(provider)
         self._async_client = self._initialize_async_milvus_client(provider)
@@ -183,12 +189,14 @@ class VDBService(object):
                 "user": self.connection_args.get("user"),
                 "password": self.connection_args.get("password"),
                 "db_name": self.connection_args.get("db_name"),
-                "alias": self.connection_args.get("alias")
             }
             return {k: v for k, v in client_args.items() if v is not None}
 
         elif provider == "zilliz":
-            return self.connection_args
+            # Return a copy without the default alias, as it will be added specifically for sync/async clients.
+            zilliz_args = self.connection_args.copy()
+            zilliz_args.pop('alias', None)
+            return zilliz_args
         else:
             self.logger.error(f"Unsupported vector store provider: {provider}")
             raise NotImplementedError(f"Vector store provider '{provider}' is not supported.")
@@ -198,6 +206,7 @@ class VDBService(object):
         Initializes and returns a MilvusClient with a retry mechanism.
         """
         client_args = self._get_milvus_client_args(provider)
+        client_args["alias"] = self.sync_alias
 
         try:
             # First attempt to connect
@@ -216,6 +225,8 @@ class VDBService(object):
         Initializes and returns an AsyncMilvusClient with a retry mechanism.
         """
         client_args = self._get_milvus_client_args(provider)
+        client_args["alias"] = self.async_alias
+        
         try:
             return AsyncMilvusClient(**client_args)
         except Exception as e:
@@ -245,6 +256,30 @@ class VDBService(object):
         """
         return self._async_client
 
+    def get_vector_field(self, collection_name: str) -> str:
+        """
+        Retrieves the vector field name for a given collection from a cached instance.
+
+        Args:
+            collection_name (str): The name of the collection.
+
+        Returns:
+            str: The name of the vector field.
+        
+        Raises:
+            ValueError: If no cached instance is found for the collection.
+        """
+        if collection_name in self._instances:
+            return self._instances[collection_name]._vector_field
+        if collection_name in self._async_instances:
+            return self._async_instances[collection_name]._vector_field
+        
+        self.logger.warning(f"No cached instance found for collection '{collection_name}' to get vector field. Creating a temporary sync instance.")
+        # As a fallback, create a temporary sync instance to fetch the schema info.
+        # This is less efficient but ensures the method is robust.
+        temp_instance = self.get_vector_store(collection_name)
+        return temp_instance._vector_field
+
     def get_embeddings(self, from_model_balancer: bool = False, provider: Optional[str] = "azure-openai", model_type: Optional[str] = "embedding-large") -> Embeddings:
         """
         Gets an embedding function, either from the model balancer or directly from settings.
@@ -472,21 +507,21 @@ class VDBService(object):
             raise ValueError("collection_name must be provided.")
 
         check_existence = True
-        if collection_name in self._instances:
-            instance = self._instances[collection_name]
+        if collection_name in self._async_instances:
+            instance = self._async_instances[collection_name]
             is_connected, collection_exists = await self._ais_good_connection(instance, collection_name)
             
             if is_connected and collection_exists:
-                self.logger.info(f"Returning existing vector store instance for collection: {collection_name}")
+                self.logger.info(f"Returning existing async vector store instance for collection: {collection_name}")
                 return instance
             
-            self.logger.warning(f"Cached instance for '{collection_name}' is invalid. Removing it from cache.")
-            del self._instances[collection_name]
+            self.logger.warning(f"Cached async instance for '{collection_name}' is invalid. Removing it from cache.")
+            del self._async_instances[collection_name]
             
             if is_connected and not collection_exists:
                 check_existence = False
 
-        self.logger.info(f"Creating new vector store instance for collection: {collection_name}")
+        self.logger.info(f"Creating new async vector store instance for collection: {collection_name}")
         if embeddings is None:
             embeddings = self.get_embeddings()
 
@@ -509,28 +544,34 @@ class VDBService(object):
             "params": {}
         }
         
+        # Prepare connection args with the specific async alias
+        async_connection_args = self.connection_args.copy()
+        async_connection_args["alias"] = self.async_alias
+
         vdb = await asyncio.to_thread(
             self._create_milvus_instance_with_retry,
             collection_name=collection_name,
             embeddings=embeddings,
-            index_params=index_params
+            index_params=index_params,
+            connection_args=async_connection_args
         )
 
-        self._instances[collection_name] = vdb
+        self._async_instances[collection_name] = vdb
 
         return vdb
 
-    def _create_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict) -> Milvus:
+    def _create_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict, connection_args: Optional[dict] = None) -> Milvus:
         """
         Creates a Milvus instance with a retry mechanism for connection failures.
         """
         retries = 2
+        conn_args = connection_args if connection_args is not None else self.connection_args
         for attempt in range(retries + 1):
             try:
                 vdb = Milvus(
                     embedding_function=embeddings,
                     collection_name=collection_name,
-                    connection_args=self.connection_args,
+                    connection_args=conn_args,
                     index_params=index_params
                 )
                 self.logger.info(f"Successfully connected to Milvus for collection '{collection_name}' on attempt {attempt + 1}.")
@@ -602,8 +643,8 @@ class VDBService(object):
             raise RuntimeError(f"An error occurred while dropping collection '{collection_name}' asynchronously.") from e
         finally:
             # Whether successful or not, remove the stale instance from the cache.
-            if collection_name in self._instances:
-                del self._instances[collection_name]
+            if collection_name in self._async_instances:
+                del self._async_instances[collection_name]
                 self.logger.info(f"Removed '{collection_name}' from instance cache.")
 
     def delete_data_by_filter(self, collection_name: str = None, filter: str = None) -> None:

{crewplus-0.2.52 → crewplus-0.2.54}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "crewplus"
-version = "0.2.52"
+version = "0.2.54"
 description = "Base services for CrewPlus AI applications"
 authors = [
     { name = "Tim Liu", email = "tim@opsmateai.com" },