crewplus 0.1.6__py3-none-any.whl → 0.2.1__py3-none-any.whl

crewplus/services/init_services.py

@@ -0,0 +1,16 @@
+import os
+from crewplus.services.model_load_balancer import ModelLoadBalancer
+
+model_balancer = None
+
+def init_load_balancer():
+    global model_balancer
+    if model_balancer is None:
+        config_path = os.getenv("MODEL_CONFIG_PATH", "config/models_config.json")
+        model_balancer = ModelLoadBalancer(config_path)
+        model_balancer.load_config()  # Load initial configuration synchronously
+
+def get_model_balancer() -> ModelLoadBalancer:
+    if model_balancer is None:
+        raise RuntimeError("ModelLoadBalancer not initialized")
+    return model_balancer

crewplus/services/model_load_balancer.py

@@ -81,7 +81,7 @@ class ModelLoadBalancer:
 
         Args:
             provider: The model provider (e.g., 'azure-openai', 'google-genai').
-            model_type: The type of model (e.g., 'inference', 'embedding').
+            model_type: The type of model (e.g., 'inference', 'embedding', 'embedding-large').
             deployment_name: The unique name for the model deployment.
 
         Returns:

crewplus/utils/schema_action.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+class Action(Enum):
+    UPSERT = "upsert"  # Update existing fields; if a match is found, it updates, otherwise, it inserts. Does not delete unmatched existing fields.
+    DELETE = "delete"  # Clear data from fields in the schema.
+    UPDATE = "update"  # Update only the matching original fields.
+    INSERT = "insert"  # Insert data, clearing the original fields before inserting new values.

crewplus/utils/schema_document_updater.py

