crewplus 0.2.34__py3-none-any.whl → 0.2.37__py3-none-any.whl

crewplus/services/gemini_chat_model.py

@@ -294,6 +294,7 @@ class GeminiChatModel(BaseChatModel):
         if creds is None:
             # Get service account file from env if not provided
             sa_file = self.service_account_file or os.getenv("GCP_SERVICE_ACCOUNT_FILE")
+            self.logger.debug(f"Service account file: {sa_file}")
             if sa_file:
                 try:
                     creds = service_account.Credentials.from_service_account_file(
@@ -307,15 +308,21 @@ class GeminiChatModel(BaseChatModel):
         
         # If creds is still None, the client will use Application Default Credentials (ADC).
         
-        self._client = genai.Client(
-            project=self.project_id,
-            location=self.location,
-            credentials=creds,
-        )
-        self.logger.info(
-            f"Initialized GeminiChatModel with model: {self.model_name} for Vertex AI "
-            f"(Project: {self.project_id}, Location: {self.location})"
-        )
+        try:
+            self._client = genai.Client(
+                vertexai=True,
+                project=self.project_id,
+                location=self.location,
+                credentials=creds,
+            )
+            self.logger.info(
+                f"Initialized GeminiChatModel with model: {self.model_name} for Vertex AI "
+                f"(Project: {self.project_id}, Location: {self.location})"
+            )
+        except Exception as e:
+            error_msg = f"Failed to initialize GenAI Client for Vertex AI: {e}"
+            self.logger.error(error_msg, exc_info=True)
+            raise ValueError(error_msg)
 
     def get_model_identifier(self) -> str:
         """Return a string identifying this model for tracing and logging."""

crewplus/services/init_services.py

@@ -1,15 +1,27 @@
 import os
+import logging
+from typing import Optional
 from .model_load_balancer import ModelLoadBalancer
 
 model_balancer = None
 
-def init_load_balancer(config_path: str = None):
+def init_load_balancer(
+    config_path: Optional[str] = None, 
+    logger: Optional[logging.Logger] = None
+):
     """
     Initializes the global ModelLoadBalancer instance.
 
     This function is idempotent. If the balancer is already initialized,
     it does nothing. It follows a safe initialization pattern where the
     global instance is only assigned after successful configuration loading.
+    
+    Args:
+        config_path (Optional[str]): The path to the model configuration file.
+            If not provided, it's determined by the `MODEL_CONFIG_PATH`
+            environment variable, or defaults to "config/models_config.json".
+        logger (Optional[logging.Logger]): An optional logger instance to be
+            used by the model balancer.
     """
     global model_balancer
     if model_balancer is None:
@@ -20,7 +32,10 @@ def init_load_balancer(config_path: str = None):
         )
         try:
             # 1. Create a local instance first.
-            balancer = ModelLoadBalancer(final_config_path)
+            balancer = ModelLoadBalancer(
+                config_path=final_config_path, 
+                logger=logger
+            )
             # 2. Attempt to load its configuration.
             balancer.load_config()
             # 3. Only assign to the global variable on full success.

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.34.dist-info → crewplus-0.2.37.dist-info}/METADATA

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

{crewplus-0.2.34.dist-info → crewplus-0.2.37.dist-info}/RECORD

@@ -1,13 +1,13 @@
-crewplus-0.2.34.dist-info/METADATA,sha256=UNWEOkl43TfQMN5AtWXFlMKy2Fdy7zBv6oKfN8_y_qM,5327
-crewplus-0.2.34.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.34.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.34.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.37.dist-info/METADATA,sha256=BzOLmuTMYx92TsT05qZXnQVBRhJO6t6XMYbLDh3JqJc,5327
+crewplus-0.2.37.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.37.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.37.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/services/__init__.py,sha256=V1CG8b2NOmRzNgQH7BPl4KVxWSYJH5vfEsW1wVErKNE,375
 crewplus/services/azure_chat_model.py,sha256=WMSf4BDO8UcP7ZASNGRJxdTEnuWBmCRSY_4yx_VMbok,5499
-crewplus/services/gemini_chat_model.py,sha256=XlJeaic1lq31lMx1EPKCe4AlYaHDR57tXOdu8V8j6dU,39623
-crewplus/services/init_services.py,sha256=EBpDkIwzuujmdlqjyWvdLQCfhQmfS_OKFz-9Ji8nmAU,1628
-crewplus/services/model_load_balancer.py,sha256=PU3wn8lh6pGeVFn62SURW_1lIGXbUnAklM1EWcJMhLU,11752
+crewplus/services/gemini_chat_model.py,sha256=2M0kXbjnPfQm3O9BZaGzJFlHF0pBdAFR3kIRS1bALoI,39983
+crewplus/services/init_services.py,sha256=7oZ1GmesK32EDB_DYnTzW17MEpXjXK41_U_1pmqu_m4,2183
+crewplus/services/model_load_balancer.py,sha256=Q9Gx3GrbKworU-Ytxeqp0ggHSgZ1Q6brtTk-nCl4sak,12095
 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.34.dist-info/RECORD,,
+crewplus-0.2.37.dist-info/RECORD,,