Fix tests for Azure OpenAI models, add GPT-4.1-nano (#286)

* fix tests for openai azure models after changes. add gpt-4.1-nano

* update test models to 4.1 nano

* formatting

Author: patricebechard

src/agentlab/agents/generic_agent/__init__.py

@@ -28,8 +28,10 @@
     AGENT_AZURE_4o_MINI_VISION,
     AGENT_AZURE_41,
     AGENT_AZURE_41_MINI,
+    AGENT_AZURE_41_NANO,
     AGENT_AZURE_41_VISION,
     AGENT_AZURE_41_MINI_VISION,
+    AGENT_AZURE_41_NANO_VISION,
     AGENT_AZURE_5,
     AGENT_AZURE_5_MINI,
     AGENT_AZURE_5_NANO,
@@ -66,8 +68,10 @@
     "AGENT_AZURE_4o_MINI_VISION",
     "AGENT_AZURE_41",
     "AGENT_AZURE_41_MINI",
+    "AGENT_AZURE_41_NANO",
     "AGENT_AZURE_41_VISION",
     "AGENT_AZURE_41_MINI_VISION",
+    "AGENT_AZURE_41_NANO_VISION",
     "AGENT_AZURE_5",
     "AGENT_AZURE_5_MINI",
     "AGENT_AZURE_5_NANO",

src/agentlab/agents/generic_agent/agent_configs.py

@@ -279,6 +279,10 @@
     chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-mini-2025-04-14"],
     flags=FLAGS_GPT_4o,
 )
+AGENT_AZURE_41_NANO = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-nano-2025-04-14"],
+    flags=FLAGS_GPT_4o,
+)
 
 AGENT_AZURE_5 = GenericAgentArgs(
     chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-2025-08-07"],
@@ -350,6 +354,10 @@
     chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-mini-2025-04-14"],
     flags=FLAGS_GPT_4o_VISION,
 )
+AGENT_AZURE_41_NANO_VISION = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-4.1-nano-2025-04-14"],
+    flags=FLAGS_GPT_4o_VISION,
+)
 
 AGENT_AZURE_5_VISION = GenericAgentArgs(
     chat_model_args=CHAT_MODEL_ARGS_DICT["azure/gpt-5-2025-08-07"],

src/agentlab/llm/chat_api.py

@@ -110,6 +110,10 @@ 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.
+    )
+
     def make_model(self):
         return AzureChatModel(
             model_name=self.model_name,
@@ -396,13 +400,25 @@ def __init__(
         model_name,
         api_key=None,
         temperature=0.5,
+        deployment_name=None,
         max_tokens=100,
         max_retry=4,
         min_retry_wait_time=60,
         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"
         endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
-        assert endpoint, "AZURE_OPENAI_ENDPOINT has to be defined in the environment"
+        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}."
+            )
 
         client_args = {
             "base_url": endpoint,
@@ -411,7 +427,6 @@ def __init__(
         super().__init__(
             model_name=model_name,
             api_key=api_key,
-            api_key_env_var="AZURE_OPENAI_API_KEY",
             temperature=temperature,
             max_tokens=max_tokens,
             max_retry=max_retry,

src/agentlab/llm/llm_configs.py

@@ -44,6 +44,13 @@
         temperature=1,  # gpt-5 supports temperature of 1 only
         vision_support=True,
     ),
+    "openai/gpt-4.1-nano-2025-04-14": OpenAIModelArgs(
+        model_name="gpt-4.1-nano-2025-04-14",
+        max_total_tokens=128_000,
+        max_input_tokens=128_000,
+        max_new_tokens=16_384,
+        vision_support=True,
+    ),
     "openai/gpt-4.1-mini-2025-04-14": OpenAIModelArgs(
         model_name="gpt-4.1-mini-2025-04-14",
         max_total_tokens=128_000,
@@ -160,6 +167,13 @@
         max_new_tokens=16_384,
         vision_support=True,
     ),
+    "azure/gpt-4.1-nano-2025-04-14": AzureModelArgs(
+        model_name="gpt-4.1-nano",
+        max_total_tokens=128_000,
+        max_input_tokens=128_000,
+        max_new_tokens=16_384,
+        vision_support=True,
+    ),
     "azure/gpt-5-2025-08-07": AzureModelArgs(
         model_name="gpt-5",
         max_total_tokens=400_000,

tests/llm/test_chat_api.py

@@ -21,10 +21,13 @@
 
 @pytest.mark.pricy
 @pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
+@pytest.mark.skipif(
+    not os.getenv("AZURE_OPENAI_API_KEY"), reason="Skipping as Azure API key not set"
+)
 def test_api_model_args_azure():
     model_args = AzureModelArgs(
-        model_name="gpt-35-turbo",
-        deployment_name="gpt-35-turbo",
+        model_name="gpt-4.1-nano",
+        deployment_name="gpt-4.1-nano",
         max_total_tokens=8192,
         max_input_tokens=8192 - 512,
         max_new_tokens=512,
@@ -43,6 +46,7 @@ def test_api_model_args_azure():
 
 @pytest.mark.pricy
 @pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
 def test_api_model_args_openai():
     model_args = OpenAIModelArgs(
         model_name="gpt-4o-mini",
@@ -64,6 +68,9 @@ def test_api_model_args_openai():
 
 @pytest.mark.pricy
 @pytest.mark.skipif(skip_tests, reason="Skipping on remote as Anthropic is pricy")
+@pytest.mark.skipif(
+    not os.getenv("ANTHROPIC_API_KEY"), reason="Skipping as Anthropic API key not set"
+)
 def test_api_model_args_anthropic():
     model_args = AnthropicModelArgs(
         model_name="claude-3-haiku-20240307",

tests/llm/test_litellm_api.py

@@ -1,7 +1,7 @@
+import os
 from functools import partial
 
 import pytest
-
 from agentlab.llm.litellm_api import LiteLLMModelArgs
 from agentlab.llm.response_api import APIPayload, LLMOutput
 
@@ -106,6 +106,7 @@ def test_multi_action_tool_calls():
 
 
 @pytest.mark.pricy
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
 def test_single_tool_call():
     """
     Test that the LLMOutput contains only one tool call when use_only_first_toolcall is True.
@@ -137,6 +138,7 @@ def test_single_tool_call():
 
 
 @pytest.mark.pricy
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
 def test_force_tool_call():
     """
     Test that the model can produce a specific tool call when requested.

tests/llm/test_response_api.py

@@ -686,6 +686,7 @@ def test_claude_model_with_multiple_messages_pricy_call():
 
 ## Test multiaction
 @pytest.mark.pricy
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
 def test_multi_action_tool_calls():
     """
     Test that the model can produce multiple tool calls in parallel.

tests/llm/test_tracking.py

@@ -146,7 +146,7 @@ def test_openai_chat_model():
     not AZURE_OPENAI_API_KEY_AVAILABLE, reason="Azure OpenAI API key is not available"
 )
 def test_azure_chat_model():
-    chat_model = AzureChatModel(model_name="gpt-35-turbo", deployment_name="gpt-35-turbo")
+    chat_model = AzureChatModel(model_name="gpt-4.1-nano", deployment_name="gpt-4.1-nano")
     assert chat_model.input_cost > 0
     assert chat_model.output_cost > 0