ui-improvements (#11)

- feat: updates dependencies (54bdb7fbbdc4bf60daeb470f6ed2796ba387f8e5)
- feat: removes annotations (e6cfbb2fd6fe8c0429994c1bead8eb0ad8674987)
- feat: updates dependencies (f20f6154b6d31972d9b64b8f180317ab32c426ed)
- feat: improves ui (87d04296fe3766c5759390c1a70689505d02b30e)

.pre-commit-config.yaml

@@ -65,6 +65,7 @@ repos:
             "--format",
             "requirements-txt",
             "--no-hashes",
+            "--no-annotate",
             "--no-dev",
             "-o",
             "requirements.txt",
@@ -82,6 +83,7 @@ repos:
             "--format",
             "requirements-txt",
             "--no-hashes",
+            "--no-annotate",
             "--group",
             "dev",
             "--group",
@@ -91,7 +93,7 @@ repos:
           ]
       - id: mypy
         name: Running mypy
-        stages: [commit]
+        stages: [pre-commit]
         language: system
         entry: uv run mypy
         args: [--install-types, --non-interactive]

pyproject.toml

@@ -16,8 +16,10 @@ dependencies = [
     "gradio[mcp]~=5.31",
     "huggingface-hub>=0.32.3",
     "langchain-aws>=0.2.24",
+    "langchain-huggingface>=0.2.0",
     "langchain-mcp-adapters>=0.1.1",
     "langgraph>=0.4.7",
+    "openai>=1.84.0",
 ]
 
 [project.scripts]

requirements-dev.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv export --format requirements-txt --no-hashes --group dev --group test -o requirements-dev.txt
+#    uv export --format requirements-txt --no-hashes --no-annotate --group dev --group test -o requirements-dev.txt
 aiofiles==24.1.0
 annotated-types==0.7.0
 anyio==4.9.0
@@ -18,6 +18,7 @@ coverage==7.8.2
 cyclonedx-python-lib==9.1.0
 defusedxml==0.7.1
 distlib==0.3.9
+distro==1.9.0
 exceptiongroup==1.3.0 ; python_full_version < '3.11'
 fastapi==0.115.12
 ffmpy==0.6.0
@@ -36,11 +37,14 @@ identify==2.6.12
 idna==3.10
 iniconfig==2.1.0
 jinja2==3.1.6
+jiter==0.10.0
 jmespath==1.0.1
+joblib==1.5.1
 jsonpatch==1.33
 jsonpointer==3.0.0
 langchain-aws==0.2.24
 langchain-core==0.3.63
+langchain-huggingface==0.2.0
 langchain-mcp-adapters==0.1.1
 langgraph==0.4.7
 langgraph-checkpoint==2.0.26
@@ -52,12 +56,30 @@ markdown-it-py==3.0.0
 markupsafe==3.0.2
 mcp==1.9.0
 mdurl==0.1.2
+mpmath==1.3.0
 msgpack==1.1.0
 mypy==1.16.0
 mypy-extensions==1.1.0
+networkx==3.4.2 ; python_full_version < '3.11'
+networkx==3.5 ; python_full_version >= '3.11'
 nodeenv==1.9.1
 numpy==1.26.4 ; python_full_version < '3.12'
 numpy==2.2.6 ; python_full_version >= '3.12'
+nvidia-cublas-cu12==12.6.4.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-cupti-cu12==12.6.80 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-nvrtc-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-runtime-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cudnn-cu12==9.5.1.17 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cufft-cu12==11.3.0.4 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cufile-cu12==1.11.1.6 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-curand-cu12==10.3.7.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusolver-cu12==11.7.1.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusparse-cu12==12.5.4.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusparselt-cu12==0.6.3 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nccl-cu12==2.26.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nvjitlink-cu12==12.6.85 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nvtx-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+openai==1.84.0
 orjson==3.10.18
 ormsgpack==1.10.0
 packageurl-python==0.16.0