@@ -0,0 +1,173 @@
+from typing import List
+from langchain_core.documents import Document
+import random
+
+class SchemaDocumentUpdater:
+    """A utility class for updating and creating LangChain Documents with specific metadata schemas."""
+    
+    @staticmethod
+    def update_document_metadata(document: Document, metadata: dict) -> Document:
+        """
+        Updates the metadata of a LangChain Document.
+
+        Args:
+            document (Document): The document to update.
+            metadata (dict): A dictionary containing the metadata to add or update.
+
+        Returns:
+            Document: The updated document with the new metadata.
+        """
+        metadata_updates = document.metadata
+
+        for key, value in metadata.items():
+            metadata_updates[key] = value
+
+        return Document(
+            page_content=document.page_content,
+            metadata=metadata_updates
+        )
+
+    @staticmethod
+    def delete_document_metadata(document: Document, keys_to_delete: List[str]) -> Document:
+        """
+        Deletes specified keys from the metadata of a LangChain Document.
+
+        Args:
+            document (Document): The document to update.
+            keys_to_delete (List[str]): A list of keys to delete from the metadata.
+
+        Returns:
+            Document: The updated document with the specified metadata keys removed.
+        """
+        metadata = document.metadata
+
+        for key in keys_to_delete:
+            if key in metadata:
+                del metadata[key]
+
+        return Document(
+            page_content=document.page_content,
+            metadata=metadata
+        )
+
+    @staticmethod
+    def add_sample_metadata(document: Document, type: str) -> Document:
+        """
+        Adds sample metadata to a document based on a specified type.
+
+        The metadata schema is tailored for either "Reg Wheel" or "Robot" types.
+
+        Args:
+            document (Document): The document to which sample metadata will be added.
+            type (str): The type of sample metadata to add ("Reg Wheel" or "Robot").
+
+        Returns:
+            Document: The document with added sample metadata.
+        """
+        if type == "Reg Wheel":
+            meta = {
+                "keywords": "Reg Wheel",
+                "plant_metadata": {
+                    "entity_id": "EQUIP_123",
+                    "entity_type": "Machine",
+                    "hierarchy_path": "/EnterpriseA/SITE_A/LINE_003/",
+                    "entity_tags": ["nickname_for_EQUIP_123", "PB3"],
+                    "parent_entity": None,
+                    "linked_entities": []
+                },
+                "version_metadata": {
+                    "version_id": "V2.0",
+                    "version_tags": ["global"],
+                    "version_date": "2024/05/23"
+                },
+                "other_metadata": {}
+            }
+        else:  # Robot
+            meta = {
+                "keywords": "Robot",
+                "plant_metadata": {
+                    "entity_id": "EQUIP_124",
+                    "entity_type": "Robot",
+                    "hierarchy_path": "/EnterpriseA/SITE_A/LINE_002/",
+                    "entity_tags": ["nickname_for_EQUIP_124", "RB2"],
+                    "parent_entity": None,
+                    "linked_entities": []
+                },
+                "version_metadata": {
+                    "version_id": "R1.0",
+                    "version_tags": ["prototype"],
+                    "version_date": "2024/05/23"
+                },
+                "other_metadata": {}
+            }
+
+        updated_document = SchemaDocumentUpdater.update_document_metadata(document, meta)
+        return updated_document
+
+    @staticmethod
+    def create_test_document(index: int, type: str) -> Document:
+        """
+        Creates a test document with sample content and metadata.
+
+        The content and metadata are generated based on the specified type ("Reg Wheel" or "Robot").
+
+        Args:
+            index (int): An index number to make the document unique.
+            type (str): The type of test document to create ("Reg Wheel" or "Robot").
+
+        Returns:
+            A new test document.
+        """
+        meta = {
+            "title": f"{type} Maintenance Record {index}",
+            "source_url": f"http://example.com/{type.lower()}_maintenance_{index}",
+            "file_type": "xlsx",
+            "page": index
+        }
+
+        if type == "Reg Wheel":
+            page_content = ["| Date       | Maintenance Performed | Technician | Notes                      |",
+                            "|------------|-----------------------|------------|----------------------------|"]
+            for _ in range(random.randint(10, 20)):
+                day = random.randint(1, 28)
+                maintenance_performed = random.choice(["Oil Change", "Belt Replacement", "Alignment Check", "General Inspection"])
+                technician = random.choice(["John Doe", "Jane Smith", "Jim Brown"])
+                notes = random.choice(["Changed oil and filter", "Replaced worn-out belt", "Checked and adjusted align", "No issues found"])
+                page_content.append(f"| 2023-05-{day:02} | {maintenance_performed} | {technician} | {notes} |")
+            page_content = "\n".join(page_content)
+        else:  # Robot
+            technicians = ["Bob", "Tim", "Alice"]
+            page_content = ["| Date       | Maintenance Performed | Technician | Notes                               |",
+                            "|------------|-----------------------|------------|-------------------------------------|"]
+            for _ in range(random.randint(10, 20)):
+                day = random.randint(1, 28)
+                maintenance_performed = random.choice(["Sensor Calibration", "Actuator Testing", "Software Update", "Battery Replacement"])
+                technician = random.choice(technicians)
+                notes = random.choice(["Calibrated all sensors", "Tested and replaced faulty actuators", "Updated robot software to v2.1", "Replaced old battery with new one"])
+                page_content.append(f"| 2023-05-{day:02} | {maintenance_performed} | {technician} | {notes} |")
+            page_content = "\n".join(page_content)
+
+        document = Document(page_content=page_content, metadata=meta)
+        return SchemaDocumentUpdater.add_sample_metadata(document, type)
+
+    @staticmethod
+    def create_test_documents(doc_num: int) -> List[Document]:
+        """
+        Creates a list of test documents.
+
+        It generates a mix of "Reg Wheel" and "Robot" documents.
+
+        Args:
+            doc_num (int): The total number of documents to create.
+
+        Returns:
+            List[Document]: A list of created test documents.
+        """
+
+        reg_wheel_docs_num = doc_num * 2 // 3
+        robot_docs_num = doc_num - reg_wheel_docs_num
+        
+        documents = [SchemaDocumentUpdater.create_test_document(i+1, "Reg Wheel") for i in range(reg_wheel_docs_num)]
+        documents += [SchemaDocumentUpdater.create_test_document(i+1 + reg_wheel_docs_num, "Robot") for i in range(robot_docs_num)]
+        
+        return documents

