format

Author: patricebechard

src/agentlab/llm/chat_api.py

@@ -6,13 +6,12 @@
 from functools import partial
 from typing import Optional
 
+import agentlab.llm.tracking as tracking
 import anthropic
 import openai
-from openai import NOT_GIVEN, OpenAI
-
-import agentlab.llm.tracking as tracking
 from agentlab.llm.base_api import AbstractChatModel, BaseModelArgs
 from agentlab.llm.llm_utils import AIMessage, Discussion
+from openai import NOT_GIVEN, OpenAI
 
 
 def make_system_message(content: str) -> dict:
@@ -89,6 +88,7 @@ def make_model(self):
             log_probs=self.log_probs,
         )
 
+
 @dataclass
 class LiteLLMModelArgs(BaseModelArgs):
 
@@ -119,9 +119,7 @@ def make_model(self):
 class AzureModelArgs(BaseModelArgs):
     """Serializable object for instantiating a generic chat model with an Azure model."""
 
-    deployment_name: str = (
-        None  # NOTE: deployment_name is deprecated for Azure OpenAI and won't be used.
-    )
+    deployment_name: str = None  # NOTE: deployment_name is deprecated for Azure OpenAI and won't be used.
 
     def make_model(self):
         return AzureChatModel(
@@ -219,9 +217,7 @@ class RetryError(Exception):
 def handle_error(error, itr, min_retry_wait_time, max_retry):
     if not isinstance(error, openai.OpenAIError):
         raise error
-    logging.warning(
-        f"Failed to get a response from the API: \n{error}\n" f"Retrying... ({itr+1}/{max_retry})"
-    )
+    logging.warning(f"Failed to get a response from the API: \n{error}\n" f"Retrying... ({itr+1}/{max_retry})")
     wait_time = _extract_wait_time(
         error.args[0],
         min_retry_wait_time=min_retry_wait_time,
@@ -320,18 +316,13 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                 self.error_types.append(error_type)
 
         if not completion:
-            raise RetryError(
-                f"Failed to get a response from the API after {self.max_retry} retries\n"
-                f"Last error: {error_type}"
-            )
+            raise RetryError(f"Failed to get a response from the API after {self.max_retry} retries\n" f"Last error: {error_type}")
 
         input_tokens = completion.usage.prompt_tokens
         output_tokens = completion.usage.completion_tokens
         cost = input_tokens * self.input_cost + output_tokens * self.output_cost
 
-        if hasattr(tracking.TRACKER, "instance") and isinstance(
-            tracking.TRACKER.instance, tracking.LLMTracker
-        ):
+        if hasattr(tracking.TRACKER, "instance") and isinstance(tracking.TRACKER.instance, tracking.LLMTracker):
             tracking.TRACKER.instance(input_tokens, output_tokens, cost)
 
         if n_samples == 1:
@@ -404,6 +395,7 @@ def __init__(
             log_probs=log_probs,
         )
 
+
 class AzureChatModel(ChatModel):
     def __init__(
         self,
@@ -417,18 +409,12 @@ def __init__(
         log_probs=False,
     ):
         api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
-        assert (
-            api_key
-        ), "AZURE_OPENAI_API_KEY has to be defined in the environment when using AzureChatModel"
+        assert api_key, "AZURE_OPENAI_API_KEY has to be defined in the environment when using AzureChatModel"
         endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
-        assert (
-            endpoint
-        ), "AZURE_OPENAI_ENDPOINT has to be defined in the environment when using AzureChatModel"
+        assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment when using AzureChatModel"
 
         if deployment_name is not None:
-            logging.info(
-                f"Deployment name is deprecated for Azure OpenAI and won't be used. Using model name: {model_name}."
-            )
+            logging.info(f"Deployment name is deprecated for Azure OpenAI and won't be used. Using model name: {model_name}.")
 
         client_args = {
             "base_url": endpoint,
@@ -560,12 +546,8 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                 output_tokens = getattr(usage, "output_tokens", 0)
                 cache_read_tokens = getattr(usage, "cache_input_tokens", 0)
                 cache_write_tokens = getattr(usage, "cache_creation_input_tokens", 0)
-                cache_read_cost = (
-                    self.input_cost * tracking.ANTHROPIC_CACHE_PRICING_FACTOR["cache_read_tokens"]
-                )
-                cache_write_cost = (
-                    self.input_cost * tracking.ANTHROPIC_CACHE_PRICING_FACTOR["cache_write_tokens"]
-                )
+                cache_read_cost = self.input_cost * tracking.ANTHROPIC_CACHE_PRICING_FACTOR["cache_read_tokens"]
+                cache_write_cost = self.input_cost * tracking.ANTHROPIC_CACHE_PRICING_FACTOR["cache_write_tokens"]
                 cost = (
                     new_input_tokens * self.input_cost
                     + output_tokens * self.output_cost
@@ -574,9 +556,7 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                 )
 
                 # Track usage if available
-                if hasattr(tracking.TRACKER, "instance") and isinstance(
-                    tracking.TRACKER.instance, tracking.LLMTracker
-                ):
+                if hasattr(tracking.TRACKER, "instance") and isinstance(tracking.TRACKER.instance, tracking.LLMTracker):
                     tracking.TRACKER.instance(new_input_tokens, output_tokens, cost)
 
                 return AIMessage(response.content[0].text)
@@ -613,14 +593,8 @@ def __init__(
         self.max_tokens = max_tokens
         self.max_retry = max_retry
 
-        if (
-            not os.getenv("AWS_REGION")
-            or not os.getenv("AWS_ACCESS_KEY")
-            or not os.getenv("AWS_SECRET_KEY")
-        ):
-            raise ValueError(
-                "AWS_REGION, AWS_ACCESS_KEY and AWS_SECRET_KEY must be set in the environment when using BedrockChatModel"
-            )
+        if not os.getenv("AWS_REGION") or not os.getenv("AWS_ACCESS_KEY") or not os.getenv("AWS_SECRET_KEY"):
+            raise ValueError("AWS_REGION, AWS_ACCESS_KEY and AWS_SECRET_KEY must be set in the environment when using BedrockChatModel")
 
         self.client = anthropic.AnthropicBedrock(
             aws_region=os.getenv("AWS_REGION"),
@@ -638,6 +612,7 @@ def make_model(self):
             max_tokens=self.max_new_tokens,
         )
 
+
 class LiteLLMChatModel(AbstractChatModel):
     def __init__(
         self,
@@ -661,7 +636,6 @@ def __init__(
         self.max_retry = max_retry
         self.min_retry_wait_time = min_retry_wait_time
         self.log_probs = log_probs
-        self.reasoning_effort = reasoning_effort
 
         # Get pricing information
         if pricing_func:
@@ -679,9 +653,9 @@ def __init__(
             self.input_cost = 0.0
             self.output_cost = 0.0
 
-
     def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float = None) -> dict:
         from litellm import completion as litellm_completion
+
         # Initialize retry tracking attributes
         self.retries = 0
         self.success = False
@@ -696,10 +670,6 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                 completion = litellm_completion(
                     model=self.model_name,
                     messages=messages,
-                    # n=n_samples,
-                    # temperature=temperature,
-                    # max_completion_tokens=self.max_tokens,
-                    reasoning_effort=self.reasoning_effort,
                 )
 
                 if completion.usage is None:
@@ -714,18 +684,13 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                 self.error_types.append(error_type)
 
         if not completion:
-            raise RetryError(
-                f"Failed to get a response from the API after {self.max_retry} retries\n"
-                f"Last error: {error_type}"
-            )
+            raise RetryError(f"Failed to get a response from the API after {self.max_retry} retries\n" f"Last error: {error_type}")
 
         input_tokens = completion.usage.prompt_tokens
         output_tokens = completion.usage.completion_tokens
         cost = input_tokens * self.input_cost + output_tokens * self.output_cost
 
-        if hasattr(tracking.TRACKER, "instance") and isinstance(
-            tracking.TRACKER.instance, tracking.LLMTracker
-        ):
+        if hasattr(tracking.TRACKER, "instance") and isinstance(tracking.TRACKER.instance, tracking.LLMTracker):
             tracking.TRACKER.instance(input_tokens, output_tokens, cost)
 
         if n_samples == 1:
@@ -745,4 +710,4 @@ def get_stats(self):
         return {
             "n_retry_llm": self.retries,
             # "busted_retry_llm": int(not self.success), # not logged if it occurs anyways
-        }
+        }

src/agentlab/llm/llm_configs.py

@@ -1,14 +1,13 @@
-from openai import NOT_GIVEN
-
 from agentlab.llm.chat_api import (
     AnthropicModelArgs,
     AzureModelArgs,
     BedrockModelArgs,
+    LiteLLMModelArgs,
     OpenAIModelArgs,
     OpenRouterModelArgs,
     SelfHostedModelArgs,
-    LiteLLMModelArgs,
 )
+from openai import NOT_GIVEN
 
 default_oss_llms_args = {
     "n_retry_server": 4,