feat: adds langchain-openai

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,14 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
                 # ),
             },
         ),
+        (
+            "Azure OpenAI",
+            {
+                "GPT-4o": ("ggpt-4o-global-standard"),
+                "GPT-4o Mini": ("o4-mini"),
+                "GPT-4.5 Preview": ("gpt-4.5-preview"),
+            },
+        ),
     ),
 )
 
@@ -128,12 +138,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 +179,34 @@ 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
+        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,
+            **kwargs,
+        )
+    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 +271,18 @@ 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}"
@@ -260,6 +308,51 @@ async def gr_connect_to_hf(
     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}"
+    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 Azure OpenAI!"
+
+
 async def gr_connect_to_nebius(
     model_id: str,
     nebius_access_token_textbox: str,
@@ -334,6 +427,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 +447,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,43 +456,62 @@ 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),
     )
 
 
 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,
     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:
     """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)
-        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
 
@@ -468,13 +585,39 @@ 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(
-                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],
@@ -485,8 +628,21 @@ with (
                     aws_session_token_textbox,
                     aws_region_dropdown,
                     hf_token,
+                    azure_endpoint,
+                    azure_api_token,
+                    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",
@@ -510,13 +666,16 @@ with (
             update_connection_status,
             inputs=[
                 model_provider,
-                model_display_id,
+                model_id_textbox,
                 mcp_list.state,
                 aws_access_key_textbox,
                 aws_secret_key_textbox,
                 aws_session_token_textbox,
                 aws_region_dropdown,
                 hf_token,
+                azure_endpoint,
+                azure_api_token,
+                azure_api_version,
                 temperature,
                 max_tokens,
             ],