crewplus 0.2.29__py3-none-any.whl → 0.2.30__py3-none-any.whl

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, with_metadata: bool = False):
+    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.
@@ -118,7 +120,16 @@ 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']
             model = self.models[model_id]
             if with_metadata:
@@ -186,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
@@ -197,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.29.dist-info → crewplus-0.2.30.dist-info}/METADATA

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

{crewplus-0.2.29.dist-info → crewplus-0.2.30.dist-info}/RECORD

@@ -1,13 +1,13 @@
-crewplus-0.2.29.dist-info/METADATA,sha256=Gjp8NS1iow7bCkTVYioXaFmWPr6RKFBqiCrMEzTdVJg,5327
-crewplus-0.2.29.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.29.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.29.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.30.dist-info/METADATA,sha256=1Mhnmb9kTEFG0dyGfMTiGaBrNpJxx36lPGg8GmhsmQc,5327
+crewplus-0.2.30.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.30.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.30.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/services/__init__.py,sha256=zUM4ZwUfGMBDx-j7Wehf_KC5yYXPTK8BK_oeO5veIXQ,398
 crewplus/services/azure_chat_model.py,sha256=WMSf4BDO8UcP7ZASNGRJxdTEnuWBmCRSY_4yx_VMbok,5499
 crewplus/services/gemini_chat_model.py,sha256=oVLL07VEjwrHOeP56YSPLnldbfjvTVRrsTI6xcPNt1E,35224
 crewplus/services/init_services.py,sha256=EBpDkIwzuujmdlqjyWvdLQCfhQmfS_OKFz-9Ji8nmAU,1628
-crewplus/services/model_load_balancer.py,sha256=meJlmyJoKPXA6AlIYiVJETd8bS6zeaDevMuVLs6LFhk,9715
+crewplus/services/model_load_balancer.py,sha256=2LOiPUslRKreBZTndKOx6WH8VtJT8mep37GsDINQTJs,10780
 crewplus/services/tracing_manager.py,sha256=aCU9N4Jvh8pDD3h8kWX4O-Ax8xwdLHnQ4wJ3sf-vLwA,6289
 crewplus/utils/__init__.py,sha256=2Gk1n5srFJQnFfBuYTxktdtKOVZyNrFcNaZKhXk35Pw,142
 crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
@@ -20,4 +20,4 @@ docs/GeminiChatModel.md,sha256=zZYyl6RmjZTUsKxxMiC9O4yV70MC4TD-IGUmWhIDBKA,8677
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.29.dist-info/RECORD,,
+crewplus-0.2.30.dist-info/RECORD,,