feat: adds tabs and warning when selecting diff model config after being connected to the model

tdagent/grchat.py

@@ -5,6 +5,7 @@ import enum
 import os
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
+from pathlib import Path
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any
 
@@ -13,6 +14,7 @@ import botocore
 import botocore.exceptions
 import gradio as gr
 import gradio.themes as gr_themes
+import markdown
 from langchain_aws import ChatBedrock
 from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
@@ -119,6 +121,8 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
     ),
 )
 
+CONNECT_STATE_DEFAULT = gr.State()
+
 
 @dataclasses.dataclass
 class ToolInvocationInfo:
@@ -344,6 +348,7 @@ async def gr_connect_to_bedrock(  # noqa: PLR0913
 ) -> str:
     """Initialize Bedrock agent."""
     global llm_agent  # noqa: PLW0603
+    CONNECT_STATE_DEFAULT.value = True
 
     if not access_key or not secret_key:
         return "❌ Please provide both Access Key ID and Secret Access Key"
@@ -384,7 +389,7 @@ async def gr_connect_to_hf(
 ) -> str:
     """Initialize Hugging Face agent."""
     global llm_agent  # noqa: PLW0603
-
+    CONNECT_STATE_DEFAULT.value = True
     llm, error = create_hf_llm(
         model_id,
         hf_access_token_textbox,
@@ -420,6 +425,7 @@ async def gr_connect_to_azure(  # noqa: PLR0913
 ) -> str:
     """Initialize Hugging Face agent."""
     global llm_agent  # noqa: PLW0603
+    CONNECT_STATE_DEFAULT.value = True
 
     llm, error = create_azure_llm(
         model_id,
@@ -449,6 +455,7 @@ async def gr_connect_to_azure(  # noqa: PLR0913
 # ) -> str:
 #     """Initialize Hugging Face agent."""
 #     global llm_agent
+#     connected_state.value = True
 
 #     llm, error = create_openai_llm(model_id, nebius_access_token_textbox)
 
@@ -527,6 +534,21 @@ def _add_tools_trace_to_message(message: str) -> str:
     return f"{message}\n\n# Tools Trace\n\n" + "\n".join(traces)
 
 
+def _read_markdown_body_as_html(path: str = "README.md") -> str:
+    with Path(path).open(encoding="utf-8") as f:  # Default mode is "r"
+        lines = f.readlines()
+
+    # Skip YAML front matter if present
+    if lines and lines[0].strip() == "---":
+        for i in range(1, len(lines)):
+            if lines[i].strip() == "---":
+                lines = lines[i + 1 :]  # skip metadata block
+                break
+
+    markdown_body = "".join(lines).strip()
+    return markdown.markdown(markdown_body)
+
+
 ## UI components ##
 
 
@@ -538,160 +560,178 @@ with (
             font="sans-serif",
         ),
         title="TDAgent",
+        fill_height=True,
+        fill_width=True,
     ) as gr_app,
-    gr.Row(),
+    gr.Tabs(),
 ):
-    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",
-                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):
-            model_provider = gr.Dropdown(
-                choices=list(MODEL_OPTIONS.keys()),
-                value=None,
-                label="Select Model Provider",
-            )
-
-            ## Amazon Bedrock Configuration ##
-            with gr.Group(visible=False) as aws_bedrock_conf_group:
-                aws_access_key_textbox = gr.Textbox(
-                    label="AWS Access Key ID",
-                    type="password",
-                    placeholder="Enter your AWS Access Key ID",
-                )
-                aws_secret_key_textbox = gr.Textbox(
-                    label="AWS Secret Access Key",
-                    type="password",
-                    placeholder="Enter your AWS Secret Access Key",
-                )
-                aws_region_dropdown = gr.Dropdown(
-                    label="AWS Region",
-                    choices=[
-                        "us-east-1",
-                        "us-west-2",
-                        "eu-west-1",
-                        "eu-central-1",
-                        "ap-southeast-1",
+    with gr.TabItem("About"):
+        html_content = _read_markdown_body_as_html("README.md")
+        gr.Markdown(html_content, elem_id="about-content")
+
+    with gr.TabItem("TDAgent"), 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",
+                        ),
                     ],
-                    value="eu-west-1",
-                )
-                aws_session_token_textbox = gr.Textbox(
-                    label="AWS Session Token",
-                    type="password",
-                    placeholder="Enter your AWS session token",
+                    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",
                 )
 
