crewplus 0.2.28__tar.gz → 0.2.30__tar.gz

{crewplus-0.2.28 → crewplus-0.2.30}/PKG-INFO

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

{crewplus-0.2.28 → crewplus-0.2.30}/crewplus/services/model_load_balancer.py

@@ -76,21 +76,23 @@ class ModelLoadBalancer:
             self.logger.error(f"Failed to load model configuration: {e}", exc_info=True)
             raise RuntimeError(f"Failed to load model configuration: {e}")
 
-    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None):
+    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None, with_metadata: bool = False, selection_strategy: str = 'random'):
         """
         Get a model instance.
 
         Can fetch a model in two ways:
         1. By its specific `deployment_name`.
-        2. By `provider` and `model_type`, which will select a model using round-robin.
+        2. By `provider` and `model_type`, which will select a model using a specified strategy.
 
         Args:
             provider: The model provider (e.g., 'azure-openai', 'google-genai').
             model_type: The type of model (e.g., 'inference', 'embedding', 'embedding-large').
             deployment_name: The unique name for the model deployment.
+            with_metadata: If True, returns a tuple of (model, deployment_name).
+            selection_strategy: The selection strategy ('random', 'round_robin', or 'least_used'). Defaults to 'random'.
 
         Returns:
-            An instantiated language model object.
+            An instantiated language model object, or a tuple if with_metadata is True.
 
         Raises:
             RuntimeError: If the model configuration has not been loaded.
@@ -104,7 +106,11 @@ class ModelLoadBalancer:
             for model_config in self.models_config:
                 if model_config.get('deployment_name') == deployment_name:
                     model_id = model_config['id']
-                    return self.models[model_id]
+                    model = self.models[model_id]
+                    if with_metadata:
+                        return model, deployment_name
+                    return model
+                
             self.logger.error(f"No model found for deployment name: {deployment_name}")
             raise ValueError(f"No model found for deployment name: {deployment_name}")
 
@@ -114,9 +120,21 @@ class ModelLoadBalancer:
                 self.logger.error(f"No models found for provider '{provider}' and type '{model_type}'")
                 raise ValueError(f"No models found for provider '{provider}' and type '{model_type}'")
 
-            selected_model_config = self._round_robin_selection(candidates)
+            if selection_strategy == 'random':
+                selected_model_config = self._random_selection(candidates)
+            elif selection_strategy == 'round_robin':
+                selected_model_config = self._round_robin_selection(candidates)
+            elif selection_strategy == 'least_used':
+                selected_model_config = self._least_used_selection(candidates)
+            else:
+                self.logger.warning(f"Unsupported selection strategy: '{selection_strategy}'. Defaulting to 'random'.")
+                selected_model_config = self._random_selection(candidates)
+                
             model_id = selected_model_config['id']
-            return self.models[model_id]
+            model = self.models[model_id]
+            if with_metadata:
+                return model, selected_model_config.get('deployment_name')
+            return model
 
         raise ValueError("Either 'deployment_name' or both 'provider' and 'model_type' must be provided.")
 
@@ -179,6 +197,12 @@ class ModelLoadBalancer:
         self.usage_counter = defaultdict(int)
         self.current_indices = {}
 
+    def _random_selection(self, candidates: list) -> Dict:
+        """Selects a model randomly from a list of candidates."""
+        model = random.choice(candidates)
+        self.usage_counter[model['id']] += 1
+        return model
+
     def _round_robin_selection(self, candidates: list) -> Dict:
         if id(candidates) not in self.current_indices:
             self.current_indices[id(candidates)] = 0
@@ -190,8 +214,8 @@ class ModelLoadBalancer:
         return model
 
     def _least_used_selection(self, candidates: list) -> Dict:
-        min_usage = min(self.usage_counter[m['model_id']] for m in candidates)
-        least_used = [m for m in candidates if self.usage_counter[m['model_id']] == min_usage]
+        min_usage = min(self.usage_counter[m['id']] for m in candidates)
+        least_used = [m for m in candidates if self.usage_counter[m['id']] == min_usage]
         model = random.choice(least_used)
         self.usage_counter[model['id']] += 1
         return model

{crewplus-0.2.28 → crewplus-0.2.30}/crewplus/vectorstores/milvus/vdb_service.py

@@ -10,6 +10,7 @@ from langchain_milvus import Milvus
 from langchain_core.embeddings import Embeddings
 from langchain_openai import AzureOpenAIEmbeddings
 from pymilvus import MilvusClient
+import time
 
 from ...services.init_services import get_model_balancer
 from .schema_milvus import SchemaMilvus, DEFAULT_SCHEMA
@@ -361,10 +362,9 @@ class VDBService(object):
             "params": {}
         }
         
-        vdb = Milvus(
-            embedding_function=embeddings,
+        vdb = self._create_milvus_instance_with_retry(
             collection_name=collection_name,
-            connection_args=self.connection_args,
+            embeddings=embeddings,
             index_params=index_params
         )
 
@@ -373,6 +373,32 @@ class VDBService(object):
 
         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.
+        """
+        retries = 2
+        for attempt in range(retries + 1):
+            try:
+                vdb = Milvus(
+                    embedding_function=embeddings,
+                    collection_name=collection_name,
+                    connection_args=self.connection_args,
+                    index_params=index_params
+                )
+                self.logger.info(f"Successfully connected to Milvus for collection '{collection_name}' on attempt {attempt + 1}.")
+                return vdb  # Return on success
+            except Exception as e:
+                self.logger.warning(
+                    f"Attempt {attempt + 1}/{retries + 1} to connect to Milvus for collection '{collection_name}' failed: {e}"
+                )
+                if attempt < retries:
+                    self.logger.info("Retrying in 3 seconds...")
+                    time.sleep(3)
+                else:
+                    self.logger.error(f"Failed to connect to Milvus for collection '{collection_name}' after {retries + 1} attempts.")
+                    raise RuntimeError(f"Could not connect to Milvus after {retries + 1} attempts.") from e
+
     def drop_collection(self, collection_name: str) -> None:
         """
         Deletes a collection from the vector database and removes it from the cache.

{crewplus-0.2.28 → crewplus-0.2.30}/pyproject.toml

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