crewplus/vectorstores/milvus/milvus_schema_manager.py

@@ -0,0 +1,221 @@
+from pymilvus import DataType, MilvusClient
+import json
+import logging
+from typing import Any
+
+class MilvusSchemaManager:
+    """
+    Manages Milvus/Milvus collection schemas.
+
+    This class provides functionalities to create and validate collection schemas
+    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):
+        """
+        Initializes the MilvusSchemaManager.
+
+        Args:
+            client (MilvusClient): An instance of the 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.logger = logger or logging.getLogger(__name__)
+
+    def bind_client(self, client: MilvusClient):
+        """
+        Binds a new MilvusClient instance to the manager.
+
+        Args:
+            client (MilvusClient): The Milvus client instance to use.
+        """
+        self.client = client
+
+    def _add_array_field(self, schema, field_name, field_info):
+        """
+        Adds an ARRAY field to the schema based on field information.
+
+        This is a helper method to handle the specific logic for creating ARRAY fields.
+
+        Args:
+            schema: The Milvus schema object to add the field to.
+            field_name (str): The name of the field.
+            field_info (dict): A dictionary containing information about the field,
+                               such as element type and max capacity.
+
+        Raises:
+            ValueError: If required information like 'element' or 'max_capacity'
+                        is missing from field_info, or if an unsupported element
+                        type is specified.
+        """
+        element_type_str = field_info.get("element")
+        if not element_type_str:
+            raise ValueError(f"Array field '{field_name}' must have 'element' type specified.")
+
+        element_type = None
+        if element_type_str in ["STRING", "VARCHAR", "TEXT"]:
+            element_type = DataType.VARCHAR
+        elif element_type_str == "INT64":
+            element_type = DataType.INT64
+        else:
+            raise ValueError(f"Unsupported element type '{element_type_str}' for ARRAY field '{field_name}'.")
+
+        max_capacity = field_info.get("max_capacity")
+        if max_capacity is None:
+            raise ValueError(f"Array field '{field_name}' must have 'max_capacity' specified.")
+
+        nullable = field_info.get('nullable', True)
+
+        field_args = {
+            "field_name": field_name,
+            "datatype": DataType.ARRAY,
+            "element_type": element_type,
+            "max_capacity": int(max_capacity),
+            "nullable": nullable,
+        }
+
+        if element_type == DataType.VARCHAR:
+            max_length = field_info.get('max_length', 65535)
+            field_args["max_length"] = int(max_length)
+        
+        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)
+        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":
+                max_length = field_info.get('max_length', 256)  # Default max_length if not provided
+                nullable = field_info.get('nullable', False)    # Default nullable if not provided
+                schema.add_field(field_name=field_name, datatype=DataType.VARCHAR, max_length=max_length, nullable=nullable)
+            elif field_type == "JSON":
+                nullable = field_info.get('nullable', True)
+                schema.add_field(field_name=field_name, datatype=DataType.JSON, nullable=nullable)
+            elif field_type == "INT64":
+                is_primary = field_info.get('is_primary', False)
+                auto_id = field_info.get('auto_id', False)
+                nullable = field_info.get('nullable', False)
+                schema.add_field(field_name=field_name, datatype=DataType.INT64, is_primary=is_primary, auto_id=auto_id, nullable=nullable)
+            elif field_type == "ARRAY":
+                self._add_array_field(schema, field_name, field_info)
+            elif field_type == "FLOAT_VECTOR":
+                dim = field_info.get('dim', 1536)  # Default dimension if not provided
+                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):
+        """
+        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)
+        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():
+                field_name = index_details['fieldname']
+                index_type = index_details['type']
+                params = index_details['params']
+                index_params.add_index(
+                    field_name=field_name,
+                    index_type=index_type,
+                    index_name=index_name,
+                    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(
+                    field_name=field_name,
+                    index_name="vector",
+                    index_type="AUTOINDEX",
+                    metric_type="L2"
+                )
+
+        return index_params
+
+    def create_collection(self, collection_name: str, json_schema: str):
+        """
+        Creates a new collection in Milvus.
+
+        This method orchestrates the creation of the schema and index parameters
+        before creating the collection itself.
+
+        Args:
+            collection_name (str): The name for the new collection.
+            json_schema (str): The JSON string defining the collection's schema
+                               and indexes.
+        """
+        schema = self.create_collection_schema(json_schema)
+        index_params = self.create_index_params(json_schema)
+
+        self.client.create_collection(
+            collection_name=collection_name,
+            schema=schema,
+            index_params=index_params,
+            enable_dynamic_fields=True   # we need to enable dynamic fields for schema updates
+        )
+
+    def validate_schema(self, json_schema: str) -> bool:
+        """
+        Validates the given schema by attempting to create a collection schema and index params.
+
+        Args:
+            json_schema (str): The schema JSON string to validate.
+
+        Returns:
+            bool: True if the schema is valid, False if any exceptions are caught.
+        """
+        try:
+            self.create_collection_schema(json_schema)
+            self.create_index_params(json_schema)
+            return True
+        except Exception as e:
+            self.logger.error(f"Schema validation failed: {e}")
+            return False
+
+
+class ZillizSchemaManager(MilvusSchemaManager):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        import warnings
+
+        warnings.warn(
+            "The ZillizSchemaManager class will be deprecated in the future. "
+            "Please use the MilvusSchemaManager class instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(*args, **kwargs)    

crewplus/vectorstores/milvus/schema_milvus.py

@@ -0,0 +1,253 @@
+from typing import List, Optional
+import logging
+import json
+
+from pymilvus import DataType
+from langchain_milvus import Milvus
+from langchain_core.documents import Document
+from crewplus.utils.schema_document_updater import SchemaDocumentUpdater
+from crewplus.utils.schema_action import Action
+from .milvus_schema_manager import MilvusSchemaManager
+
+
+class SchemaMilvus(Milvus):
+    """
+    SchemaMilvus is a subclass of the Milvus class from langchain_milvus. This class is responsible for updating metadata of documents in a Milvus vector store.
+
+    Attributes:
+        embedding_function: Embedding function used by the Milvus vector store.
+        collection_name: Name of the collection in the Milvus vector store.
+        connection_args: Connection arguments for the Milvus vector store.
+        index_params: Index parameters for the Milvus vector store.
+        auto_id: Flag to specify if auto ID generation is enabled.
+        primary_field: The primary field of the collection.
+        vector_field: The vector field of the collection.
+        consistency_level: The consistency level for the Milvus vector store.
+        collection_schema: Schema JSON string associated with the Milvus existing collection name.
+    """
+    def __init__(
+        self, 
+        embedding_function, 
+        collection_name, 
+        connection_args, 
+        index_params=None, 
+        auto_id=True, 
+        primary_field="pk", 
+        text_field: str = "text",
+        vector_field=["vector"], 
+        consistency_level="Session",
+        logger: Optional[logging.Logger] = None
+    ):
+        """
+        Initializes the SchemaMilvus class with the provided parameters.
+
+        Args:
+            embedding_function: Embedding function used by the Milvus vector store.
+            collection_name: Name of the collection in the Milvus vector store.
+            connection_args: Connection arguments for the Milvus vector store.
+            index_params: Index parameters for the Milvus vector store.
+            auto_id: Flag to specify if auto ID generation is enabled.
+            primary_field: The primary field of the collection.
+            text_field: The text field of the collection.
+            vector_field: The vector field of the collection.
+            consistency_level: The consistency level for the Milvus vector store.
+            logger: Optional logger instance. If not provided, a default logger is created.
+        """
+        super().__init__(
+            embedding_function=embedding_function,
+            collection_name=collection_name,
+            connection_args=connection_args,
+            index_params=index_params,
+            auto_id=auto_id,
+            primary_field=primary_field,
+            text_field=text_field,
+            vector_field=vector_field,
+            consistency_level=consistency_level
+        )
+        self.logger = logger or logging.getLogger(__name__)
+        self.collection_schema = None
+        self.schema_manager = MilvusSchemaManager(client=self.client)
+
+    def set_schema(self, schema: str):
+        """
+        Sets the collection schema.
+
+        Args:
+            schema: The schema JSON string.
+        """
+        self.collection_schema = schema
+    
+    def get_fields(self, collection_name: Optional[str] = None) -> Optional[List[str]]:
+        """
+        Retrieves and returns the fields from the collection schema.
+
+        Args:
+            collection_name: The name of the collection to describe. If None, use self.collection_name.
+
+        Returns:
+            List[str] | None: The list of field names from the collection schema (excluding vector and text fields), or None if collection_name is not provided or an error occurs.
+        """
+        if collection_name is None:
+            collection_name = self.collection_name
+        if collection_name is None:
+            return None
+
+        try:
+            schema = self.client.describe_collection(collection_name)
+            fields = [field["name"] for field in schema["fields"] if field["type"] != DataType.FLOAT_VECTOR ]
+            return fields
+        except Exception as e:
+            self.logger.warning(f"Failed to retrieve schema fields: {e}")
+            return None
+    
+    def create_collection(self) -> bool:
+        """
+        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_client(self.client)
+        if not self.schema_manager.validate_schema(self.collection_schema):
+            self.logger.error("Failed to validate schema")
+            return False
+        try:
+            self.schema_manager.create_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: {e}")
+            return False
+
+    def drop_collection(self, collection_name: Optional[str] = None) -> bool:
+        """
+        Drops the collection using the Milvus client.
+
+        Returns:
+            bool: True if the collection is successfully dropped, False otherwise.
+        """
+        if collection_name is None:
+            collection_name = self.collection_name
+            
+        try:
+            self.client.drop_collection(collection_name)
+            self.logger.info(f"Collection {collection_name} dropped successfully")
+            return True
+        except Exception as e:
+            self.logger.error(f"Failed to drop collection {self.collection_name}: {e}")
+            return False
+
+    def _handle_upsert(self, doc: Document, metadata_dict: dict) -> Document:
+        """
+        Handles the UPSERT action for a single document by merging metadata.
+        """
+        existing_metadata = doc.metadata
+        for key, value in metadata_dict.items():
+            # Skip primary key and text fields to prevent modification.
+            if key in [self.primary_field, self.text_field]:
+                continue
+
+            if isinstance(value, dict):
+                # If the new value is a dictionary, handle nested updates.
+                if key not in existing_metadata or not isinstance(existing_metadata.get(key), dict):
+                    # If the key doesn't exist or its value is not a dict, replace it.
+                    existing_metadata[key] = value
+                else:
+                    # If both are dictionaries, recursively update the nested fields.
+                    for sub_key, sub_value in value.items():
+                        if isinstance(sub_value, dict) and sub_key in existing_metadata[key] and isinstance(existing_metadata[key].get(sub_key), dict):
+                            existing_metadata[key][sub_key].update(sub_value)
+                        else:
+                            existing_metadata[key][sub_key] = sub_value
+            else:
+                # For non-dictionary values, simply update or add the field.
+                existing_metadata[key] = value
+
+        doc.metadata = existing_metadata
+        return doc
+
+    def update_documents_metadata(self, expr: str, metadata: str,action:Action=Action.UPSERT) -> List[Document]:
+        """
+        Updates the metadata of documents in the Milvus vector store based on the provided expression.
+
+        Args:
+            expr: Expression to filter the target documents.
+            metadata: New metadata to update the documents with.
+
+        Returns:
+            List of updated documents.
+        """
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise ValueError("Invalid JSON string for metadata")
+        
+        # Retrieve documents that match the filter expression.
+        fields = self.get_fields()
+        documents = self.search_by_metadata(expr, fields=fields, limit=5000)
+        
+        updated_documents = []
+        for doc in documents:
+            # Preserve the original primary key and text values.
+            pk_value = doc.metadata.get(self.primary_field) # default to pk
+            text_value = doc.metadata.get(self.text_field)
+
+            # Apply the specified action to update the document's metadata.
+            if action == Action.UPSERT:
+                doc = self._handle_upsert(doc, metadata_dict)
+            elif action == Action.DELETE:
+                keys_to_delete = metadata_dict.keys()
+                doc = SchemaDocumentUpdater.delete_document_metadata(doc, list(keys_to_delete))
+            elif action == Action.UPDATE:
+                existing_metadata = doc.metadata
+                update_dict = {}
+                for key, value in metadata_dict.items():
+                    if key in existing_metadata:
+                        if isinstance(value, dict) and isinstance(existing_metadata[key], dict):
+                            merged = existing_metadata[key].copy()
+                            for sub_key, sub_value in value.items():
+                                if sub_key in merged:
+                                    merged[sub_key] = sub_value
+                            update_dict[key] = merged
+                        else:
+                            update_dict[key] = value
+                doc = SchemaDocumentUpdater.update_document_metadata(doc, update_dict)
+            elif action == Action.INSERT:
+                existing_metadata = doc.metadata
+                for key, value in metadata_dict.items():
+                    if key in ['pk', 'text']:
+                        continue
+
+                    if isinstance(value, dict) and key in existing_metadata and isinstance(existing_metadata[key], dict):
+                        existing_metadata[key] = {}
+                        existing_metadata[key] = value
+                    else:
+                        existing_metadata[key] = value
+                doc.metadata = existing_metadata
+
+            # Restore the primary key and text values to ensure they are not lost.
+            if pk_value is not None:
+                doc.metadata[self.primary_field] = pk_value
+            if text_value is not None:
+                doc.metadata[self.text_field] = text_value
+
+            updated_documents.append(doc)
+
+        # Extract the primary keys for the upsert operation.
+        updated_ids = [doc.metadata[self.primary_field] for doc in updated_documents]
+        
+        # Remove primary key and text from metadata before upserting,
+        # as they are handled separately by the vector store.
+        for doc in updated_documents:
+            doc.metadata.pop(self.primary_field, None)
+            doc.metadata.pop(self.text_field, None)
+        
+        # Perform the upsert operation to update the documents in the collection.
+        self.upsert(ids=updated_ids, documents=updated_documents)
+        
+        return updated_documents

