crewplus 0.2.47__tar.gz → 0.2.51__tar.gz

{crewplus-0.2.47 → crewplus-0.2.51}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.47
+Version: 0.2.51
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -116,6 +116,11 @@ crewplus-base/                    # GitHub repo name
 
 ```
 
+## Version Update
+
+0.2.50
+Add async aget_vector_store to enable async vector search  
+
 ## Deploy to PyPI
 
 Clean Previous Build Artifacts:
@@ -123,17 +128,17 @@ Remove the dist/, build/, and *.egg-info/ directories to ensure that no old file
 
 rm -rf dist build *.egg-info
 
-# install deployment tool
+### install deployment tool
 pip install twine
 
-# build package
+### build package
 python -m build
 
-# deploy to TestPyPI (Test first)
+### deploy to TestPyPI (Test first)
 python -m twine upload --repository testpypi dist/*
 
-# install from TestPyPI 
+### install from TestPyPI 
 pip install -i https://test.pypi.org/simple/ crewplus
 
-# Deploy to official PyPI
+### Deploy to official PyPI
 python -m twine upload dist/*

{crewplus-0.2.47 → crewplus-0.2.51}/README.md

@@ -96,6 +96,11 @@ crewplus-base/                    # GitHub repo name
 
 ```
 
+## Version Update
+
+0.2.50
+Add async aget_vector_store to enable async vector search  
+
 ## Deploy to PyPI
 
 Clean Previous Build Artifacts:
@@ -103,17 +108,17 @@ Remove the dist/, build/, and *.egg-info/ directories to ensure that no old file
 
 rm -rf dist build *.egg-info
 
-# install deployment tool
+### install deployment tool
 pip install twine
 
-# build package
+### build package
 python -m build
 
