crewplus 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

crewplus/vectorstores/milvus/schema_milvus.py

@@ -9,6 +9,30 @@ from crewplus.utils.schema_document_updater import SchemaDocumentUpdater
 from crewplus.utils.schema_action import Action
 from .milvus_schema_manager import MilvusSchemaManager
 
+DEFAULT_SCHEMA = """
+{
+    "node_types": {
+        "Document": {
+            "properties": {
+                "pk": {
+                    "type": "INT64",
+                    "is_primary": true,
+                    "auto_id": true
+                },
+                "vector": {
+                    "type": "FLOAT_VECTOR",
+                    "dim": 1536
+                },
+                "text": {
+                    "type": "VARCHAR",
+                    "max_length": 65535,
+                    "description": "The core text of the memory. This could be a user query, a documented fact, a procedural step, or a log of an event."
+                }
+	        }
+	    }
+    }
+}
+"""
 
 class SchemaMilvus(Milvus):
     """

crewplus/vectorstores/milvus/vdb_service.py

@@ -1,394 +1,427 @@
-# -*- coding: utf-8 -*-
-# @Author: Cursor
-# @Date: 2025-02-12
-# @Last Modified by: Gemini
-# @Last Modified time: 2025-07-01
-
-import logging
-from typing import List, Dict, Union, Optional
-from langchain_milvus import Zilliz
-from langchain_core.embeddings import Embeddings
-from langchain_openai import AzureOpenAIEmbeddings
-from pymilvus import MilvusClient
-
-from crewplus.services.init_services import get_model_balancer
-from crewplus.vectorstores.milvus.schema_milvus import SchemaMilvus
-
-class VDBService(object):
-    """
-    A service to manage connections to Milvus/Zilliz vector databases and embedding models.
-
-    This service centralizes the configuration and instantiation of the Milvus client
-    and provides helper methods to get embedding functions and vector store instances.
-
-    Args:
-        settings (dict, optional): A dictionary containing configuration for the vector store
-                         and embedding models.
-        endpoint (str, optional): The URI for the Zilliz cluster. Can be used for simple
-                                  initialization instead of `settings`.
-        token (str, optional): The token for authenticating with Zilliz. Must be provided
-                               with `endpoint`.
-        schema (str, optional): The schema definition for a collection. Defaults to None.
-        logger (logging.Logger, optional): An optional logger instance. Defaults to None.
-
-    Raises:
-        ValueError: If required configurations are missing.
-        NotImplementedError: If an unsupported provider is specified.
-        RuntimeError: If the MilvusClient fails to initialize after a retry.
-
-    Example:
-        >>> # Initialize with a full settings dictionary
-        >>> settings = {
-        ...     "embedder": {
-        ...         "provider": "azure-openai",
-        ...         "config": {
-        ...             "model": "text-embedding-3-small",
-        ...             "api_version": "2023-05-15",
-        ...             "api_key": "YOUR_AZURE_OPENAI_KEY",
-        ...             "openai_base_url": "YOUR_AZURE_OPENAI_ENDPOINT",
-        ...             "embedding_dims": 1536
-        ...         }
-        ...     },
-        ...     "vector_store": {
-        ...         "provider": "milvus",
-        ...         "config": {
-        ...             "host": "localhost",
-        ...             "port": 19530,
-        ...             "user": "root",
-        ...             "password": "password",
-        ...             "db_name": "default"
-        ...         }
-        ...     },
-        ...     "index_params": {
-        ...         "metric_type": "L2",
-        ...         "index_type": "AUTOINDEX",
-        ...         "params": {}
-        ...     }
-        ... }
-        >>> vdb_service = VDBService(settings=settings)
-        >>>
-        >>> # Alternatively, initialize with an endpoint and token for Zilliz
-        >>> # vdb_service_zilliz = VDBService(endpoint="YOUR_ZILLIZ_ENDPOINT", token="YOUR_ZILLIZ_TOKEN")
-        >>>
-        >>> # Get the raw Milvus client
-        >>> client = vdb_service.get_vector_client()
-        >>> print(client.list_collections())
-        >>> # Get an embedding function
-        >>> embeddings = vdb_service.get_embeddings()
-        >>> print(embeddings)
-        >>> # Get a LangChain vector store instance (will be cached)
-        >>> vector_store = vdb_service.get_vector_store(collection_name="my_collection")
-        >>> print(vector_store)
-        >>> same_vector_store = vdb_service.get_vector_store(collection_name="my_collection")
-        >>> assert vector_store is same_vector_store
-    """
-    _client: MilvusClient
-    _instances: Dict[str, Zilliz] = {}
-
-    schema: str
-    embedding_function: Embeddings
-    index_params: dict
-    connection_args: dict
-    settings: dict
-    
-    def __init__(self, settings: dict = None, endpoint: str = None, token: str = None, schema: str = None, logger: logging.Logger = None):
-        """
-        Initializes the VDBService.
-
-        Can be initialized in two ways:
-        1. By providing a full `settings` dictionary for complex configurations.
-        2. By providing `endpoint` and `token` for a direct Zilliz connection.
-           Note: When using this method, an `embedder` configuration is not created.
-           You must either use the `ModelLoadBalancer` or pass an `Embeddings` object
-           directly to methods like `get_vector_store`.
-        
-        Args:
-            settings (dict, optional): Configuration dictionary for the service. Defaults to None.
-            endpoint (str, optional): The URI for the Zilliz cluster. Used if `settings` is not provided.
-            token (str, optional): The token for authenticating with the Zilliz cluster.
-            schema (str, optional): Default schema for new collections. Defaults to None.
-            logger (logging.Logger, optional): Logger instance. Defaults to None.
-        """
-        self.logger = logger or logging.getLogger(__name__)
-
-        if settings:
-            self.settings = settings
-        elif endpoint and token:
-            self.logger.info("Initializing VDBService with endpoint and token for a Zilliz connection.")
-            self.settings = {
-                "vector_store": {
-                    "provider": "zilliz",
-                    "config": {
-                        "uri": endpoint,
-                        "token": token
-                    }
-                }
-            }
-        else:
-            raise ValueError("VDBService must be initialized with either a 'settings' dictionary or both 'endpoint' and 'token'.")
-
-        vector_store_settings = self.settings.get("vector_store")
-        if not vector_store_settings:
-            msg = "'vector_store' not found in settings"
-            self.logger.error(msg)
-            raise ValueError(msg)
-
-        provider = vector_store_settings.get("provider")
-        self.connection_args = vector_store_settings.get("config")
-
-        if not provider or not self.connection_args:
-            msg = "'provider' or 'config' not found in 'vector_store' settings"
-            self.logger.error(msg)
-            raise ValueError(msg)
-
-        self._client = self._initialize_milvus_client(provider)
-
-        self.schema = schema
-        self.index_params = self.settings.get("index_params")
-        
-        self.logger.info("VDBService initialized successfully")
-
-    def _initialize_milvus_client(self, provider: str) -> MilvusClient:
-        """
-        Initializes and returns a MilvusClient with a retry mechanism.
-        """
-        client_args = {}
-        if provider == "milvus":
-            host = self.connection_args.get("host", "localhost")
-            port = self.connection_args.get("port", 19530)
-            
-            # Use https for remote hosts, and http for local connections.
-            scheme = "https" if host not in ["localhost", "127.0.0.1"] else "http"
-            uri = f"{scheme}://{host}:{port}"
-            
-            client_args = {
-                "uri": uri,
-                "user": self.connection_args.get("user"),
-                "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}
-
-        elif provider == "zilliz":
-            client_args = self.connection_args
-        else:
-            self.logger.error(f"Unsupported vector store provider: {provider}")
-            raise NotImplementedError(f"Vector store provider '{provider}' is not supported.")
-
-        try:
-            # First attempt to connect
-            return MilvusClient(**client_args)
-        except Exception as e:
-            self.logger.error(f"Failed to initialize MilvusClient, trying again. Error: {e}")
-            # Second attempt after failure
-            try:
-                return MilvusClient(**client_args)
-            except Exception as e_retry:
-                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 get_vector_client(self) -> MilvusClient:
-        """
-        Returns the active MilvusClient instance.
-
-        Returns:
-            MilvusClient: The initialized client for interacting with the vector database.
-        """
-        return self._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.
-
-        Args:
-            from_model_balancer (bool): If True, uses the central model balancer service.
-                                        If False, creates a new instance based on 'embedder' settings.
-            model_type (str, optional): The type of model to get from the balancer. Defaults to "embedding-large".
-
-        Returns:
-            Embeddings: An instance of a LangChain embedding model.
-        """
-        if from_model_balancer:
-            model_balancer = get_model_balancer()
-            return model_balancer.get_model(provider=provider, model_type=model_type)
-
-        embedder_config = self.settings.get("embedder")
-        if not embedder_config:
-            self.logger.error("'embedder' configuration not found in settings.")
-            raise ValueError("'embedder' configuration not found in settings.")
-
-        provider = embedder_config.get("provider")
-        config = embedder_config.get("config")
-
-        if not provider or not config:
-            self.logger.error("Embedder 'provider' or 'config' not found in settings.")
-            raise ValueError("Embedder 'provider' or 'config' not found in settings.")
-
-        if provider == "azure-openai":
-            # Map the settings config to AzureOpenAIEmbeddings parameters.
-            azure_config = {
-                "azure_deployment": config.get("model"),
-                "openai_api_version": config.get("api_version"),
-                "api_key": config.get("api_key"),
-                "azure_endpoint": config.get("openai_base_url"),
-                "dimensions": config.get("embedding_dims"),
-                "chunk_size": config.get("chunk_size", 16),
-                "request_timeout": config.get("request_timeout", 60),
-                "max_retries": config.get("max_retries", 2)
-            }
-            # Filter out None values to use client defaults.
-            azure_config = {k: v for k, v in azure_config.items() if v is not None}
-            
-            return AzureOpenAIEmbeddings(**azure_config)
-        else:
-            self.logger.error(f"Unsupported embedding provider: {provider}")
-            raise NotImplementedError(f"Embedding provider '{provider}' is not supported yet.")
-        
-    def get_vector_store(self, collection_name: str, embeddings: Embeddings = None, metric_type: str = "L2") -> Zilliz:
-        """
-        Gets a vector store instance, creating it if it doesn't exist for the collection.
-        This method validates both the embedding function and the vector store connection
-        before caching the instance to prevent faulty instances from being reused.
-
-        Args:
-            collection_name (str): The name of the collection in the vector database.
-            embeddings (Embeddings, optional): An embedding model instance. If None, one is created.
-            metric_type (str): The distance metric for the index. Defaults to "L2".
-
-        Returns:
-            Zilliz: LangChain Zilliz instance, which is compatible with both Zilliz and Milvus.
-        """
-        if not collection_name:
-            self.logger.error("get_vector_store called with no collection_name.")
-            raise ValueError("collection_name must be provided.")
-
-        # Return the cached instance if it already exists.
-        if collection_name in self._instances:
-            self.logger.info(f"Returning existing vector store instance for collection: {collection_name}")
-            return self._instances[collection_name]
-
-        self.logger.info(f"Creating new vector store instance for collection: {collection_name}")
-        if embeddings is None:
-            embeddings = self.get_embeddings()
-
-        # 1. Validate the embedding function before proceeding.
-        try:
-            self.logger.debug(f"Testing embedding function for collection '{collection_name}'...")
-            embeddings.embed_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
-
-        # If embeddings are valid, proceed to create the Zilliz instance.
-        index_params = self.index_params or {
-            "metric_type": metric_type,
-            "index_type": "AUTOINDEX",
-            "params": {}
-        }
-        
-        vdb = Zilliz(
-            embedding_function=embeddings,
-            collection_name=collection_name,
-            connection_args=self.connection_args,
-            index_params=index_params
-        )
-
-        # Cache the newly created instance.
-        self._instances[collection_name] = vdb
-
-        return vdb
-
-    def delete_old_indexes(self, url: str = None, vdb: Zilliz = None) -> (bool | None):
-        """ Delete old indexes of the same source_url
-
-        Args:
-            url (str): source url
-            vdb (Zilliz): Zilliz instance
-        """
-        self.logger.info(f"Delete old indexes of the same source_url:{url}")
-
-        if url is None or vdb is None:
-            return None
-
-        # Delete indexes of the same source_url
-        expr = f'source_url == "{url}" or source == "{url}"'
-        pks = vdb.get_pks(expr)
-
-        # Delete entities by pks
-        if pks is not None and len(pks) > 0 :
-            res = vdb.delete(pks)
-            self.logger.info("Deleted old indexes result: " + str(res))
-            return res
-
-    def delete_old_indexes_by_id(self, source_id: str = None, vdb: Zilliz = None) -> (bool | None):
-        """ Delete old indexes of the same source_id
-
-        Args:
-            source_id (str): source id
-        """
-        self.logger.info(f"Delete old indexes of the same source_id:{source_id}")
-
-        if source_id is None or vdb is None:
-            return None
-
-        # Delete indexes of the same source_id
-        expr = f'source_id == "{source_id}"'
-        pks = vdb.get_pks(expr)
-
-        # Delete entities by pks
-        if pks is not None and len(pks) > 0 :
-            res = vdb.delete(pks)
-            self.logger.info("Deleted old indexes result: " + str(res))
-            return res
-
-    def drop_collection(self, collection_name: str) -> None:
-        """
-        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("drop_collection called without a collection_name.")
-            raise ValueError("collection_name must be provided.")
-
-        self.logger.info(f"Attempting to drop collection: {collection_name}")
-
-        try:
-            client = self.get_vector_client()
-            client.drop_collection(collection_name=collection_name)
-            self.logger.info(f"Successfully dropped collection: {collection_name}")
-        except Exception as e:
-            self.logger.error(f"Failed to drop collection '{collection_name}': {e}")
-            raise RuntimeError(f"An error occurred while dropping collection '{collection_name}'.") 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
-
-        Args:
-            collection_name (str): collection_name
-            filter (str): filter
-        """
-        self.logger.info(f"Delete data by filter:{filter}")
-
-        try:
-            client=self.get_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")
-            client.delete(collection_name=collection_name, filter=filter)
-        except Exception as e:
+# -*- coding: utf-8 -*-
+# @Author: Cursor
+# @Date: 2025-02-12
+# @Last Modified by: Gemini
+# @Last Modified time: 2025-07-04
+
+import logging
+from typing import List, Dict, Union, Optional
+from langchain_milvus import Zilliz
+from langchain_core.embeddings import Embeddings
+from langchain_openai import AzureOpenAIEmbeddings
+from pymilvus import MilvusClient
+
+from ...services.init_services import get_model_balancer
+from .schema_milvus import SchemaMilvus, DEFAULT_SCHEMA
+
+class VDBService(object):
+    """
+    A service to manage connections to Milvus/Zilliz vector databases and embedding models.
+
+    This service centralizes the configuration and instantiation of the Milvus client
+    and provides helper methods to get embedding functions and vector store instances.
+
+    Args:
+        settings (dict, optional): A dictionary containing configuration for the vector store
+                         and embedding models.
+        endpoint (str, optional): The URI for the Zilliz cluster. Can be used for simple
+                                  initialization instead of `settings`.
+        token (str, optional): The token for authenticating with Zilliz. Must be provided
+                               with `endpoint`.
+        schema (str, optional): The schema definition for a collection. Defaults to None.
+        logger (logging.Logger, optional): An optional logger instance. Defaults to None.
+
+    Raises:
+        ValueError: If required configurations are missing.
+        NotImplementedError: If an unsupported provider is specified.
+        RuntimeError: If the MilvusClient fails to initialize after a retry.
+
+    Example:
+        >>> # Initialize with a full settings dictionary
+        >>> settings = {
+        ...     "embedder": {
+        ...         "provider": "azure-openai",
+        ...         "config": {
+        ...             "model": "text-embedding-3-small",
+        ...             "api_version": "2023-05-15",
+        ...             "api_key": "YOUR_AZURE_OPENAI_KEY",
+        ...             "openai_base_url": "YOUR_AZURE_OPENAI_ENDPOINT",
+        ...             "embedding_dims": 1536
+        ...         }
+        ...     },
+        ...     "vector_store": {
+        ...         "provider": "milvus",
+        ...         "config": {
+        ...             "host": "localhost",
+        ...             "port": 19530,
+        ...             "user": "root",
+        ...             "password": "password",
+        ...             "db_name": "default"
+        ...         }
+        ...     },
+        ...     "index_params": {
+        ...         "metric_type": "L2",
+        ...         "index_type": "AUTOINDEX",
+        ...         "params": {}
+        ...     }
+        ... }
+        >>> vdb_service = VDBService(settings=settings)
+        >>>
+        >>> # Alternatively, initialize with an endpoint and token for Zilliz
+        >>> # vdb_service_zilliz = VDBService(endpoint="YOUR_ZILLIZ_ENDPOINT", token="YOUR_ZILLIZ_TOKEN")
+        >>>
+        >>> # Get the raw Milvus client
+        >>> client = vdb_service.get_vector_client()
+        >>> print(client.list_collections())
+        >>> # Get an embedding function
+        >>> embeddings = vdb_service.get_embeddings()
+        >>> print(embeddings)
+        >>> # Get a LangChain vector store instance (will be cached)
+        >>> vector_store = vdb_service.get_vector_store(collection_name="my_collection")
+        >>> print(vector_store)
+        >>> same_vector_store = vdb_service.get_vector_store(collection_name="my_collection")
+        >>> assert vector_store is same_vector_store
+    """
+    _client: MilvusClient
+    _instances: Dict[str, Zilliz] = {}
+
+    schema: str
+    embedding_function: Embeddings
+    index_params: dict
+    connection_args: dict
+    settings: dict
+    
+    def __init__(self, settings: dict = None, endpoint: str = None, token: str = None, schema: str = None, logger: logging.Logger = None):
+        """
+        Initializes the VDBService.
+
+        Can be initialized in two ways:
+        1. By providing a full `settings` dictionary for complex configurations.
+        2. By providing `endpoint` and `token` for a direct Zilliz connection.
+           Note: When using this method, an `embedder` configuration is not created.
+           You must either use the `ModelLoadBalancer` or pass an `Embeddings` object
+           directly to methods like `get_vector_store`.
+        
+        Args:
+            settings (dict, optional): Configuration dictionary for the service. Defaults to None.
+            endpoint (str, optional): The URI for the Zilliz cluster. Used if `settings` is not provided.
+            token (str, optional): The token for authenticating with the Zilliz cluster.
+            schema (str, optional): Default schema for new collections. Defaults to None.
+            logger (logging.Logger, optional): Logger instance. Defaults to None.
+        """
+        self.logger = logger or logging.getLogger(__name__)
+
+        if settings:
+            self.settings = settings
+        elif endpoint and token:
+            self.logger.info("Initializing VDBService with endpoint and token for a Zilliz connection.")
+            self.settings = {
+                "vector_store": {
+                    "provider": "zilliz",
+                    "config": {
+                        "uri": endpoint,
+                        "token": token
+                    }
+                }
+            }
+        else:
+            raise ValueError("VDBService must be initialized with either a 'settings' dictionary or both 'endpoint' and 'token'.")
+
+        vector_store_settings = self.settings.get("vector_store")
+        if not vector_store_settings:
+            msg = "'vector_store' not found in settings"
+            self.logger.error(msg)
+            raise ValueError(msg)
+
+        provider = vector_store_settings.get("provider")
+        self.connection_args = vector_store_settings.get("config")
+
+        if not provider or not self.connection_args:
+            msg = "'provider' or 'config' not found in 'vector_store' settings"
+            self.logger.error(msg)
+            raise ValueError(msg)
+
+        self._client = self._initialize_milvus_client(provider)
+
+        self.schema = schema
+        self.index_params = self.settings.get("index_params")
+        
+        self.logger.info("VDBService initialized successfully")
+
+    def _initialize_milvus_client(self, provider: str) -> MilvusClient:
+        """
+        Initializes and returns a MilvusClient with a retry mechanism.
+        """
+        client_args = {}
+        if provider == "milvus":
+            host = self.connection_args.get("host", "localhost")
+            port = self.connection_args.get("port", 19530)
+            
+            # Use https for remote hosts, and http for local connections.
+            scheme = "https" if host not in ["localhost", "127.0.0.1"] else "http"
+            uri = f"{scheme}://{host}:{port}"
+            
+            client_args = {
+                "uri": uri,
+                "user": self.connection_args.get("user"),
+                "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}
+
+        elif provider == "zilliz":
+            client_args = self.connection_args
+        else:
+            self.logger.error(f"Unsupported vector store provider: {provider}")
+            raise NotImplementedError(f"Vector store provider '{provider}' is not supported.")
+
+        try:
+            # First attempt to connect
+            return MilvusClient(**client_args)
+        except Exception as e:
+            self.logger.error(f"Failed to initialize MilvusClient, trying again. Error: {e}")
+            # Second attempt after failure
+            try:
+                return MilvusClient(**client_args)
+            except Exception as e_retry:
+                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 get_vector_client(self) -> MilvusClient:
+        """
+        Returns the active MilvusClient instance.
+
+        Returns:
+            MilvusClient: The initialized client for interacting with the vector database.
+        """
+        return self._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.
+
+        Args:
+            from_model_balancer (bool): If True, uses the central model balancer service.
+                                        If False, creates a new instance based on 'embedder' settings.
+            model_type (str, optional): The type of model to get from the balancer. Defaults to "embedding-large".
+
+        Returns:
+            Embeddings: An instance of a LangChain embedding model.
+        """
+        if from_model_balancer:
+            model_balancer = get_model_balancer()
+            return model_balancer.get_model(provider=provider, model_type=model_type)
+
+        embedder_config = self.settings.get("embedder")
+        if not embedder_config:
+            self.logger.error("'embedder' configuration not found in settings.")
+            raise ValueError("'embedder' configuration not found in settings.")
+
+        provider = embedder_config.get("provider")
+        config = embedder_config.get("config")
+
+        if not provider or not config:
+            self.logger.error("Embedder 'provider' or 'config' not found in settings.")
+            raise ValueError("Embedder 'provider' or 'config' not found in settings.")
+
+        if provider == "azure-openai":
+            # Map the settings config to AzureOpenAIEmbeddings parameters.
+            azure_config = {
+                "azure_deployment": config.get("model"),
+                "openai_api_version": config.get("api_version"),
+                "api_key": config.get("api_key"),
+                "azure_endpoint": config.get("openai_base_url"),
+                "dimensions": config.get("embedding_dims"),
+                "chunk_size": config.get("chunk_size", 16),
+                "request_timeout": config.get("request_timeout", 60),
+                "max_retries": config.get("max_retries", 2)
+            }
+            # Filter out None values to use client defaults.
+            azure_config = {k: v for k, v in azure_config.items() if v is not None}
+            
+            return AzureOpenAIEmbeddings(**azure_config)
+        else:
+            self.logger.error(f"Unsupported embedding provider: {provider}")
+            raise NotImplementedError(f"Embedding provider '{provider}' is not supported yet.")
+        
+    def _ensure_collection_exists(self, collection_name: str, embeddings: Embeddings):
+        """
+        Checks if a collection exists and creates it if it doesn't.
+        This operation is wrapped in a try-except block to handle potential failures
+        during collection creation.
+        """
+        try:
+            client = self.get_vector_client()
+            if not 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 schema_milvus.create_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}' exists: {e}")
+            raise RuntimeError(f"Failed to ensure collection '{collection_name}' exists.") from e
+
+    def get_vector_store(self, collection_name: str, embeddings: Embeddings = None, metric_type: str = "L2") -> Zilliz:
+        """
+        Gets a vector store instance, creating it if it doesn't exist for the collection.
+        This method validates both the embedding function and the vector store connection
+        before caching the instance to prevent faulty instances from being reused.
+
+        Args:
+            collection_name (str): The name of the collection in the vector database.
+            embeddings (Embeddings, optional): An embedding model instance. If None, one is created.
+            metric_type (str): The distance metric for the index. Defaults to "L2".
+
+        Returns:
+            Zilliz: LangChain Zilliz instance, which is compatible with both Zilliz and Milvus.
+        """
+        if not collection_name:
+            self.logger.error("get_vector_store called with no collection_name.")
+            raise ValueError("collection_name must be provided.")
+
+        # Return the cached instance if it already exists.
+        if collection_name in self._instances:
+            self.logger.info(f"Returning existing vector store instance for collection: {collection_name}")
+            return self._instances[collection_name]
+
+        self.logger.info(f"Creating new vector store instance for collection: {collection_name}")
+        if embeddings is None:
+            embeddings = self.get_embeddings()
+
+        # Ensure the collection exists before proceeding.
+        self._ensure_collection_exists(collection_name, embeddings)
+
+        # 1. Validate the embedding function before proceeding.
+        try:
+            self.logger.debug(f"Testing embedding function for collection '{collection_name}'...")
+            embeddings.embed_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
+
+        # If embeddings are valid, proceed to create the Zilliz instance.
+        index_params = self.index_params or {
+            "metric_type": metric_type,
+            "index_type": "AUTOINDEX",
+            "params": {}
+        }
+        
+        vdb = Zilliz(
+            embedding_function=embeddings,
+            collection_name=collection_name,
+            connection_args=self.connection_args,
+            index_params=index_params
+        )
+
+        # Cache the newly created instance.
+        self._instances[collection_name] = vdb
+
+        return vdb
+
+    def delete_old_indexes(self, url: str = None, vdb: Zilliz = None) -> (bool | None):
+        """ Delete old indexes of the same source_url
+
+        Args:
+            url (str): source url
+            vdb (Zilliz): Zilliz instance
+        """
+        self.logger.info(f"Delete old indexes of the same source_url:{url}")
+
+        if url is None or vdb is None:
+            return None
+
+        # Delete indexes of the same source_url
+        expr = f'source_url == "{url}" or source == "{url}"'
+        pks = vdb.get_pks(expr)
+
+        # Delete entities by pks
+        if pks is not None and len(pks) > 0 :
+            res = vdb.delete(pks)
+            self.logger.info("Deleted old indexes result: " + str(res))
+            return res
+
+    def delete_old_indexes_by_id(self, source_id: str = None, vdb: Zilliz = None) -> (bool | None):
+        """ Delete old indexes of the same source_id
+
+        Args:
+            source_id (str): source id
+        """
+        self.logger.info(f"Delete old indexes of the same source_id:{source_id}")
+
+        if source_id is None or vdb is None:
+            return None
+
+        # Delete indexes of the same source_id
+        expr = f'source_id == "{source_id}"'
+        pks = vdb.get_pks(expr)
+
+        # Delete entities by pks
+        if pks is not None and len(pks) > 0 :
+            res = vdb.delete(pks)
+            self.logger.info("Deleted old indexes result: " + str(res))
+            return res
+
+    def drop_collection(self, collection_name: str) -> None:
+        """
+        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("drop_collection called without a collection_name.")
+            raise ValueError("collection_name must be provided.")
+
+        self.logger.info(f"Attempting to drop collection: {collection_name}")
+
+        try:
+            client = self.get_vector_client()
+            client.drop_collection(collection_name=collection_name)
+            self.logger.info(f"Successfully dropped collection: {collection_name}")
+        except Exception as e:
+            self.logger.error(f"Failed to drop collection '{collection_name}': {e}")
+            raise RuntimeError(f"An error occurred while dropping collection '{collection_name}'.") 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
+
+        Args:
+            collection_name (str): collection_name
+            filter (str): filter
+        """
+        self.logger.info(f"Delete data by filter:{filter}")
+
+        try:
+            client=self.get_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")
+            client.delete(collection_name=collection_name, filter=filter)
+        except Exception as e:
             raise RuntimeError(f"delete collection data failed: {str(e)}")

{crewplus-0.2.7.dist-info → crewplus-0.2.9.dist-info}/METADATA

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

{crewplus-0.2.7.dist-info → crewplus-0.2.9.dist-info}/RECORD

@@ -1,7 +1,7 @@
-crewplus-0.2.7.dist-info/METADATA,sha256=MAKdIos08Osw84jGk2hENFuJQPHAPG-H0GmFeqxPERk,4909
-crewplus-0.2.7.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-crewplus-0.2.7.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.7.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.9.dist-info/METADATA,sha256=3a-vd0vbsjeIv53XWjBShM14_UeusmAP4f_rtOb8iIo,4909
+crewplus-0.2.9.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.9.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.9.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 crewplus/services/__init__.py,sha256=MmH2v3N0ZMsuqFNAupkXENjUqvgf5ehQ99H6EzPqLZU,48
 crewplus/services/gemini_chat_model.py,sha256=i9p5KvSJYaHSUBLPKM_bpyGVLWCDQoNeah_WjQVJRXs,26227
@@ -11,10 +11,10 @@ crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSn
 crewplus/utils/schema_document_updater.py,sha256=frvffxn2vbi71fHFPoGb9hq7gH2azmmdq17p-Fumnvg,7322
 crewplus/vectorstores/milvus/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 crewplus/vectorstores/milvus/milvus_schema_manager.py,sha256=qHMVIM0NS3rLfACb8d3-tQS9hJo6_7_YP8AxVx4t1Cc,9019
-crewplus/vectorstores/milvus/schema_milvus.py,sha256=GhHTtCH5HsIJc3RHa25RXl3aZdkS3Rba5KeuUk_Hi0k,11425
-crewplus/vectorstores/milvus/vdb_service.py,sha256=iCiSXcmtpD9tW1WYyJ1P0SRJEFEsUoH3lY65mnwf6R8,17524
+crewplus/vectorstores/milvus/schema_milvus.py,sha256=ZRe6g4U6dMaUqZOEqpYukKIBZoGLMMVVdVUvNPiT0YA,12129
+crewplus/vectorstores/milvus/vdb_service.py,sha256=I6rlEA9hYB3POWSlaVAzxl9Hdq4zLljTnmpdvIWLYZo,18777
 docs/GeminiChatModel.md,sha256=_IQyup3ofAa2HxfSurO1GYUEezTHYYt5Q1khYNVThGM,8040
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.7.dist-info/RECORD,,
+crewplus-0.2.9.dist-info/RECORD,,

{crewplus-0.2.7.dist-info → crewplus-0.2.9.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: pdm-backend (2.4.4)
+Generator: pdm-backend (2.4.5)
 Root-Is-Purelib: true
 Tag: py3-none-any