crewplus/vectorstores/milvus/vdb_service.py

@@ -0,0 +1,342 @@
+# -*- 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): A dictionary containing configuration for the vector store
+                         and embedding models.
+        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 from the settings dictionary.
+        NotImplementedError: If an unsupported provider is specified.
+        RuntimeError: If the MilvusClient fails to initialize after a retry.
+
+    Example:
+        >>> 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)
+        >>> # 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, schema: str = None, logger: logging.Logger = None):
+        """
+        Initializes the VDBService.
+        
+        Args:
+            settings (dict): Configuration dictionary for the service.
+            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__)
+        self.settings = settings
+
+        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, 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(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 caches instances by collection name to avoid re-instantiation.
+
+        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()
+
+        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) -> None:
+        """ Delete old indexes of the same source_url
+
+        Args:
+            url (str): source url
+        """
+        if url is None or vdb is None:
+            return
+
+        # Delete indexes of the same source_url
+        expr = "source in [\"" + url + "\"]"
+        pks = vdb.get_pks(expr)
+
+        # Delete entities by pks
+        if pks is not None and len(pks) > 0 :
+            old_items = vdb.delete(pks)
+            self.logger.info("ingesting document -- delete old indexes -- " + str(old_items))
+
+    def delete_old_indexes_by_id(self, id: str = None, vdb: Zilliz = None) -> None:
+        """ Delete old indexes of the same source_id
+
+        Args:
+            id (str): source id
+        """
+        self.logger.info(f"Delete old indexes of the same source_id:{id}")
+
+        if id is None or vdb is None:
+            return
+
+        # Delete indexes of the same source_id
+        expr = "source_id in [\"" + id + "\"]"
+        pks = vdb.get_pks(expr)
+
+        # Delete entities by pks
+        if pks is not None and len(pks) > 0 :
+            old_items = vdb.delete(pks)
+            self.logger.info("ingesting document -- delete old indexes -- " + str(old_items))
+
+    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 a collection
+
+        Args:
+            collection_name (str): scollection_name
+        """
+        self.logger.info(f"drop a collection by name:{collection_name}")
+
+        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.1.6.dist-info → crewplus-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.1.6
+Version: 0.2.1
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -15,6 +15,7 @@ Requires-Dist: google-genai==1.21.1
 Requires-Dist: mkdocs<2.0.0,>=1.6.1
 Requires-Dist: mkdocs-material<10.0.0,>=9.6.14
 Requires-Dist: mkdocstrings-python<2.0.0,>=1.16.12
