Fix: Replace deprecated alias-fast with alias-large in fallback cascade

tinytroupe/openai_utils.py

@@ -31,6 +31,8 @@ class OpenAIClient:
     def __init__(self, cache_api_calls=default["cache_api_calls"], cache_file_name=default["cache_file_name"]) -> None:
         logger.debug("Initializing OpenAIClient")
 
+        self.client = None
+
         # should we cache api calls and reuse them?
         self.set_api_cache(cache_api_calls, cache_file_name)
     
@@ -52,7 +54,8 @@ class OpenAIClient:
         """
         Sets up the OpenAI API configurations for this client.
         """
-        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY", os.getenv("BLABLADOR_API_KEY", "dummy_token")), base_url=os.getenv("HELMHOLTZ_BLABLADOR_ENDPOINT", "https://api.helmholtz-blablador.fz-juelich.de/v1"))
+        if self.client is None:
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
     @config_manager.config_defaults(
         model="model",
@@ -141,7 +144,6 @@ class OpenAIClient:
             "messages": current_messages,
             "temperature": temperature,
             "max_tokens":max_tokens,
-            "top_p": top_p,
             "frequency_penalty": frequency_penalty,
             "presence_penalty": presence_penalty,
             "stop": stop,
@@ -150,18 +152,40 @@ class OpenAIClient:
             "n": n,
         }
 
+        if top_p is not None and top_p > 0:
+            chat_api_params["top_p"] = top_p
+
         if response_format is not None:
             chat_api_params["response_format"] = response_format
 
         i = 0
-        while i < max_attempts:
+        while True:
             try:
                 i += 1
 
+                #
+                # Model fallback and retry strategy requested by the user:
+                # 1. alias-fast for 3 attempts, 35s wait
+                # 2. alias-large for 2 attempts, 35s wait
+                # 3. alias-huge until success, 60s wait
+                #
+                # Model fallback strategy using config
+                if i <= 3:
+                    current_model = config["OpenAI"].get("MODEL", "alias-large")
+                    current_wait_time = 35
+                elif i <= 5:
+                    current_model = config["OpenAI"].get("FALLBACK_MODEL_LARGE", "alias-large")
+                    current_wait_time = 35
+                else:
+                    current_model = config["OpenAI"].get("FALLBACK_MODEL_HUGE", "alias-huge")
+                    current_wait_time = 60
+
+                chat_api_params["model"] = current_model
+
                 try:
-                    logger.debug(f"Sending messages to OpenAI API. Token count={self._count_tokens(current_messages, model)}.")
+                    logger.debug(f"Sending messages to OpenAI API. Model={current_model}. Token count={self._count_tokens(current_messages, current_model)}.")
                 except NotImplementedError:
-                    logger.debug(f"Token count not implemented for model {model}.")
+                    logger.debug(f"Token count not implemented for model {current_model}.")
                     
                 start_time = time.monotonic()
                 logger.debug(f"Calling model with client class {self.__class__.__name__}.")
@@ -169,15 +193,11 @@ class OpenAIClient:
                 ###############################################################
                 # call the model, either from the cache or from the API
                 ###############################################################
-                cache_key = str((model, chat_api_params)) # need string to be hashable
+                cache_key = str((current_model, chat_api_params)) # need string to be hashable
                 if self.cache_api_calls and (cache_key in self.api_cache):
                     response = self.api_cache[cache_key]
                 else:
-                    if waiting_time > 0:
-                        logger.info(f"Waiting {waiting_time} seconds before next API request (to avoid throttling)...")
-                        time.sleep(waiting_time)
-                    
-                    response = self._raw_model_call(model, chat_api_params)
+                    response = self._raw_model_call(current_model, chat_api_params)
                     if self.cache_api_calls:
                         self.api_cache[cache_key] = response
                         self._save_cache()
@@ -193,35 +213,21 @@ class OpenAIClient:
                 else:
                     return utils.sanitize_dict(self._raw_model_response_extractor(response))
 
-            except InvalidRequestError as e:
+            except (InvalidRequestError, openai.BadRequestError) as e:
                 logger.error(f"[{i}] Invalid request error, won't retry: {e}")
-
-                # there's no point in retrying if the request is invalid
-                # so we return None right away
                 return None
             