-# deploy to TestPyPI (Test first)
+### deploy to TestPyPI (Test first)
 python -m twine upload --repository testpypi dist/*
 
-# install from TestPyPI 
+### install from TestPyPI 
 pip install -i https://test.pypi.org/simple/ crewplus
 
-# Deploy to official PyPI
+### Deploy to official PyPI
 python -m twine upload dist/*

{crewplus-0.2.47 → crewplus-0.2.51}/crewplus/vectorstores/milvus/milvus_schema_manager.py

@@ -1,4 +1,4 @@
-from pymilvus import DataType, MilvusClient
+from pymilvus import DataType, MilvusClient, AsyncMilvusClient
 import json
 import logging
 from typing import Any
@@ -11,17 +11,19 @@ class MilvusSchemaManager:
     and index parameters based on a JSON definition. It interacts with a
     MilvusClient instance to perform these operations.
     """
-    def __init__(self, client: MilvusClient, logger=None):
+    def __init__(self, client: MilvusClient = None, async_client: AsyncMilvusClient = None, logger=None):
         """
         Initializes the MilvusSchemaManager.
 
         Args:
             client (MilvusClient): An instance of the Milvus client.
+            async_client (AsyncMilvusClient): An instance of the async Milvus client.
             logger (logging.Logger, optional): A logger instance. If not provided,
                                                a default logger will be created.
                                                Defaults to None.
         """
         self.client = client
+        self.async_client = async_client
         self.logger = logger or logging.getLogger(__name__)
 
     def bind_client(self, client: MilvusClient):
@@ -33,6 +35,15 @@ class MilvusSchemaManager:
         """
         self.client = client
 
+    def bind_async_client(self, async_client: AsyncMilvusClient):
+        """
+        Binds a new AsyncMilvusClient instance to the manager.
+
+        Args:
+            async_client (AsyncMilvusClient): The Milvus async client instance to use.
+        """
+        self.async_client = async_client
+
     def _add_array_field(self, schema, field_name, field_info):
         """
         Adds an ARRAY field to the schema based on field information.
@@ -82,23 +93,8 @@ class MilvusSchemaManager:
         
         schema.add_field(**field_args)
 
-    def create_collection_schema(self, json_schema: str):
-        """
-        Creates a Milvus collection schema from a JSON string.
-
-        Args:
-            json_schema (str): A JSON string defining the schema.
-
-        Returns:
-            A Milvus schema object.
-
-        Raises:
-            ValueError: If an unknown field type is encountered in the schema.
-        """
-        schema_data = json.loads(json_schema)
+    def _build_collection_schema_from_dict(self, schema, schema_data):
         fields = schema_data['node_types']['Document']['properties']
-
-        schema = self.client.create_schema(auto_id=False, enable_dynamic_fields=True)
         for field_name, field_info in fields.items():
             field_type = field_info['type']
             if field_type == "STRING" or field_type == "VARCHAR" or field_type == "TEXT":
@@ -123,28 +119,36 @@ class MilvusSchemaManager:
                 schema.add_field(field_name=field_name, datatype=DataType.FLOAT_VECTOR, dim=dim)
             else:
                 raise ValueError(f"Unknown field type: {field_type}")
-
         return schema
 
-    def create_index_params(self, json_schema: str):
+    def create_collection_schema(self, json_schema: str):
         """
-        Creates index parameters from a JSON schema string.
-
-        This method defines indexes based on the 'indexes' section of the schema
-        and automatically creates an 'AUTOINDEX' for any FLOAT_VECTOR fields.
+        Creates a Milvus collection schema from a JSON string.
 
         Args:
-            json_schema (str): A JSON string defining the schema and indexes.
+            json_schema (str): A JSON string defining the schema.
 
         Returns:
-            Milvus index parameters object.
+            A Milvus schema object.
+
+        Raises:
+            ValueError: If an unknown field type is encountered in the schema.
         """
         schema_data = json.loads(json_schema)
+        schema = self.client.create_schema(auto_id=False, enable_dynamic_fields=True)
+        return self._build_collection_schema_from_dict(schema, schema_data)
+
+    async def acreate_collection_schema(self, json_schema: str):
+        """
+        Asynchronously creates a Milvus collection schema from a JSON string.
+        """
+        schema_data = json.loads(json_schema)
+        schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_fields=True)
+        return self._build_collection_schema_from_dict(schema, schema_data)
+
+    def _build_index_params_from_dict(self, index_params, schema_data):
         fields = schema_data['node_types']['Document']['properties']
-        
-        index_params = self.client.prepare_index_params()
 
-        # Check if 'indexes' key exists
         if 'indexes' in schema_data['node_types']['Document']:
             indexes = schema_data['node_types']['Document']['indexes']
             for index_name, index_details in indexes.items():
@@ -158,7 +162,6 @@ class MilvusSchemaManager:
                     params=params
                 )
 
-        # Automatic indexing for FLOAT_VECTOR fields
         for field_name, field_info in fields.items():
             if field_info['type'] == "FLOAT_VECTOR":
                 index_params.add_index(
@@ -167,9 +170,33 @@ class MilvusSchemaManager:
                     index_type="AUTOINDEX",
                     metric_type="L2"
                 )
-
         return index_params
 
+    def create_index_params(self, json_schema: str):
+        """
+        Creates index parameters from a JSON schema string.
+
+        This method defines indexes based on the 'indexes' section of the schema
+        and automatically creates an 'AUTOINDEX' for any FLOAT_VECTOR fields.
+
+        Args:
+            json_schema (str): A JSON string defining the schema and indexes.
+
+        Returns:
+            Milvus index parameters object.
+        """
+        schema_data = json.loads(json_schema)
+        index_params = self.client.prepare_index_params()
+        return self._build_index_params_from_dict(index_params, schema_data)
+
+    async def acreate_index_params(self, json_schema: str):
+        """
+        Asynchronously creates index parameters from a JSON schema string.
+        """
+        schema_data = json.loads(json_schema)
+        index_params = AsyncMilvusClient.prepare_index_params()
+        return self._build_index_params_from_dict(index_params, schema_data)
+
     def create_collection(self, collection_name: str, json_schema: str):
         """
         Creates a new collection in Milvus.
@@ -182,6 +209,8 @@ class MilvusSchemaManager:
             json_schema (str): The JSON string defining the collection's schema
                                and indexes.
         """
+        if not self.client:
+            raise ValueError("Synchronous client not provided for create_collection.")
         schema = self.create_collection_schema(json_schema)
         index_params = self.create_index_params(json_schema)
 
@@ -192,6 +221,23 @@ class MilvusSchemaManager:
             enable_dynamic_fields=True   # we need to enable dynamic fields for schema updates
         )
 
+    async def acreate_collection(self, collection_name: str, json_schema: str):
+        """
+        Asynchronously creates a new collection in Milvus.
+        """
+        if not self.async_client:
+            raise ValueError("Asynchronous client not provided for acreate_collection.")
+            
+        schema = await self.acreate_collection_schema(json_schema)
+        index_params = await self.acreate_index_params(json_schema)
+
+        await self.async_client.create_collection(
+            collection_name=collection_name,
+            schema=schema,
+            index_params=index_params,
+            enable_dynamic_fields=True
+        )
+
     def validate_schema(self, json_schema: str) -> bool:
         """
         Validates the given schema by attempting to create a collection schema and index params.

{crewplus-0.2.47 → crewplus-0.2.51}/crewplus/vectorstores/milvus/schema_milvus.py

@@ -91,7 +91,7 @@ class SchemaMilvus(Milvus):
         )
         self.logger = logger or logging.getLogger(__name__)
         self.collection_schema = None
-        self.schema_manager = MilvusSchemaManager(client=self.client)
+        self.schema_manager = MilvusSchemaManager(client=self.client, async_client=self.aclient)
 
     def set_schema(self, schema: str):
         """
@@ -149,6 +149,29 @@ class SchemaMilvus(Milvus):
             self.logger.error(f"Failed to create collection: {e}")
             return False
 
+    async def acreate_collection(self) -> bool:
+        """
+        Asynchronously validates the schema and creates the collection using the MilvusSchemaManager.
+
+        Returns:
+            bool: True if the collection is successfully created, False otherwise.
+        """
+        if self.collection_schema is None:
+            self.logger.error("Collection schema is not set. Please set a schema using set_schema().")
+            return False
+        
+        self.schema_manager.bind_async_client(self.aclient)
+        if not self.schema_manager.validate_schema(self.collection_schema):
+            self.logger.error("Failed to validate schema")
+            return False
+        try:
+            await self.schema_manager.acreate_collection(self.collection_name, self.collection_schema)
+            self.logger.info(f"Collection {self.collection_name} created successfully")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to create collection asynchronously: {e}")
+            return False
+
     def drop_collection(self, collection_name: Optional[str] = None) -> bool:
         """
         Drops the collection using the Milvus client.

{crewplus-0.2.47 → crewplus-0.2.51}/crewplus/vectorstores/milvus/vdb_service.py

@@ -2,18 +2,20 @@
 # @Author: Cursor
 # @Date: 2025-02-12
 # @Last Modified by: Gemini
-# @Last Modified time: 2025-07-04
+# @Last Modified time: 2025-10-09
 
 import logging
 from typing import List, Dict, Union, Optional
 from langchain_milvus import Milvus
 from langchain_core.embeddings import Embeddings
 from langchain_openai import AzureOpenAIEmbeddings
-from pymilvus import MilvusClient
+from pymilvus import MilvusClient, AsyncMilvusClient
 import time
+import asyncio
 
 from ...services.init_services import get_model_balancer
 from .schema_milvus import SchemaMilvus, DEFAULT_SCHEMA
+from .milvus_schema_manager import MilvusSchemaManager
 
 class VDBService(object):
     """
@@ -84,6 +86,7 @@ class VDBService(object):
         >>> assert vector_store is same_vector_store
     """
     _client: MilvusClient
+    _async_client: AsyncMilvusClient
     _instances: Dict[str, Milvus] = {}
 
     schema: str
@@ -111,6 +114,7 @@ class VDBService(object):
             logger (logging.Logger, optional): Logger instance. Defaults to None.
         """
         self.logger = logger or logging.getLogger(__name__)
+        self.collection_schema = None
 
         if settings:
             self.settings = settings
@@ -143,17 +147,19 @@ class VDBService(object):
             raise ValueError(msg)
 
         self._client = self._initialize_milvus_client(provider)
+        self._async_client = self._initialize_async_milvus_client(provider)
 
         self.schema = schema
         self.index_params = self.settings.get("index_params")
         
+        self.schema_manager = MilvusSchemaManager(client=self._client, async_client=self._async_client)
+        
         self.logger.info("VDBService initialized successfully")
 
-    def _initialize_milvus_client(self, provider: str) -> MilvusClient:
+    def _get_milvus_client_args(self, provider: str) -> dict:
         """
-        Initializes and returns a MilvusClient with a retry mechanism.
+        Constructs the arguments for Milvus/AsyncMilvus client initialization based on the provider.
         """
-        client_args = {}
         if provider == "milvus":
             host = self.connection_args.get("host", "localhost")
             port = self.connection_args.get("port", 19530)
@@ -168,15 +174,20 @@ class VDBService(object):
                 "password": self.connection_args.get("password"),
                 "db_name": self.connection_args.get("db_name")
             }
-            # Filter out None values to use client defaults
-            client_args = {k: v for k, v in client_args.items() if v is not None}
+            return {k: v for k, v in client_args.items() if v is not None}
 
         elif provider == "zilliz":
-            client_args = self.connection_args
+            return self.connection_args
         else:
             self.logger.error(f"Unsupported vector store provider: {provider}")
             raise NotImplementedError(f"Vector store provider '{provider}' is not supported.")
 
+    def _initialize_milvus_client(self, provider: str) -> MilvusClient:
+        """
+        Initializes and returns a MilvusClient with a retry mechanism.
+        """
+        client_args = self._get_milvus_client_args(provider)
+
         try:
             # First attempt to connect
             return MilvusClient(**client_args)
@@ -189,6 +200,22 @@ class VDBService(object):
                 self.logger.error(f"Failed to initialize MilvusClient on retry. Final error: {e_retry}")
                 raise RuntimeError(f"Could not initialize MilvusClient after retry: {e_retry}")
 
+    def _initialize_async_milvus_client(self, provider: str) -> AsyncMilvusClient:
+        """
+        Initializes and returns an AsyncMilvusClient with a retry mechanism.
+        """
+        client_args = self._get_milvus_client_args(provider)
+        try:
+            return AsyncMilvusClient(**client_args)
+        except Exception as e:
+            self.logger.error(f"Failed to initialize AsyncMilvusClient, trying again. Error: {e}")
+            # Second attempt after failure
+            try:
+                return AsyncMilvusClient(**client_args)
+            except Exception as e_retry:
+                self.logger.error(f"Failed to initialize AsyncMilvusClient on retry. Final error: {e_retry}")
+                raise RuntimeError(f"Could not initialize AsyncMilvusClient after retry: {e_retry}") from e_retry
+
     def get_vector_client(self) -> MilvusClient:
         """
         Returns the active MilvusClient instance.
@@ -198,6 +225,15 @@ class VDBService(object):
         """
         return self._client
 
+    def get_async_vector_client(self) -> AsyncMilvusClient:
+        """
+        Returns the active AsyncMilvusClient instance.
+
+        Returns:
+            AsyncMilvusClient: The initialized async client for interacting with the vector database.
+        """
+        return self._async_client
+
     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.
@@ -276,6 +312,34 @@ class VDBService(object):
             self.logger.error(f"An error occurred while ensuring collection '{collection_name}' : {e}")
             raise RuntimeError(f"Failed to ensure collection '{collection_name}' .") from e
 
+    async def _aensure_collection_exists(self, collection_name: str, embeddings: Embeddings, check_existence: bool = True):
+        """
+        Asynchronously checks if a collection exists and creates it if it doesn't.
+        """
+        try:
+            client = self.get_async_vector_client()
+            if check_existence and not await client.has_collection(collection_name):
+                self.logger.info(f"Collection '{collection_name}' does not exist. Creating it.")
+                
+                schema_milvus = SchemaMilvus(
+                    embedding_function=embeddings,
+                    collection_name=collection_name,
+                    connection_args=self.connection_args,
+                    index_params=self.index_params
+                )
+
+                schema_to_use = self.schema or DEFAULT_SCHEMA
+                if not self.schema:
+                    self.logger.warning(f"No schema provided for VDBService. Using DEFAULT_SCHEMA for collection '{collection_name}'.")
+                
+                schema_milvus.set_schema(schema_to_use)
+                
+                if not await schema_milvus.acreate_collection():
+                    raise RuntimeError(f"SchemaMilvus failed to create collection '{collection_name}'.")
+        except Exception as e:
+            self.logger.error(f"An error occurred while ensuring collection '{collection_name}' : {e}")
+            raise RuntimeError(f"Failed to ensure collection '{collection_name}' .") from e
+
     def _is_good_connection(self, vdb_instance: Milvus, collection_name: str) -> tuple[bool, bool | None]:
         """
         Checks if the Milvus instance has a good connection by verifying collection existence.
@@ -301,6 +365,21 @@ class VDBService(object):
             self.logger.warning(f"Connection check failed for cached instance of '{collection_name}': {e}")
             return False, None
 
+    async def _ais_good_connection(self, vdb_instance: Milvus, collection_name: str) -> tuple[bool, bool | None]:
+        """
+        Asynchronously checks if the Milvus instance has a good connection.
+        """
+        try:
+            collection_exists = await vdb_instance.aclient.has_collection(collection_name)
+            if collection_exists:
+                self.logger.debug(f"Connection for cached instance of '{collection_name}' is alive.")
+            else:
+                self.logger.warning(f"Collection '{collection_name}' not found for cached instance. It may have been dropped.")
+            return True, collection_exists
+        except Exception as e:
+            self.logger.warning(f"Connection check failed for cached instance of '{collection_name}': {e}")
+            return False, None
+
     def get_vector_store(self, collection_name: str, embeddings: Embeddings = None, metric_type: str = "IP") -> Milvus:
         """
         Gets a vector store instance, creating it if it doesn't exist for the collection.
@@ -373,6 +452,63 @@ class VDBService(object):
 
         return vdb
 
+    async def aget_vector_store(self, collection_name: str, embeddings: Embeddings = None, metric_type: str = "IP") -> Milvus:
+        """
+        Asynchronously gets a vector store instance, creating it if it doesn't exist.
+        """
+        if not collection_name:
+            self.logger.error("aget_vector_store called with no collection_name.")
+            raise ValueError("collection_name must be provided.")
+
+        check_existence = True
+        if collection_name in self._instances:
+            instance = self._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}")
+                return instance
+            
+            self.logger.warning(f"Cached instance for '{collection_name}' is invalid. Removing it from cache.")
+            del self._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}")
+        if embeddings is None:
+            embeddings = self.get_embeddings()
+
+        await self._aensure_collection_exists(collection_name, embeddings, check_existence=check_existence)
+
+        try:
+            self.logger.debug(f"Testing embedding function for collection '{collection_name}'...")
+            await embeddings.aembed_query("validation_test_string")
+            self.logger.debug("Embedding function is valid.")
+        except Exception as e:
+            self.logger.error(
+                f"The provided embedding function is invalid and failed with error: {e}. "
+                f"Cannot create a vector store for collection '{collection_name}'."
+            )
+            raise RuntimeError(f"Invalid embedding function provided.") from e
+
+        index_params = self.index_params or {
+            "metric_type": metric_type,
+            "index_type": "AUTOINDEX",
+            "params": {}
+        }
+        
+        vdb = await asyncio.to_thread(
+            self._create_milvus_instance_with_retry,
+            collection_name=collection_name,
+            embeddings=embeddings,
+            index_params=index_params
+        )
+
+        self._instances[collection_name] = vdb
+
+        return vdb
+
     def _create_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict) -> Milvus:
         """
         Creates a Milvus instance with a retry mechanism for connection failures.
@@ -429,6 +565,36 @@ class VDBService(object):
                 del self._instances[collection_name]
                 self.logger.info(f"Removed '{collection_name}' from instance cache.")
 
+    async def adrop_collection(self, collection_name: str) -> None:
+        """
+        Asynchronously deletes a collection from the vector database and removes it from the cache.
+
+        Args:
+            collection_name (str): The name of the collection to drop.
+
+        Raises:
+            ValueError: If collection_name is not provided.
+            RuntimeError: If the operation fails on the database side.
+        """
+        if not collection_name:
+            self.logger.error("adrop_collection called without a collection_name.")
+            raise ValueError("collection_name must be provided.")
+
+        self.logger.info(f"Attempting to drop collection asynchronously: {collection_name}")
+
+        try:
+            client = self.get_async_vector_client()
+            await client.drop_collection(collection_name=collection_name)
+            self.logger.info(f"Successfully dropped collection asynchronously: {collection_name}")
+        except Exception as e:
+            self.logger.error(f"Failed to drop collection '{collection_name}' asynchronously: {e}")
+            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]
+                self.logger.info(f"Removed '{collection_name}' from instance cache.")
+
     def delete_data_by_filter(self, collection_name: str = None, filter: str = None) -> None:
         """ Delete data by filter
 
@@ -446,6 +612,23 @@ class VDBService(object):
         except Exception as e:
             raise RuntimeError(f"delete collection data failed: {str(e)}")
 
+    async def adelete_data_by_filter(self, collection_name: str = None, filter: str = None) -> None:
+        """ Asynchronously delete data by filter
+
+        Args:
+            collection_name (str): collection_name
+            filter (str): filter
+        """
+        self.logger.info(f"Delete data by filter asynchronously:{filter}")
+
+        try:
+            client=self.get_async_vector_client()
+            if collection_name is None or client is None or filter is None:
+                return RuntimeError(f"collection_name must be not null or check out your client to link milvus")
+            await client.delete(collection_name=collection_name, filter=filter)
+        except Exception as e:
+            raise RuntimeError(f"delete collection data failed: {str(e)}")
+
     @staticmethod
     def delete_old_indexes(url: str = None, vdb: Milvus = None) -> (bool | None):
         """ Delete old indexes of the same source_url

{crewplus-0.2.47 → crewplus-0.2.51}/pyproject.toml

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