add litellm chat model for genericagent

Author: patricebechard

src/agentlab/llm/chat_api.py

@@ -89,6 +89,18 @@ def make_model(self):
             log_probs=self.log_probs,
         )
 
+@dataclass
+class LiteLLMModelArgs(BaseModelArgs):
+
+    def make_model(self):
+        return LiteLLMChatModel(
+            model_name=self.model_name,
+            temperature=self.temperature,
+            max_tokens=self.max_new_tokens,
+            log_probs=self.log_probs,
+            reasoning_effort=self.reasoning_effort,
+        )
+
 
 @dataclass
 class OpenAIModelArgs(BaseModelArgs):
@@ -393,7 +405,6 @@ def __init__(
             log_probs=log_probs,
         )
 
-
 class AzureChatModel(ChatModel):
     def __init__(
         self,
@@ -627,3 +638,113 @@ def make_model(self):
             temperature=self.temperature,
             max_tokens=self.max_new_tokens,
         )
+
+class LiteLLMChatModel(AbstractChatModel):
+    def __init__(
+        self,
+        model_name,
+        api_key=None,
+        temperature=0.5,
+        max_tokens=100,
+        max_retry=4,
+        min_retry_wait_time=60,
+        api_key_env_var=None,
+        client_class=OpenAI,
+        client_args=None,
+        pricing_func=None,
+        log_probs=False,
+        reasoning_effort=None,
+    ):
+        assert max_retry > 0, "max_retry should be greater than 0"
+
+        self.model_name = model_name
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        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:
+            pricings = pricing_func()
+            try:
+                self.input_cost = float(pricings[model_name]["prompt"])
+                self.output_cost = float(pricings[model_name]["completion"])
+            except KeyError:
+                logging.warning(
+                    f"Model {model_name} not found in the pricing information, prices are set to 0. Maybe try upgrading langchain_community."
+                )
+                self.input_cost = 0.0
+                self.output_cost = 0.0
+        else:
+            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
+        self.error_types = []
+
+        completion = None
+        e = None
+        for itr in range(self.max_retry):
+            self.retries += 1
+            temperature = temperature if temperature is not None else self.temperature
+            try:
+                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:
+                    raise OpenRouterError(
+                        "The completion object does not contain usage information. This is likely a bug in the OpenRouter API."
+                    )
+
+                self.success = True
+                break
+            except openai.OpenAIError as e:
+                error_type = handle_error(e, itr, self.min_retry_wait_time, self.max_retry)
+                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}"
+            )
+
+        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
+        ):
+            tracking.TRACKER.instance(input_tokens, output_tokens, cost)
+
+        if n_samples == 1:
+            res_text = completion.choices[0].message.content
+            if res_text is not None:
+                res_text = res_text.removesuffix("<|end|>").strip()
+            else:
+                res_text = ""
+            res = AIMessage(res_text)
+            if self.log_probs:
+                res["log_probs"] = completion.choices[0].log_probs
+            return res
+        else:
+            return [AIMessage(c.message.content.removesuffix("<|end|>").strip()) for c in completion.choices]
+
+    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

@@ -7,6 +7,7 @@
     OpenAIModelArgs,
     OpenRouterModelArgs,
     SelfHostedModelArgs,
+    LiteLLMModelArgs,
 )
 
 default_oss_llms_args = {
@@ -200,6 +201,68 @@
         temperature=1,  # temperature param not supported by gpt-5
         vision_support=True,
     ),
+    "azure/gpt-5-high-2025-08-07": AzureModelArgs(
+        model_name="gpt-5",
+        max_total_tokens=400_000,
+        max_input_tokens=256_000,
+        max_new_tokens=128_000,
+        temperature=1,  # temperature param not supported by gpt-5
+        vision_support=True,
+        reasoning_effort="high",
+    ),
+    "azure/gpt-5-mini-high-2025-08-07": AzureModelArgs(
+        model_name="gpt-5-mini",
+        max_total_tokens=400_000,
+        max_input_tokens=256_000,
+        max_new_tokens=128_000,
+        temperature=1,  # temperature param not supported by gpt-5
+        vision_support=True,
+        reasoning_effort="high",
+    ),
+    "azure/gpt-5-nano-high-2025-08-07": AzureModelArgs(
+        model_name="gpt-5-nano",
+        max_total_tokens=400_000,
+        max_input_tokens=256_000,
+        max_new_tokens=128_000,
+        temperature=1,  # temperature param not supported by gpt-5
+        vision_support=True,
+        reasoning_effort="high",
+    ),
+    "azure/gpt-oss-120b": AzureModelArgs(
+        model_name="gpt-oss-120b",
+        max_total_tokens=200_000,
+        max_input_tokens=200_000,
+        max_new_tokens=100_000,
+        temperature=1,
+        vision_support=False,
+        reasoning_effort="low",
+    ),
+    "azure/o3-high-2025-04-16": AzureModelArgs(
+        model_name="o3",
+        max_total_tokens=200_000,
+        max_input_tokens=200_000,
+        max_new_tokens=100_000,
+        temperature=1,
+        vision_support=False,
+        reasoning_effort="high",
+    ),
+    "azure/o3-mini-2025-01-31": AzureModelArgs(
+        model_name="o3-mini",
+        max_total_tokens=200_000,
+        max_input_tokens=200_000,
+        max_new_tokens=100_000,
+        temperature=1,
+        vision_support=False,
+    ),
+    "azure/o3-mini-high-2025-01-31": AzureModelArgs(
+        model_name="o3-mini",
+        max_total_tokens=200_000,
+        max_input_tokens=200_000,
+        max_new_tokens=100_000,
+        temperature=1,
+        vision_support=False,
+        reasoning_effort="high",
+    ),
     # ---------------- Anthropic ----------------#
     "anthropic/claude-3-7-sonnet-20250219": AnthropicModelArgs(
         model_name="claude-3-7-sonnet-20250219",