crewplus 0.2.74__py3-none-any.whl → 0.2.76__py3-none-any.whl

crewplus/services/model_load_balancer.py

@@ -68,7 +68,7 @@ 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, selection_strategy: str = 'random'):
+    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None, with_metadata: bool = False, selection_strategy: str = 'random', disable_streaming: bool = False):
         """
         Get a model instance.
 
@@ -82,6 +82,7 @@ class ModelLoadBalancer:
             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'.
+            disable_streaming: If True, get a model instance with streaming disabled.
 
         Returns:
             An instantiated language model object, or a tuple if with_metadata is True.
@@ -97,7 +98,7 @@ class ModelLoadBalancer:
         if deployment_name:
             for model_config in self.models_config:
                 if model_config.get('deployment_name') == deployment_name:
-                    model = self._get_or_create_model(model_config)
+                    model = self._get_or_create_model(model_config, disable_streaming)
                     if with_metadata:
                         return model, deployment_name
                     return model
@@ -121,7 +122,7 @@ class ModelLoadBalancer:
                 self.logger.warning(f"Unsupported selection strategy: '{selection_strategy}'. Defaulting to 'random'.")
                 selected_model_config = self._random_selection(candidates)
                 
-            model = self._get_or_create_model(selected_model_config)
+            model = self._get_or_create_model(selected_model_config, disable_streaming)
             if with_metadata:
                 return model, selected_model_config.get('deployment_name')
             return model
@@ -134,21 +135,25 @@ class ModelLoadBalancer:
             self.thread_local.models_cache = {}
         return self.thread_local.models_cache
 
-    def _get_or_create_model(self, model_config: Dict):
+    def _get_or_create_model(self, model_config: Dict, disable_streaming: bool = False):
         """
         Gets a model instance from the thread-local cache. If it doesn't exist,
         it instantiates, caches, and returns it.
         """
         model_id = model_config['id']
+        cache_key = f"{model_id}"
+        if disable_streaming:
+            cache_key += "-non-streaming"
+            
         models_cache = self._get_thread_local_models_cache()
 
-        if model_id not in models_cache:
-            self.logger.debug(f"Creating new model instance for id {model_id} in thread {threading.get_ident()}")
-            models_cache[model_id] = self._instantiate_model(model_config)
+        if cache_key not in models_cache:
+            self.logger.debug(f"Creating new model instance for id {cache_key} in thread {threading.get_ident()}")
+            models_cache[cache_key] = self._instantiate_model(model_config, disable_streaming)
         
-        return models_cache[model_id]
+        return models_cache[cache_key]
 
-    def _instantiate_model(self, model_config: Dict):
+    def _instantiate_model(self, model_config: Dict, disable_streaming: bool = False):
         """Instantiate and return an LLM object based on the model configuration"""
         provider = model_config['provider']
         self.logger.debug(f"Model balancer: instantiating {provider} -- {model_config.get('deployment_name')}")
@@ -162,8 +167,13 @@ class ModelLoadBalancer:
             }
             if 'temperature' in model_config:
                 kwargs['temperature'] = model_config['temperature']
-            if model_config.get('deployment_name') == 'o1-mini':
+            
+            # The 'disable_streaming' parameter takes precedence
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
+            elif model_config.get('deployment_name') == 'o1-mini':
                 kwargs['disable_streaming'] = True
+                
             return TracedAzureChatOpenAI(**kwargs)
         elif provider == 'openai':
             kwargs = {
@@ -197,6 +207,8 @@ class ModelLoadBalancer:
                 kwargs['temperature'] = model_config['temperature']
             if 'max_tokens' in model_config:
                 kwargs['max_tokens'] = model_config['max_tokens']
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
             return GeminiChatModel(**kwargs)
         elif provider == 'vertex-ai':
             deployment_name = model_config['deployment_name']
@@ -216,6 +228,8 @@ class ModelLoadBalancer:
                 kwargs['temperature'] = model_config['temperature']
             if 'max_tokens' in model_config:
                 kwargs['max_tokens'] = model_config['max_tokens']
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
             return GeminiChatModel(**kwargs)
         else:
             self.logger.error(f"Unsupported provider: {provider}")

{crewplus-0.2.74.dist-info → crewplus-0.2.76.dist-info}/METADATA

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

{crewplus-0.2.74.dist-info → crewplus-0.2.76.dist-info}/RECORD

@@ -1,7 +1,7 @@
-crewplus-0.2.74.dist-info/METADATA,sha256=PAa0VeDtSbp3D9WqgVplcDaiCmWoo0H-HcJU8niPHiI,5424
-crewplus-0.2.74.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.74.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.74.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.76.dist-info/METADATA,sha256=8TMsBVGxXiVMSCQ4QHkikVTWgQWbW2zweYlpVGKnTrY,5424
+crewplus-0.2.76.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.76.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.76.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/callbacks/__init__.py,sha256=YG7ieeb91qEjp1zF0-inEN7mjZ7yT_D2yzdWFT8Z1Ws,63
 crewplus/callbacks/async_langfuse_handler.py,sha256=8_p7ctgcmDNQgF5vOqA47I0x-3GWsm7zioZcZHgedZk,7163
@@ -9,7 +9,7 @@ crewplus/services/__init__.py,sha256=V1CG8b2NOmRzNgQH7BPl4KVxWSYJH5vfEsW1wVErKNE
 crewplus/services/azure_chat_model.py,sha256=iWzJ2GQFSNmwJx-2O5_xKPSB6VVc-7T6bcfFI8_WezA,5521
 crewplus/services/gemini_chat_model.py,sha256=DYqz01H2TIHiCDQesSozVfOsMigno6QGwOtIweg7UHk,40103
 crewplus/services/init_services.py,sha256=tc1ti8Yufo2ixlJpwg8uH0KmoyQ4EqxCOe4uTEWnlRM,2413
-crewplus/services/model_load_balancer.py,sha256=Q9Gx3GrbKworU-Ytxeqp0ggHSgZ1Q6brtTk-nCl4sak,12095
+crewplus/services/model_load_balancer.py,sha256=HIx-k-FiizJSF4e88SFxfFVNS93vJR2zrOdU_fg26FU,12826
 crewplus/services/tracing_manager.py,sha256=pwNFeA77vnoZMh_AUOnK5TvAaPOOLg5oDnVOe1yUa9A,8502
 crewplus/utils/__init__.py,sha256=2Gk1n5srFJQnFfBuYTxktdtKOVZyNrFcNaZKhXk35Pw,142
 crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
@@ -23,4 +23,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.74.dist-info/RECORD,,
+crewplus-0.2.76.dist-info/RECORD,,