feat: git merge

tdagent/grchat.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any
@@ -8,9 +9,10 @@ import boto3
 import botocore
 import botocore.exceptions
 import gradio as gr
+import gradio.themes as gr_themes
 from langchain_aws import ChatBedrock
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph.prebuilt import create_react_agent
 from openai import OpenAI
@@ -51,15 +53,32 @@ GRADIO_ROLE_TO_LG_MESSAGE_TYPE = MappingProxyType(
     },
 )
 
-MODEL_OPTIONS = {
-    "AWS Bedrock": {
-        "Anthropic Claude 3.5 Sonnet": "eu.anthropic.claude-3-5-sonnet-20240620-v1:0",
-        #  "Anthropic Claude 3.7 Sonnet": "anthropic.claude-3-7-sonnet-20250219-v1:0",
-    },
-    "HuggingFace": {
-        "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct",
-    },
-}
+MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
+    (
+        (
+            "HuggingFace",
+            {
+                "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.3",
+                "Llama 3.1 8B Instruct": "meta-llama/Llama-3.1-8B-Instruct",
+                # "Qwen3 235B A22B": "Qwen/Qwen3-235B-A22B",  # Slow inference
+                "Microsoft Phi-3.5-mini Instruct": "microsoft/Phi-3.5-mini-instruct",
+                # "Deepseek R1 distill-llama 70B": "deepseek-ai/DeepSeek-R1-Distill-Llama-70B",  # noqa: E501
+                # "Deepseek V3": "deepseek-ai/DeepSeek-V3",
+            },
+        ),
+        (
+            "AWS Bedrock",
+            {
+                "Anthropic Claude 3.5 Sonnet (EU)": (
+                    "eu.anthropic.claude-3-5-sonnet-20240620-v1:0"
+                ),
+                #  "Anthropic Claude 3.7 Sonnet": (
+                #   "anthropic.claude-3-7-sonnet-20250219-v1:0"
+                # ),
+            },
+        ),
+    ),
+)
 
 #### Shared variables ####
 
@@ -109,18 +128,20 @@ def create_bedrock_llm(
 def create_hf_llm(
     hf_model_id: str,
     huggingfacehub_api_token: str | None = None,
-) -> tuple[HuggingFaceEndpoint | None, str]:
+) -> tuple[ChatHuggingFace | None, str]:
     """Create a LangGraph Hugging Face agent."""
     try:
         llm = HuggingFaceEndpoint(
             model=hf_model_id,
-            huggingfacehub_api_token=huggingfacehub_api_token,
             temperature=0.8,
+            task="text-generation",
+            huggingfacehub_api_token=huggingfacehub_api_token,
         )
+        chat_llm = ChatHuggingFace(llm=llm)
     except Exception as e:  # noqa: BLE001
         return None, str(e)
 
-    return llm, ""
+    return chat_llm, ""
 
 
 ## OpenAI LLM creation ##
@@ -286,14 +307,18 @@ async def gr_chat_function(  # noqa: D103
         messages.append(message_type(content=hist_msg["content"]))
 
     messages.append(HumanMessage(content=message))
-
-    llm_response = await llm_agent.ainvoke(
-        {
-            "messages": messages,
-        },
-    )
-
-    return llm_response["messages"][-1].content
+    try:
+        llm_response = await llm_agent.ainvoke(
+            {
+                "messages": messages,
+            },
+        )
+        return llm_response["messages"][-1].content
+    except Exception as err:
+        raise gr.Error(
+            f"We encountered an error while invoking the model:\n{err}",
+            print_exception=True,
+        ) from err
 
 
 ## UI components ##
@@ -314,7 +339,12 @@ def toggle_model_fields(
     # Update model choices based on the selected provider
     if provider in MODEL_OPTIONS:
         model_choices = list(MODEL_OPTIONS[provider].keys())
-        model_pretty = gr.update(choices=model_choices, visible=True, interactive=True)
+        model_pretty = gr.update(
+            choices=model_choices,
+            value=model_choices[0],
+            visible=True,
+            interactive=True,
+        )
     else:
         model_pretty = gr.update(choices=[], visible=False)
 
@@ -346,7 +376,9 @@ async def update_connection_status(  # noqa: PLR0913
     """Update the connection status based on the selected provider and model."""
     if not provider or not pretty_model:
         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(
@@ -363,15 +395,21 @@ async def update_connection_status(  # noqa: PLR0913
             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)
-        else:
-            return "❌ Invalid provider"
-    return connection if connection else "❌ Invalid provider"
+
+    return connection
 
 
-with gr.Blocks(
-    theme=gr.themes.Origin(primary_hue="teal", spacing_size="sm", font="sans-serif"),
-    title="TDAgent",
-) as gr_app, gr.Row():
+with (
+    gr.Blocks(
+        theme=gr_themes.Origin(
+            primary_hue="teal",
+            spacing_size="sm",
+            font="sans-serif",
+        ),
+        title="TDAgent",
+    ) as gr_app,
+    gr.Row(),
+):
     with gr.Column(scale=1):
         with gr.Accordion("🔌  MCP Servers", open=False):
             mcp_list = MutableCheckBoxGroup(
@@ -382,6 +420,10 @@ with gr.Blocks(
                     ),
                 ],
                 label="MCP Servers",
+                new_value_label="MCP endpoint",
+                new_name_label="MCP endpoint name",
+                new_value_placeholder="https://my-cool-mcp-server.com/mcp/sse",
+                new_name_placeholder="Swiss army knife of MCPs",
             )
 
         with gr.Accordion("⚙️  Provider Configuration", open=True):