-            except openai.BadRequestError as e:
-                logger.error(f"[{i}] Invalid request error, won't retry: {e}")
-                
-                # there's no point in retrying if the request is invalid
-                # so we return None right away
-                return None
-            
-            except openai.RateLimitError:
-                logger.warning(
-                    f"[{i}] Rate limit error, waiting a bit and trying again.")
-                aux_exponential_backoff()
-            
-            except NonTerminalError as e:
-                logger.error(f"[{i}] Non-terminal error: {e}")
-                aux_exponential_backoff()
-                
-            except Exception as e:
-                logger.error(f"[{i}] {type(e).__name__} Error: {e}")
-                aux_exponential_backoff()
-
-        logger.error(f"Failed to get response after {max_attempts} attempts.")
-        return None
+            except (openai.RateLimitError,
+                    openai.APITimeoutError,
+                    openai.APIConnectionError,
+                    openai.InternalServerError,
+                    NonTerminalError,
+                    Exception) as e:
+                msg = f"[{i}] {type(e).__name__} Error with {current_model}: {e}. Waiting {current_wait_time} seconds before next attempt..."
+                logger.warning(msg)
+
+                time.sleep(current_wait_time)
+                continue
     
     def _raw_model_call(self, model, chat_api_params):
         """
@@ -244,8 +250,12 @@ class OpenAIClient:
             chat_api_params["reasoning_effort"] = default["reasoning_effort"]
 
 
-        # To make the log cleaner, we remove the messages from the logged parameters
-        logged_params = {k: v for k, v in chat_api_params.items() if k != "messages"} 
+        # To make the log cleaner, we remove the messages from the logged parameters,
+        # unless we are in debug mode
+        if logger.getEffectiveLevel() <= logging.DEBUG:
+            logged_params = chat_api_params
+        else:
+            logged_params = {k: v for k, v in chat_api_params.items() if k != "messages"}
 
         if "response_format" in chat_api_params:
             # to enforce the response format via pydantic, we need to use a different method
@@ -312,8 +322,8 @@ class OpenAIClient:
             elif "gpt-3.5-turbo" in model:
                 logger.debug("Token count: gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
                 return self._count_tokens(messages, model="gpt-3.5-turbo-0613")
-            elif ("gpt-4" in model) or ("ppo" in model) :
-                logger.debug("Token count: gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
+            elif ("gpt-4" in model) or ("ppo" in model) or ("alias-large" in model) or ("alias-huge" in model) or ("alias-large" in model):
+                logger.debug("Token count: gpt-4/alias-large may update over time. Returning num tokens assuming gpt-4-0613.")
                 return self._count_tokens(messages, model="gpt-4-0613")
             else:
                 raise NotImplementedError(
@@ -394,23 +404,40 @@ class AzureClient(OpenAIClient):
         Sets up the Azure OpenAI Service API configurations for this client,
         including the API endpoint and key.
         """
-        if os.getenv("AZURE_OPENAI_KEY"):
-            logger.info("Using Azure OpenAI Service API with key.")
-            self.client = AzureOpenAI(azure_endpoint= os.getenv("AZURE_OPENAI_ENDPOINT"),
-                                    api_version = config["OpenAI"]["AZURE_API_VERSION"],
-                                    api_key = os.getenv("AZURE_OPENAI_KEY"))
-        else:  # Use Entra ID Auth
-            logger.info("Using Azure OpenAI Service API with Entra ID Auth.")
-            from azure.identity import DefaultAzureCredential, get_bearer_token_provider
-
-            credential = DefaultAzureCredential()
-            token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
-            self.client = AzureOpenAI(
-                azure_endpoint= os.getenv("AZURE_OPENAI_ENDPOINT"),
-                api_version = config["OpenAI"]["AZURE_API_VERSION"],
-                azure_ad_token_provider=token_provider
+        if self.client is None:
+            if os.getenv("AZURE_OPENAI_KEY"):
+                logger.info("Using Azure OpenAI Service API with key.")
+                self.client = AzureOpenAI(azure_endpoint= os.getenv("AZURE_OPENAI_ENDPOINT"),
+                                        api_version = config["OpenAI"]["AZURE_API_VERSION"],
+                                        api_key = os.getenv("AZURE_OPENAI_KEY"))
+            else:  # Use Entra ID Auth
+                logger.info("Using Azure OpenAI Service API with Entra ID Auth.")
+                from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+                credential = DefaultAzureCredential()
+                token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+                self.client = AzureOpenAI(
+                    azure_endpoint= os.getenv("AZURE_OPENAI_ENDPOINT"),
+                    api_version = config["OpenAI"]["AZURE_API_VERSION"],
+                    azure_ad_token_provider=token_provider
+                )
+
+
+class HelmholtzBlabladorClient(OpenAIClient):
+
+    def __init__(self, cache_api_calls=default["cache_api_calls"], cache_file_name=default["cache_file_name"]) -> None:
+        logger.debug("Initializing HelmholtzBlabladorClient")
+        super().__init__(cache_api_calls, cache_file_name)
+
+    def _setup_from_config(self):
+        """
+        Sets up the Helmholtz Blablador API configurations for this client.
+        """
+        if self.client is None:
+            self.client = OpenAI(
+                base_url="https://api.helmholtz-blablador.fz-juelich.de/v1",
+                api_key=os.getenv("BLABLADOR_API_KEY", "dummy"),
             )
-    
 
 ###########################################################################
 # Exceptions
@@ -502,6 +529,7 @@ def force_api_cache(cache_api_calls, cache_file_name=default["cache_file_name"])
 # default client
 register_client("openai", OpenAIClient())
 register_client("azure", AzureClient())
+register_client("helmholtz-blablador", HelmholtzBlabladorClient())