@@ -88,24 +110,36 @@ python-dotenv==1.1.0
 python-multipart==0.0.20
 pytz==2025.2
 pyyaml==6.0.2
+regex==2024.11.6
 requests==2.32.3
 requests-toolbelt==1.0.0
 rich==14.0.0
 ruff==0.11.12
 s3transfer==0.13.0
 safehttpx==0.1.6
+safetensors==0.5.3
+scikit-learn==1.6.1
+scipy==1.15.3
 semantic-version==2.10.0
+sentence-transformers==4.1.0
+setuptools==80.9.0 ; (python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
 shellingham==1.5.4 ; sys_platform != 'emscripten'
 six==1.17.0
 sniffio==1.3.1
 sortedcontainers==2.4.0
 sse-starlette==2.3.6
 starlette==0.46.2
+sympy==1.14.0
 tenacity==9.1.2
+threadpoolctl==3.6.0
+tokenizers==0.21.1
 toml==0.10.2
 tomli==2.2.1 ; python_full_version <= '3.11'
 tomlkit==0.13.2
+torch==2.7.1
 tqdm==4.67.1
+transformers==4.52.4
+triton==3.3.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
 typer==0.16.0 ; sys_platform != 'emscripten'
 typing-extensions==4.13.2
 typing-inspection==0.4.1

requirements.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv export --format requirements-txt --no-hashes --no-dev -o requirements.txt
+#    uv export --format requirements-txt --no-hashes --no-annotate --no-dev -o requirements.txt
 aiofiles==24.1.0
 annotated-types==0.7.0
 anyio==4.9.0
@@ -12,6 +12,7 @@ charset-normalizer==3.4.2
 click==8.2.1 ; sys_platform != 'emscripten'
 colorama==0.4.6 ; sys_platform == 'win32'
 coverage==7.8.2
+distro==1.9.0
 exceptiongroup==1.3.0 ; python_full_version < '3.11'
 fastapi==0.115.12
 ffmpy==0.6.0
@@ -29,11 +30,14 @@ huggingface-hub==0.32.3
 idna==3.10
 iniconfig==2.1.0
 jinja2==3.1.6
+jiter==0.10.0
 jmespath==1.0.1
+joblib==1.5.1
 jsonpatch==1.33
 jsonpointer==3.0.0
 langchain-aws==0.2.24
 langchain-core==0.3.63
+langchain-huggingface==0.2.0
 langchain-mcp-adapters==0.1.1
 langgraph==0.4.7
 langgraph-checkpoint==2.0.26
@@ -44,8 +48,26 @@ markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
 markupsafe==3.0.2
 mcp==1.9.0
 mdurl==0.1.2 ; sys_platform != 'emscripten'
+mpmath==1.3.0
+networkx==3.4.2 ; python_full_version < '3.11'
+networkx==3.5 ; python_full_version >= '3.11'
 numpy==1.26.4 ; python_full_version < '3.12'
 numpy==2.2.6 ; python_full_version >= '3.12'
+nvidia-cublas-cu12==12.6.4.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-cupti-cu12==12.6.80 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-nvrtc-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cuda-runtime-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cudnn-cu12==9.5.1.17 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cufft-cu12==11.3.0.4 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cufile-cu12==1.11.1.6 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-curand-cu12==10.3.7.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusolver-cu12==11.7.1.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusparse-cu12==12.5.4.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-cusparselt-cu12==0.6.3 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nccl-cu12==2.26.2 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nvjitlink-cu12==12.6.85 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+nvidia-nvtx-cu12==12.6.77 ; platform_machine == 'x86_64' and sys_platform == 'linux'
+openai==1.84.0
 orjson==3.10.18
 ormsgpack==1.10.0
 packaging==24.2
@@ -66,22 +88,34 @@ python-dotenv==1.1.0
 python-multipart==0.0.20
 pytz==2025.2
 pyyaml==6.0.2
+regex==2024.11.6
 requests==2.32.3
 requests-toolbelt==1.0.0
 rich==14.0.0 ; sys_platform != 'emscripten'
 ruff==0.11.12 ; sys_platform != 'emscripten'
 s3transfer==0.13.0
 safehttpx==0.1.6
+safetensors==0.5.3
+scikit-learn==1.6.1
+scipy==1.15.3
 semantic-version==2.10.0
+sentence-transformers==4.1.0
+setuptools==80.9.0 ; (python_full_version >= '3.12' and platform_machine != 'x86_64') or (python_full_version >= '3.12' and sys_platform != 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')
 shellingham==1.5.4 ; sys_platform != 'emscripten'
 six==1.17.0
 sniffio==1.3.1
 sse-starlette==2.3.6
 starlette==0.46.2
+sympy==1.14.0
 tenacity==9.1.2
+threadpoolctl==3.6.0
+tokenizers==0.21.1
 tomli==2.2.1 ; python_full_version <= '3.11'
 tomlkit==0.13.2
+torch==2.7.1
 tqdm==4.67.1
+transformers==4.52.4
+triton==3.3.1 ; platform_machine == 'x86_64' and sys_platform == 'linux'
 typer==0.16.0 ; sys_platform != 'emscripten'
 typing-extensions==4.13.2
 typing-inspection==0.4.1

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
@@ -10,8 +10,11 @@ import botocore.exceptions
 import gradio as gr
 from langchain_aws import ChatBedrock
 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
 
@@ -48,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 ####
 
@@ -56,12 +68,15 @@ llm_agent: CompiledGraph | None = None
 #### Utility functions ####
 
 
+## Bedrock LLM creation ##
 def create_bedrock_llm(
     bedrock_model_id: str,
     aws_access_key: str,
     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 = {
@@ -70,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)
@@ -83,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)
@@ -91,20 +105,59 @@ def create_bedrock_llm(
     return llm, ""
 
 
-#### UI functionality ####
+## Hugging Face LLM creation ##
+def create_hf_llm(
+    hf_model_id: str,
+    huggingfacehub_api_token: str | None = None,
+) -> tuple[HuggingFaceEndpoint | None, str]:
+    """Create a LangGraph Hugging Face agent."""
+    try:
+        llm = HuggingFaceEndpoint(
+            model=hf_model_id,
+            huggingfacehub_api_token=huggingfacehub_api_token,
+            temperature=0.8,
+        )
+    except Exception as e:  # noqa: BLE001
+        return None, str(e)
+
+    return llm, ""
 
 
-async def gr_connect_to_bedrock(
+## 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(  # 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"
 
@@ -114,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:
@@ -128,7 +183,6 @@ async def gr_connect_to_bedrock(
     #     }
     # )
     # tools = await client.get_tools()
-
     if mcp_servers:
         client = MultiServerMCPClient(
             {
@@ -142,7 +196,6 @@ async def gr_connect_to_bedrock(
         tools = await client.get_tools()
     else:
         tools = []
-
     llm_agent = create_react_agent(
         model=llm,
         tools=tools,
@@ -152,6 +205,73 @@ async def gr_connect_to_bedrock(
     return "✅ Successfully connected to AWS Bedrock!"
 
 
+async def gr_connect_to_hf(
+    model_id: str,
+    hf_access_token_textbox: str | None,
+    mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+) -> str:
+    """Initialize Hugging Face agent."""
+    global llm_agent  # noqa: PLW0603
+
+    llm, error = create_hf_llm(model_id, hf_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=llm,
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
+    )
+
+    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]],
@@ -178,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",
+) as gr_app, gr.Row():
+    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=[
@@ -231,31 +412,83 @@ 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],
         )
 
-    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.",
-    )
+    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.",
+        )
 
 
 if __name__ == "__main__":