crewplus 0.2.36__tar.gz → 0.2.37__tar.gz

{crewplus-0.2.36 → crewplus-0.2.37}/PKG-INFO

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

{crewplus-0.2.36 → crewplus-0.2.37}/crewplus/services/model_load_balancer.py

@@ -1,6 +1,7 @@
 import json
 import random
 import logging
+import threading
 from typing import Dict, List, Optional, Union
 from collections import defaultdict
 from langchain_openai import ChatOpenAI, AzureOpenAIEmbeddings
@@ -31,7 +32,7 @@ class ModelLoadBalancer:
         self.config_data = config_data
         self.logger = logger or logging.getLogger(__name__)
         self.models_config: List[Dict] = []
-        self.models: Dict[int, Union[TracedAzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings, GeminiChatModel]] = {}
+        self.thread_local = threading.local()
         self._initialize_state()
         self._config_loaded = False  # Flag to check if config is loaded
 
@@ -60,15 +61,6 @@ class ModelLoadBalancer:
 
             self.models_config = config['models']
 
-            # Instantiate models
-            for model_config in self.models_config:
-                model_id = model_config['id']
-                model_instance = self._instantiate_model(model_config)
-                if model_instance is not None:
-                    self.models[model_id] = model_instance
-                else:
-                    self.logger.warning(f"Model with id {model_id} was not loaded due to instantiation error.")
-
             self._config_loaded = True
             self.logger.debug("Model balancer: configuration loaded successfully.")
         except (FileNotFoundError, json.JSONDecodeError, ValueError) as e:
@@ -105,8 +97,7 @@ class ModelLoadBalancer:
         if deployment_name:
             for model_config in self.models_config:
                 if model_config.get('deployment_name') == deployment_name:
-                    model_id = model_config['id']
-                    model = self.models[model_id]
+                    model = self._get_or_create_model(model_config)
                     if with_metadata:
                         return model, deployment_name
                     return model
@@ -130,14 +121,33 @@ class ModelLoadBalancer:
                 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]
+            model = self._get_or_create_model(selected_model_config)
             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.")
 
+    def _get_thread_local_models_cache(self) -> Dict:
+        """Gets the model cache for the current thread, creating it if it doesn't exist."""
+        if not hasattr(self.thread_local, 'models_cache'):
+            self.thread_local.models_cache = {}
+        return self.thread_local.models_cache
+
+    def _get_or_create_model(self, model_config: Dict):
+        """
+        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']
+        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)
+        
+        return models_cache[model_id]
+
     def _instantiate_model(self, model_config: Dict):
         """Instantiate and return an LLM object based on the model configuration"""
         provider = model_config['provider']

{crewplus-0.2.36 → crewplus-0.2.37}/pyproject.toml

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