feat: adds azure llm connection

tdagent/grchat.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from types import MappingProxyType
@@ -14,6 +15,7 @@ from langchain_aws import ChatBedrock
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
+from langchain_openai import AzureChatOpenAI
 from langgraph.prebuilt import create_react_agent
 from openai import OpenAI
 from openai.types.chat import ChatCompletion
@@ -77,6 +79,13 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
                 # ),
             },
         ),
+        (
+            "Azure OpenAI",
+            {
+                "GPT-3.5 Turbo": ("gpt-35-turbo"),
+                "GPT-4o": ("gpt-4o"),
+            },
+        ),
     ),
 )
 
@@ -128,12 +137,15 @@ def create_bedrock_llm(
 def create_hf_llm(
     hf_model_id: str,
     huggingfacehub_api_token: str | None = None,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
 ) -> tuple[ChatHuggingFace | None, str]:
     """Create a LangGraph Hugging Face agent."""
     try:
         llm = HuggingFaceEndpoint(
             model=hf_model_id,
-            temperature=0.8,
+            temperature=temperature,
+            max_new_tokens=max_tokens,
             task="text-generation",
             huggingfacehub_api_token=huggingfacehub_api_token,
         )
@@ -166,6 +178,30 @@ def create_openai_llm(
     return llm, ""
 
 
+def create_azure_llm(
+    model_id: str,
+    api_version: str,
+    endpoint: str,
+    token_id: str,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
+) -> tuple[AzureChatOpenAI | None, str]:
+    """Create a LangGraph Azure OpenAI agent."""
+    try:
+        os.environ["AZURE_OPENAI_ENDPOINT"] = endpoint
+        os.environ["AZURE_OPENAI_API_KEY"] = token_id
+        llm = AzureChatOpenAI(
+            azure_deployment=model_id,
+            api_key=token_id,
+            api_version=api_version,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+    except Exception as e:  # noqa: BLE001
+        return None, str(e)
+    return llm, ""
+
+
 #### UI functionality ####
 async def gr_connect_to_bedrock(  # noqa: PLR0913
     model_id: str,
@@ -230,11 +266,63 @@ async def gr_connect_to_hf(
     model_id: str,
     hf_access_token_textbox: str | None,
     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
 ) -> str:
     """Initialize Hugging Face agent."""
     global llm_agent  # noqa: PLW0603
 
-    llm, error = create_hf_llm(model_id, hf_access_token_textbox)
+    llm, error = create_hf_llm(
+        model_id,
+        hf_access_token_textbox,
+        temperature=temperature,
+        max_tokens=max_tokens,
+    )
+
+    if llm is None:
+        return f"❌ Connection failed: {error}"
+    tools = []
+    if mcp_servers:
+        client = MultiServerMCPClient(
+            {
+                server.name.replace(" ", "-"): {
+                    "url": server.value,
+                    "transport": "sse",
+                }
+                for server in mcp_servers
+            },
+        )
+        tools = await client.get_tools()
+
+    llm_agent = create_react_agent(
+        model=llm,
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
+    )
+
+    return "✅ Successfully connected to Hugging Face!"
+
+
+async def gr_connect_to_azure(
+    model_id: str,
+    azure_endpoint: str,
+    api_key: str,
+    api_version: str,
+    mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
+) -> str:
+    """Initialize Hugging Face agent."""
+    global llm_agent  # noqa: PLW0603
+
+    llm, error = create_azure_llm(
+        model_id,
+        api_version=api_version,
+        endpoint=azure_endpoint,
+        token_id=api_key,
+        temperature=temperature,
+        max_tokens=max_tokens,
+    )
 
     if llm is None:
         return f"❌ Connection failed: {error}"
@@ -334,6 +422,9 @@ def toggle_model_fields(
     dict[str, Any],
     dict[str, Any],
     dict[str, Any],
+    dict[str, Any],
+    dict[str, Any],
+    dict[str, Any],
 ]:  # ignore: F821
     """Toggle visibility of model fields based on the selected provider."""
     # Update model choices based on the selected provider
@@ -351,6 +442,8 @@ def toggle_model_fields(
     # Visibility settings for fields specific to each provider
     is_aws = provider == "AWS Bedrock"
     is_hf = provider == "HuggingFace"
+    is_azure = provider == "Azure OpenAI"
+    # is_nebius = provider == "Nebius"
     return (
         model_pretty,
         gr.update(visible=is_aws, interactive=is_aws),
@@ -358,6 +451,9 @@ def toggle_model_fields(
         gr.update(visible=is_aws, interactive=is_aws),
         gr.update(visible=is_aws, interactive=is_aws),
         gr.update(visible=is_hf, interactive=is_hf),
+        gr.update(visible=is_azure, interactive=is_azure),
+        gr.update(visible=is_azure, interactive=is_azure),
+        gr.update(visible=is_azure, interactive=is_azure),
     )
 
 
@@ -370,6 +466,9 @@ async def update_connection_status(  # noqa: PLR0913
     aws_session_token_textbox: str,
     aws_region_dropdown: str,
     hf_token: str,
+    azure_endpoint: str,
+    azure_api_token: str,
+    azure_api_version: str,
     temperature: float,
     max_tokens: int,
 ) -> str:
@@ -392,7 +491,23 @@ async def update_connection_status(  # noqa: PLR0913
                 max_tokens,
             )
         elif provider == "HuggingFace":
-            connection = await gr_connect_to_hf(model_id, hf_token, mcp_list_state)
+            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)
 
@@ -468,6 +583,25 @@ with (
                 placeholder="Enter your Hugging Face Access Token",
                 visible=False,
             )
+            azure_endpoint = gr.Textbox(
+                label="Azure OpenAI Endpoint",
+                type="text",
+                placeholder="Enter your Azure OpenAI Endpoint",
+                visible=False,
+            )
+            azure_api_token = gr.Textbox(
+                label="Azure Access Token",
+                type="password",
+                placeholder="Enter your Azure OpenAI Access Token",
+                visible=False,
+            )
+            azure_api_version = gr.Textbox(
+                label="Azure OpenAI API Version",
+                type="text",
+                placeholder="Enter your Azure OpenAI API Version",
+                value="2024-12-01-preview",
+                visible=False,
+            )
 
         with gr.Accordion("🧠  Model Configuration", open=True):
             model_display_id = gr.Dropdown(
@@ -485,6 +619,9 @@ with (
                     aws_session_token_textbox,
                     aws_region_dropdown,
                     hf_token,
+                    azure_endpoint,
+                    azure_api_token,
+                    azure_api_version,
                 ],
             )
             # Initialize the temperature and max tokens based on model specifications
@@ -517,6 +654,9 @@ with (
                 aws_session_token_textbox,
                 aws_region_dropdown,
                 hf_token,
+                azure_endpoint,
+                azure_api_token,
+                azure_api_version,
                 temperature,
                 max_tokens,
             ],