feat: improves UI

tdagent/grchat.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import boto3
 import botocore
@@ -13,6 +13,8 @@ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph.prebuilt import create_react_agent
+from openai import OpenAI
+from openai.types.chat import ChatCompletion
 
 from tdagent.grcomponents import MutableCheckBoxGroup, MutableCheckBoxGroupEntry
 
@@ -49,6 +51,15 @@ 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",
+    },
+}
 
 #### Shared variables ####
 
@@ -64,6 +75,8 @@ def create_bedrock_llm(
     aws_secret_key: str,
     aws_session_token: str,
     aws_region: str,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
 ) -> tuple[ChatBedrock | None, str]:
     """Create a LangGraph Bedrock agent."""
     boto3_config = {
@@ -72,7 +85,6 @@ def create_bedrock_llm(
         "aws_session_token": aws_session_token if aws_session_token else None,
         "region_name": aws_region,
     }
-
     # Verify credentials
     try:
         sts = boto3.client("sts", **boto3_config)
@@ -85,7 +97,7 @@ def create_bedrock_llm(
         llm = ChatBedrock(
             model_id=bedrock_model_id,
             client=bedrock_client,
-            model_kwargs={"temperature": 0.8},
+            model_kwargs={"temperature": temperature, "max_tokens": max_tokens},
         )
     except Exception as e:  # noqa: BLE001
         return None, str(e)
@@ -111,18 +123,41 @@ def create_hf_llm(
     return llm, ""
 
 
+## OpenAI LLM creation ##
+def create_openai_llm(
+    model_id: str,
+    token_id: str,
+) -> tuple[ChatCompletion | None, str]:
+    """Create a LangGraph OpenAI agent."""
+    try:
+        client = OpenAI(
+            base_url="https://api.studio.nebius.com/v1/",
+            api_key=token_id,
+        )
+        llm = client.chat.completions.create(
+            messages=[],  # needs to be fixed
+            model=model_id,
+            max_tokens=512,
+            temperature=0.8,
+        )
+    except Exception as e:  # noqa: BLE001
+        return None, str(e)
+    return llm, ""
+
+
 #### UI functionality ####
-async def gr_connect_to_bedrock(
+async def gr_connect_to_bedrock(  # noqa: PLR0913
     model_id: str,
     access_key: str,
     secret_key: str,
     session_token: str,
     region: str,
     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    temperature: float = 0.8,
+    max_tokens: int = 512,
 ) -> str:
     """Initialize Bedrock agent."""
     global llm_agent  # noqa: PLW0603
-
     if not access_key or not secret_key:
         return "❌ Please provide both Access Key ID and Secret Access Key"
 
@@ -132,6 +167,8 @@ async def gr_connect_to_bedrock(
         secret_key.strip(),
         session_token.strip(),
         region,
+        temperature=temperature,
+        max_tokens=max_tokens,
     )
 
     if llm is None:
@@ -146,7 +183,6 @@ async def gr_connect_to_bedrock(
     #     }
     # )
     # tools = await client.get_tools()
-    tools = []
     if mcp_servers:
         client = MultiServerMCPClient(
             {
@@ -158,7 +194,8 @@ async def gr_connect_to_bedrock(
             },
         )
         tools = await client.get_tools()
-
+    else:
+        tools = []
     llm_agent = create_react_agent(
         model=llm,
         tools=tools,
@@ -202,6 +239,39 @@ async def gr_connect_to_hf(
     return "✅ Successfully connected to Hugging Face!"
 
 
+async def gr_connect_to_nebius(
+    model_id: str,
+    nebius_access_token_textbox: str,
+    mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+) -> str:
+    """Initialize Hugging Face agent."""
+    global llm_agent  # noqa: PLW0603
+
+    llm, error = create_openai_llm(model_id, nebius_access_token_textbox)
+
+    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=str(llm),
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
+    )
+    return "✅ Successfully connected to nebius!"
+
+
 async def gr_chat_function(  # noqa: D103
     message: str,
     history: list[Mapping[str, str]],
@@ -228,49 +298,110 @@ async def gr_chat_function(  # noqa: D103
 
 ## UI components ##
 
-with gr.Blocks() as gr_app:
-    gr.Markdown("# 🔐 Secure Bedrock Chatbot")
-
-    ### MCP Servers ###
-    with gr.Accordion():
-        mcp_list = MutableCheckBoxGroup(
-            values=[
-                MutableCheckBoxGroupEntry(
-                    name="TDAgent tools",
-                    value="https://agents-mcp-hackathon-tdagenttools.hf.space/gradio_api/mcp/sse",
-                ),
-            ],
-            label="MCP Servers",
-        )
 
-    # Credentials section (collapsible)
-    with gr.Accordion("🔑 Bedrock Configuration", open=True):
-        gr.Markdown(
-            "**Note**: Credentials are only stored in memory during your session.",
-        )
-        with gr.Row():
-            bedrock_model_id_textbox = gr.Textbox(
-                label="Bedrock Model Id",
-                value="eu.anthropic.claude-3-5-sonnet-20240620-v1:0",
+# Function to toggle visibility and set model IDs
+def toggle_model_fields(
+    provider: str,
+) -> tuple[
+    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
+    if provider in MODEL_OPTIONS:
+        model_choices = list(MODEL_OPTIONS[provider].keys())
+        model_pretty = gr.update(choices=model_choices, visible=True, interactive=True)
+    else:
+        model_pretty = gr.update(choices=[], visible=False)
+
+    # Visibility settings for fields specific to each provider
+    is_aws = provider == "AWS Bedrock"
+    is_hf = provider == "HuggingFace"
+    return (
+        model_pretty,
+        gr.update(visible=is_aws, interactive=is_aws),
+        gr.update(visible=is_aws, interactive=is_aws),
+        gr.update(visible=is_aws, interactive=is_aws),
+        gr.update(visible=is_aws, interactive=is_aws),
+        gr.update(visible=is_hf, interactive=is_hf),
+    )
+
+
+async def update_connection_status(  # noqa: PLR0913
+    provider: str,
+    pretty_model: 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,
+    temperature: float,
+    max_tokens: int,
+) -> str:
+    """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)
+    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)
+        else:
+            return "❌ Invalid provider"
+    return connection if connection else "❌ Invalid provider"
+
+
+with gr.Blocks(
+    theme=gr.themes.Origin(primary_hue="teal", spacing_size="sm", font="sans-serif"),
+    title="TDAgent",
+) and gr.Row() as gr_app:
+    with gr.Column(scale=1):
+        with gr.Accordion("🔌  MCP Servers", open=False):
+            mcp_list = MutableCheckBoxGroup(
+                values=[
+                    MutableCheckBoxGroupEntry(
+                        name="TDAgent tools",
+                        value="https://agents-mcp-hackathon-tdagenttools.hf.space/gradio_api/mcp/sse",
+                    ),
+                ],
+                label="MCP Servers",
+            )
+
+        with gr.Accordion("⚙️  Provider Configuration", open=True):
+            model_provider = gr.Dropdown(
+                choices=list(MODEL_OPTIONS.keys()),
+                value=None,
+                label="Select Model Provider",
             )
-        with gr.Row():
             aws_access_key_textbox = gr.Textbox(
                 label="AWS Access Key ID",
                 type="password",
                 placeholder="Enter your AWS Access Key ID",
+                visible=False,
             )
             aws_secret_key_textbox = gr.Textbox(
                 label="AWS Secret Access Key",
                 type="password",
                 placeholder="Enter your AWS Secret Access Key",
+                visible=False,
             )
-        with gr.Row():
-            aws_session_token_textbox = gr.Textbox(
-                label="AWS Session Token",
-                type="password",
-                placeholder="Enter your AWS session token",
-            )
-        with gr.Row():
             aws_region_dropdown = gr.Dropdown(
                 label="AWS Region",
                 choices=[
@@ -281,58 +412,84 @@ with gr.Blocks() as gr_app:
                     "ap-southeast-1",
                 ],
                 value="eu-west-1",
+                visible=False,
+            )
+            aws_session_token_textbox = gr.Textbox(
+                label="AWS Session Token",
+                type="password",
+                placeholder="Enter your AWS session token",
+                visible=False,
+            )
+            hf_token = gr.Textbox(
+                label="HuggingFace Token",
+                type="password",
+                placeholder="Enter your Hugging Face Access Token",
+                visible=False,
+            )
+
+        with gr.Accordion("🧠  Model Configuration", open=True):
+            model_display_id = gr.Dropdown(
+                label="Select Model ID",
+                choices=[],
+                visible=False,
+            )
+            model_provider.change(
+                toggle_model_fields,
+                inputs=[model_provider],
+                outputs=[
+                    model_display_id,
+                    aws_access_key_textbox,
+                    aws_secret_key_textbox,
+                    aws_session_token_textbox,
+                    aws_region_dropdown,
+                    hf_token,
+                ],
+            )
+            # Initialize the temperature and max tokens based on model specifications
+            temperature = gr.Slider(
+                label="Temperature",
+                minimum=0.0,
+                maximum=1.0,
+                value=0.8,
+                step=0.1,
+            )
+            max_tokens = gr.Slider(
+                label="Max Tokens",
+                minimum=64,
+                maximum=4096,
+                value=512,
+                step=64,
             )
-            connect_btn = gr.Button("🔌 Connect to Bedrock", variant="primary")
 
+        connect_btn = gr.Button("🔌  Connect to Model", variant="primary")
         status_textbox = gr.Textbox(label="Connection Status", interactive=False)
 
         connect_btn.click(
-            gr_connect_to_bedrock,
+            update_connection_status,
             inputs=[
-                bedrock_model_id_textbox,
+                model_provider,
+                model_display_id,
+                mcp_list.state,
                 aws_access_key_textbox,
                 aws_secret_key_textbox,
                 aws_session_token_textbox,
                 aws_region_dropdown,
-                mcp_list.state,
+                hf_token,
+                temperature,
+                max_tokens,
             ],
             outputs=[status_textbox],
         )
 
-    with gr.Accordion("Hugging Face Configuration", open=True):
-        with gr.Row():
-            hf_model_id_textbox = gr.Textbox(
-                label="HF Model Id",
-                value="fdtn-ai/Foundation-Sec-8B",
-            )
-        with gr.Row():
-            hf_access_token_textbox = gr.Textbox(
-                label="Hugging Face Access Token",
-                type="password",
-                placeholder="Enter your Hugging Face Access Token",
-            )
-            hf_connect_btn = gr.Button("🔌 Connect to Hugging Face", variant="primary")
-
-        status_textbox = gr.Textbox(label="Connection Status", interactive=False)
-
-        hf_connect_btn.click(
-            gr_connect_to_hf,
-            inputs=[
-                hf_model_id_textbox,
-                hf_access_token_textbox,
-                mcp_list.state,
-            ],
-            outputs=[status_textbox],
+    with gr.Column(scale=2):
+        chat_interface = gr.ChatInterface(
+            fn=gr_chat_function,
+            type="messages",
+            examples=[],  # Add examples if needed
+            title="👩‍💻 TDAgent",
+            description="This is a simple agent that uses MCP tools.",
         )
 
-    chat_interface = gr.ChatInterface(
-        fn=gr_chat_function,
-        type="messages",
-        examples=[],
-        title="Agent with MCP Tools",
-        description="This is a simple agent that uses MCP tools.",
-    )
-
 
 if __name__ == "__main__":
     gr_app.launch()