+Requires-Dist: langchain-milvus<0.3.0,>=0.2.1
 Description-Content-Type: text/markdown
 
 # CrewPlus
@@ -37,6 +38,7 @@ CrewPlus is designed as a modular and extensible ecosystem of packages. This all
 -   **`crewplus` (This package):** The core package containing foundational services for chat, model load balancing, and vector stores.
 -   **`crewplus-agents`:** An extension for creating and managing autonomous AI agents.
 -   **`crewplus-ingestion`:** Provides robust pipelines for knowledge ingestion and data processing.
+-   **`crewplus-memory`:** Provides agent memory services for Crewplus AI Agents.
 -   **`crewplus-integrations`:** A collection of third-party integrations to connect CrewPlus with other services and platforms.
 
 ## Features
@@ -94,6 +96,10 @@ crewplus-base/                    # GitHub repo name
 │       └──  gemini_chat_model.py
 │       └──  model_load_balancer.py
 │       └──  ...
+│   └──  vectorstores/milvus
+│       └──  __init__.py
+│       └──  schema_milvus.py
+│       └──  vdb_service.py
 │   └──  core/
 │       └──  __init__.py
 │       └──  config.py

crewplus-0.2.1.dist-info/RECORD

@@ -0,0 +1,16 @@
+crewplus-0.2.1.dist-info/METADATA,sha256=IKVh7qW2mV9NfYXs66enxgcBOk_GTP9xeU3kwI6lysc,4881
+crewplus-0.2.1.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+crewplus-0.2.1.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.1.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
+crewplus/services/init_services.py,sha256=5MlvZbyGPNgVA_2bPuxhhbw9TgD5rrvC0e_62YJh340,589
+crewplus/services/model_load_balancer.py,sha256=a2BfcWV-OEItVf7_tTBYY0HiXxFzA5Uk7d-6AG-jLNU,8692
+crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
+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=KiGuHWU9oz2QCCaxGECtN-F69m8ZOwjzPO0umk6ZjzA,14592
+crewplus-0.2.1.dist-info/RECORD,,

crewplus-0.1.6.dist-info/RECORD

@@ -1,9 +0,0 @@
-crewplus-0.1.6.dist-info/METADATA,sha256=lqdkw41PrxeXf9IJRMEkCZrxw1B-3jLyEmBwlU2iT-k,4609
-crewplus-0.1.6.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-crewplus-0.1.6.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.1.6.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
-crewplus/services/model_load_balancer.py,sha256=bJpSgCGPWWT1yD_nYshIPngr8Xmdq1gfq8lJ1hOEGbM,8673
-crewplus-0.1.6.dist-info/RECORD,,