-            ## Huggingface Configuration ##
-            with gr.Group(visible=False) as hf_conf_group:
-                hf_token = gr.Textbox(
-                    label="HuggingFace Token",
-                    type="password",
-                    placeholder="Enter your Hugging Face Access Token",
+            with gr.Accordion("⚙️  Provider Configuration", open=True):
+                model_provider = gr.Dropdown(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    value=None,
+                    label="Select Model Provider",
                 )
 
-            ## Azure Configuration ##
-            with gr.Group(visible=False) as azure_conf_group:
-                azure_endpoint = gr.Textbox(
-                    label="Azure OpenAI Endpoint",
+                ## Amazon Bedrock Configuration ##
+                with gr.Group(visible=False) as aws_bedrock_conf_group:
+                    aws_access_key_textbox = gr.Textbox(
+                        label="AWS Access Key ID",
+                        type="password",
+                        placeholder="Enter your AWS Access Key ID",
+                    )
+                    aws_secret_key_textbox = gr.Textbox(
+                        label="AWS Secret Access Key",
+                        type="password",
+                        placeholder="Enter your AWS Secret Access Key",
+                    )
+                    aws_region_dropdown = gr.Dropdown(
+                        label="AWS Region",
+                        choices=[
+                            "us-east-1",
+                            "us-west-2",
+                            "eu-west-1",
+                            "eu-central-1",
+                            "ap-southeast-1",
+                        ],
+                        value="eu-west-1",
+                    )
+                    aws_session_token_textbox = gr.Textbox(
+                        label="AWS Session Token",
+                        type="password",
+                        placeholder="Enter your AWS session token",
+                    )
+
+                ## Huggingface Configuration ##
+                with gr.Group(visible=False) as hf_conf_group:
+                    hf_token = gr.Textbox(
+                        label="HuggingFace Token",
+                        type="password",
+                        placeholder="Enter your Hugging Face Access Token",
+                    )
+
+                ## Azure Configuration ##
+                with gr.Group(visible=False) as azure_conf_group:
+                    azure_endpoint = gr.Textbox(
+                        label="Azure OpenAI Endpoint",
+                        type="text",
+                        placeholder="Enter your Azure OpenAI Endpoint",
+                    )
+                    azure_api_token = gr.Textbox(
+                        label="Azure Access Token",
+                        type="password",
+                        placeholder="Enter your Azure OpenAI Access Token",
+                    )
+                    azure_api_version = gr.Textbox(
+                        label="Azure OpenAI API Version",
+                        type="text",
+                        placeholder="Enter your Azure OpenAI API Version",
+                        value="2024-12-01-preview",
+                    )
+
+            with gr.Accordion("🧠  Model Configuration", open=True):
+                model_id_dropdown = gr.Dropdown(
+                    label="Select known model id or type your own below",
+                    choices=[],
+                    visible=False,
+                )
+                model_id_textbox = gr.Textbox(
+                    label="Model ID",
                     type="text",
-                    placeholder="Enter your Azure OpenAI Endpoint",
+                    placeholder="Enter the model ID",
+                    visible=False,
+                    interactive=True,
                 )
-                azure_api_token = gr.Textbox(
-                    label="Azure Access Token",
-                    type="password",
-                    placeholder="Enter your Azure OpenAI Access Token",
+
+                # Agent configuration options
+                with gr.Group():
+                    agent_system_message_radio = gr.Radio(
+                        choices=list(AGENT_SYSTEM_MESSAGES.keys()),
+                        value=next(iter(AGENT_SYSTEM_MESSAGES.keys())),
+                        label="Agent type",
+                        info=(
+                            "Changes the system message to pre-condition the agent"
+                            " to act in a desired way."
+                        ),
+                    )
+                    agent_trace_tools_checkbox = gr.Checkbox(
+                        value=False,
+                        label="Trace tool calls",
+                        info="Add the invoked tools trace at the end of the message",
+                    )
+
+                # Initialize the temperature and max tokens based on model specs
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.8,
+                    step=0.1,
                 )
-                azure_api_version = gr.Textbox(
-                    label="Azure OpenAI API Version",
-                    type="text",
-                    placeholder="Enter your Azure OpenAI API Version",
-                    value="2024-12-01-preview",
+                max_tokens = gr.Slider(
+                    label="Max Tokens",
+                    minimum=128,
+                    maximum=8192,
+                    value=2048,
+                    step=64,
                 )
 
-        with gr.Accordion("🧠  Model Configuration", open=True):
-            model_id_dropdown = gr.Dropdown(
-                label="Select known model id or type your own below",
-                choices=[],
+            connect_aws_bedrock_btn = gr.Button(
+                "🔌  Connect to Bedrock",
+                variant="primary",
                 visible=False,
             )
-            model_id_textbox = gr.Textbox(
-                label="Model ID",
-                type="text",
-                placeholder="Enter the model ID",
+            connect_hf_btn = gr.Button(
+                "🔌  Connect to Huggingface 🤗",
+                variant="primary",
+                visible=False,
+            )
+            connect_azure_btn = gr.Button(
+                "🔌  Connect to Azure",
+                variant="primary",
                 visible=False,
-                interactive=True,
             )
 
-            # Agent configuration options
-            with gr.Group():
-                agent_system_message_radio = gr.Radio(
-                    choices=list(AGENT_SYSTEM_MESSAGES.keys()),
-                    value=next(iter(AGENT_SYSTEM_MESSAGES.keys())),
-                    label="Agent type",
-                    info=(
-                        "Changes the system message to pre-condition the agent"
-                        " to act in a desired way."
-                    ),
-                )
-                agent_trace_tools_checkbox = gr.Checkbox(
-                    value=False,
-                    label="Trace tool calls",
-                    info="Add the invoked tools trace at the end of the message",
-                )
+            status_textbox = gr.Textbox(label="Connection Status", interactive=False)
 
-            # 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=128,
-                maximum=8192,
-                value=2048,
-                step=64,
+        with gr.Column(scale=2):
+            chat_interface = gr.ChatInterface(
+                fn=gr_chat_function,
+                type="messages",
+                examples=[],  # Add examples if needed
+                title="👩‍💻 TDAgent 👨‍💻",
+                description="A simple threat analyst agent with MCP tools.",
             )
-
-        connect_aws_bedrock_btn = gr.Button(
-            "🔌  Connect to Bedrock",
-            variant="primary",
-            visible=False,
-        )
-        connect_hf_btn = gr.Button(
-            "🔌  Connect to Huggingface 🤗",
-            variant="primary",
-            visible=False,
-        )
-        connect_azure_btn = gr.Button(
-            "🔌  Connect to Azure",
-            variant="primary",
-            visible=False,
+    with gr.TabItem("Demo"):
+        gr.Markdown(
+            """
+        This is a demo of TDAgent, a simple threat analyst agent with MCP tools.
+        You can configure the agent to use different LLM providers and connect to
+        various MCP servers to access tools.
+        """,
         )
-
-        status_textbox = gr.Textbox(label="Connection Status", interactive=False)
-
-    with gr.Column(scale=2):
-        chat_interface = gr.ChatInterface(
-            fn=gr_chat_function,
-            type="messages",
-            examples=[],  # Add examples if needed
-            title="��‍💻 TDAgent 👨‍💻",
-            description="A simple threat analyst agent with MCP tools.",
+        gr.HTML(
+            """<iframe width="560" height="315" src="https://youtu.be/C6Z9EOW-3lE?feature=shared" frameborder="0" allowfullscreen></iframe>""",  # noqa: E501
         )
 
     ## UI Events ##
@@ -728,6 +768,19 @@ with (
         is_azure = provider == "Azure OpenAI"
         return gr.update(visible=is_azure), gr.update(visible=is_azure)
 
+    # Initialize a flag to check if connected
+
+    def _on_change_model_configuration() -> (
+        Any
+    ):  # If model configuration changes after connecting, issue a warning
+        if CONNECT_STATE_DEFAULT.value:
+            CONNECT_STATE_DEFAULT.value = False  # Reset the state
+            return gr.Warning(
+                "When changing model configuration, you need to reconnect.",
+                duration=5,
+            )
+        return gr.update()
+
     ## Connect Event Listeners ##
 
     model_provider.change(
@@ -810,6 +863,14 @@ with (
         inputs=[model_id_dropdown, model_provider],
         outputs=[model_id_textbox],
     )
+    model_provider.change(
+        _on_change_model_configuration,
+        inputs=[model_provider],
+    )
+    model_id_dropdown.change(
+        _on_change_model_configuration,
+        inputs=[model_id_dropdown, model_provider],
+    )
 
 ## Entry Point ##