feat: adds azure models

tdagent/grchat.py

@@ -82,8 +82,9 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
         (
             "Azure OpenAI",
             {
-                "GPT-3.5 Turbo": ("gpt-35-turbo"),
-                "GPT-4o": ("gpt-4o"),
+                "GPT-4o": ("ggpt-4o-global-standard"),
+                "GPT-4o Mini": ("o4-mini"),
+                "GPT-4.5 Preview": ("gpt-4.5-preview"),
             },
         ),
     ),
@@ -190,12 +191,16 @@ def create_azure_llm(
     try:
         os.environ["AZURE_OPENAI_ENDPOINT"] = endpoint
         os.environ["AZURE_OPENAI_API_KEY"] = token_id
+        if "o4-mini" in model_id:
+            kwargs = {"max_completion_tokens": max_tokens}
+        else:
+            kwargs = {"max_tokens": max_tokens}
         llm = AzureChatOpenAI(
             azure_deployment=model_id,
             api_key=token_id,
             api_version=api_version,
             temperature=temperature,
-            max_tokens=max_tokens,
+            **kwargs,
         )
     except Exception as e:  # noqa: BLE001
         return None, str(e)
@@ -345,7 +350,7 @@ async def gr_connect_to_azure(
         prompt=SYSTEM_MESSAGE,
     )
 
-    return "✅ Successfully connected to Hugging Face!"
+    return "✅ Successfully connected to Azure OpenAI!"
 
 
 async def gr_connect_to_nebius(
@@ -459,7 +464,7 @@ def toggle_model_fields(
 
 async def update_connection_status(  # noqa: PLR0913
     provider: str,
-    pretty_model: str,
+    model_id: str,
     mcp_list_state: Sequence[MutableCheckBoxGroupEntry] | None,
     aws_access_key_textbox: str,
     aws_secret_key_textbox: str,
@@ -473,43 +478,40 @@ async def update_connection_status(  # noqa: PLR0913
     max_tokens: int,
 ) -> str:
     """Update the connection status based on the selected provider and model."""
-    if not provider or not pretty_model:
+    if not provider or not model_id:
         return "❌ Please select a provider and model."
-
-    model_id = MODEL_OPTIONS.get(provider, {}).get(pretty_model)
     connection = "❌ Invalid provider"
-    if model_id:
-        if provider == "AWS Bedrock":
-            connection = await gr_connect_to_bedrock(
-                model_id,
-                aws_access_key_textbox,
-                aws_secret_key_textbox,
-                aws_session_token_textbox,
-                aws_region_dropdown,
-                mcp_list_state,
-                temperature,
-                max_tokens,
-            )
-        elif provider == "HuggingFace":
-            connection = await gr_connect_to_hf(
-                model_id,
-                hf_token,
-                mcp_list_state,
-                temperature,
-                max_tokens,
-            )
-        elif provider == "Azure OpenAI":
-            connection = await gr_connect_to_azure(
-                model_id,
-                azure_endpoint,
-                azure_api_token,
-                azure_api_version,
-                mcp_list_state,
-                temperature,
-                max_tokens,
-            )
-        elif provider == "Nebius":
-            connection = await gr_connect_to_nebius(model_id, hf_token, mcp_list_state)
+    if provider == "AWS Bedrock":
+        connection = await gr_connect_to_bedrock(
+            model_id,
+            aws_access_key_textbox,
+            aws_secret_key_textbox,
+            aws_session_token_textbox,
+            aws_region_dropdown,
+            mcp_list_state,
+            temperature,
+            max_tokens,
+        )
+    elif provider == "HuggingFace":
+        connection = await gr_connect_to_hf(
+            model_id,
+            hf_token,
+            mcp_list_state,
+            temperature,
+            max_tokens,
+        )
+    elif provider == "Azure OpenAI":
+        connection = await gr_connect_to_azure(
+            model_id,
+            azure_endpoint,
+            azure_api_token,
+            azure_api_version,
+            mcp_list_state,
+            temperature,
+            max_tokens,
+        )
+    elif provider == "Nebius":
+        connection = await gr_connect_to_nebius(model_id, hf_token, mcp_list_state)
 
     return connection
 
@@ -605,10 +607,17 @@ with (
 
         with gr.Accordion("🧠  Model Configuration", open=True):
             model_display_id = gr.Dropdown(
-                label="Select Model ID",
+                label="Select Model from the list",
                 choices=[],
                 visible=False,
             )
+            model_id_textbox = gr.Textbox(
+                label="Model ID",
+                type="text",
+                placeholder="Enter the model ID",
+                visible=False,
+                interactive=True,
+            )
             model_provider.change(
                 toggle_model_fields,
                 inputs=[model_provider],
@@ -624,6 +633,16 @@ with (
                     azure_api_version,
                 ],
             )
+            model_display_id.change(
+                lambda x, y: gr.update(
+                    value=MODEL_OPTIONS.get(y, {}).get(x),
+                    visible=True,
+                )
+                if x
+                else model_id_textbox.value,
+                inputs=[model_display_id, model_provider],
+                outputs=[model_id_textbox],
+            )
             # Initialize the temperature and max tokens based on model specifications
             temperature = gr.Slider(
                 label="Temperature",
@@ -647,7 +666,7 @@ with (
             update_connection_status,
             inputs=[
                 model_provider,
-                model_display_id,
+                model_id_textbox,
                 mcp_list.state,
                 aws_access_key_textbox,
                 aws_secret_key_textbox,