Tools tracing

pyproject.toml

@@ -146,4 +146,5 @@ convention = "google"
 [tool.ruff.lint.per-file-ignores]
 "*/__init__.py" = ["F401"]
 "tdagent/cli/**/*.py" = ["D103", "T201"]
+"tdagent/grchat.py" = ["ANN401", "FBT001"]
 "tests/*.py" = ["D103", "PLR2004", "S101"]

tdagent/grchat.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import dataclasses
+import enum
 import os
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
@@ -12,7 +14,9 @@ import botocore.exceptions
 import gradio as gr
 import gradio.themes as gr_themes
 from langchain_aws import ChatBedrock
+from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_core.tools import BaseTool
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langchain_openai import AzureChatOpenAI
@@ -29,22 +33,46 @@ if TYPE_CHECKING:
 
 #### Constants ####
 
-SYSTEM_MESSAGE = SystemMessage(
-    """
+
+class AgentType(str, enum.Enum):
+    """TDAgent type."""
+
+    INCIDENT_HANDLER = "Incident handler"
+    DATA_ENRICHER = "Data enricher"
+
+    def __str__(self) -> str:  # noqa: D105
+        return self.value
+
+
+AGENT_SYSTEM_MESSAGES = OrderedDict(
+    (
+        (
+            AgentType.INCIDENT_HANDLER,
+            """
 You are a security analyst assistant responsible for collecting, analyzing
 and disseminating actionable intelligence related to cyber threats,
 vulnerabilities and threat actors.
 
 When presented with potential incidents information or tickets, you should
-evaluate the presented evidence, decide what is missing and gather
-additional data using any tool at your disposal. After gathering more
-information you must evaluate if the incident is a threat or
-not and, if possible, remediation actions.
+evaluate the presented evidence, gather additional data using any tool at
+your disposal and take corrective actions if possible.
+
+Afterwards, generate a cybersecurity report including: key findings, challenges,
+actions taken and recommendations.
 
-You must always present the conducted analysis and final conclusion.
 Never use external means of communication, like emails or SMS, unless
 instructed to do so.
 """.strip(),
+        ),
+        (
+            AgentType.DATA_ENRICHER,
+            """
+You are a cybersecurity incidence data enriching assistant. Analysts
+will present information about security incidents and you must use
+all the tools at your disposal to enrich the data as much as possible.
+""".strip(),
+        ),
+    ),
 )
 
 
@@ -55,6 +83,7 @@ GRADIO_ROLE_TO_LG_MESSAGE_TYPE = MappingProxyType(
     },
 )
 
+
 MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
     (
         (
@@ -90,9 +119,50 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
     ),
 )
 
+
+@dataclasses.dataclass
+class ToolInvocationInfo:
+    """Information related to a tool invocation by the LLM."""
+
+    name: str
+    inputs: Mapping[str, Any]
+
+
+class ToolsTracerCallback(BaseCallbackHandler):
+    """Callback that registers tools invoked by the Agent."""
+
+    def __init__(self) -> None:
+        self._tools_trace: list[ToolInvocationInfo] = []
+
+    def on_tool_start(  # noqa: D102
+        self,
+        serialized: dict[str, Any],
+        *args: Any,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        self._tools_trace.append(
+            ToolInvocationInfo(
+                name=serialized.get("name", "<unknown-function-name>"),
+                inputs=inputs if inputs else {},
+            ),
+        )
+        return super().on_tool_start(serialized, *args, inputs=inputs, **kwargs)
+
+    @property
+    def tools_trace(self) -> Sequence[ToolInvocationInfo]:
+        """Tools trace information."""
+        return self._tools_trace
+
+    def clear(self) -> None:
+        """Clear tools trace."""
+        self._tools_trace.clear()
+
+
 #### Shared variables ####
 
 llm_agent: CompiledGraph | None = None
+llm_tools_tracer: ToolsTracerCallback | None = None
 
 #### Utility functions ####
 
@@ -158,6 +228,8 @@ def create_hf_llm(
 
 
 ## OpenAI LLM creation ##
+
+
 def create_openai_llm(
     model_id: str,
     token_id: str,
@@ -208,6 +280,56 @@ def create_azure_llm(
 
 
 #### UI functionality ####
+
+
+async def gr_fetch_mcp_tools(
+    mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    *,
+    trace_tools: bool,
+) -> list[BaseTool]:
+    """Fetch tools from MCP servers."""
+    global llm_tools_tracer  # noqa: PLW0603
+
+    if mcp_servers:
+        client = MultiServerMCPClient(
+            {
+                server.name.replace(" ", "-"): {
+                    "url": server.value,
+                    "transport": "sse",
+                }
+                for server in mcp_servers
+            },
+        )
+        tools = await client.get_tools()
+        if trace_tools:
+            llm_tools_tracer = ToolsTracerCallback()
+            for tool in tools:
+                if tool.callbacks is None:
+                    tool.callbacks = [llm_tools_tracer]
+                elif isinstance(tool.callbacks, list):
+                    tool.callbacks.append(llm_tools_tracer)
+                else:
+                    tool.callbacks.add_handler(llm_tools_tracer)
+        else:
+            llm_tools_tracer = None
+
+        return tools
+
+    return []
+
+
+def gr_make_system_message(
+    agent_type: AgentType,
+) -> SystemMessage:
+    """Make agent's system message."""
+    try:
+        system_msg = AGENT_SYSTEM_MESSAGES[agent_type]
+    except KeyError as err:
+        raise gr.Error(f"Unknown agent type '{agent_type}'") from err
+
+    return SystemMessage(system_msg)
+
+
 async def gr_connect_to_bedrock(  # noqa: PLR0913
     model_id: str,
     access_key: str,
@@ -215,11 +337,14 @@ async def gr_connect_to_bedrock(  # noqa: PLR0913
     session_token: str,
     region: str,
     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    agent_type: AgentType,
+    trace_tool_calls: bool,
     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"
 
@@ -236,32 +361,13 @@ async def gr_connect_to_bedrock(  # noqa: PLR0913
     if llm is None:
         return f"❌ Connection failed: {error}"
 
-    # client = MultiServerMCPClient(
-    #     {
-    #         "toolkit": {
-    #             "url": "https://agents-mcp-hackathon-tdagenttools.hf.space/gradio_api/mcp/sse",
-    #             "transport": "sse",
-    #         },
-    #     }
-    # )
-    # tools = await client.get_tools()
-    if mcp_servers:
-        client = MultiServerMCPClient(
-            {
-                server.name.replace(" ", "-"): {
-                    "url": server.value,
-                    "transport": "sse",
-                }
-                for server in mcp_servers
-            },
-        )
-        tools = await client.get_tools()
-    else:
-        tools = []
     llm_agent = create_react_agent(
         model=llm,
-        tools=tools,
-        prompt=SYSTEM_MESSAGE,
+        tools=await gr_fetch_mcp_tools(
+            mcp_servers,
+            trace_tools=trace_tool_calls,
+        ),
+        prompt=gr_make_system_message(agent_type=agent_type),
     )
 
     return "✅ Successfully connected to AWS Bedrock!"
@@ -271,6 +377,8 @@ async def gr_connect_to_hf(
     model_id: str,
     hf_access_token_textbox: str | None,
     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    agent_type: AgentType,
+    trace_tool_calls: bool,
     temperature: float = 0.8,
     max_tokens: int = 512,
 ) -> str:
@@ -286,34 +394,27 @@ async def gr_connect_to_hf(
 
     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,
+        tools=await gr_fetch_mcp_tools(
+            mcp_servers,
+            trace_tools=trace_tool_calls,
+        ),
+        prompt=gr_make_system_message(agent_type=agent_type),
     )
 
     return "✅ Successfully connected to Hugging Face!"
 
 
-async def gr_connect_to_azure(
+async def gr_connect_to_azure(  # noqa: PLR0913
     model_id: str,
     azure_endpoint: str,
     api_key: str,
     api_version: str,
     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+    agent_type: AgentType,
+    trace_tool_calls: bool,
     temperature: float = 0.8,
     max_tokens: int = 512,
 ) -> str:
@@ -331,59 +432,47 @@ async def gr_connect_to_azure(
 
     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,
+        tools=await gr_fetch_mcp_tools(mcp_servers, trace_tools=trace_tool_calls),
+        prompt=gr_make_system_message(agent_type=agent_type),
     )
 
     return "✅ Successfully connected to Azure OpenAI!"
 
 
-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_connect_to_nebius(
+#     model_id: str,
+#     nebius_access_token_textbox: str,
+#     mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+# ) -> str:
+#     """Initialize Hugging Face agent."""
+#     global llm_agent
+
+#     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
@@ -401,12 +490,17 @@ async def gr_chat_function(  # noqa: D103
 
     messages.append(HumanMessage(content=message))
     try:
+        if llm_tools_tracer is not None:
+            llm_tools_tracer.clear()
+
         llm_response = await llm_agent.ainvoke(
             {
                 "messages": messages,
             },
         )
-        return llm_response["messages"][-1].content
+        return _add_tools_trace_to_message(
+            llm_response["messages"][-1].content,
+        )
     except Exception as err:
         raise gr.Error(
             f"We encountered an error while invoking the model:\n{err}",
@@ -414,106 +508,26 @@ async def gr_chat_function(  # noqa: D103
         ) from err
 
 
-## UI components ##
+def _add_tools_trace_to_message(message: str) -> str:
+    if not llm_tools_tracer or not llm_tools_tracer.tools_trace:
+        return message
+    import json
 
+    traces = []
+    for index, tool_info in enumerate(llm_tools_tracer.tools_trace):
+        trace_msg = f"  {index}. {tool_info.name}"
+        if tool_info.inputs:
+            trace_msg += "\n"
+            trace_msg += "    * Arguments:\n"
+            trace_msg += "      ```json\n"
+            trace_msg += f"      {json.dumps(tool_info.inputs, indent=4)}\n"
+            trace_msg += "      ```\n"
+        traces.append(trace_msg)
 
-# 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],
-    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,
-            value=model_choices[0],
-            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"
-    is_azure = provider == "Azure OpenAI"
-    # is_nebius = provider == "Nebius"
-    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),
-        gr.update(visible=is_azure, interactive=is_azure),
-        gr.update(visible=is_azure, interactive=is_azure),
-        gr.update(visible=is_azure, interactive=is_azure),
-    )
-
+    return f"{message}\n\n# Tools Trace\n\n" + "\n".join(traces)
 
-async def update_connection_status(  # noqa: PLR0913
-    provider: 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 model_id:
-        return "❌ Please select a provider and model."
-    connection = "❌ Invalid provider"
-    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
+## UI components ##
 
 
 with (
@@ -549,65 +563,66 @@ with (
                 value=None,
                 label="Select Model Provider",
             )
-            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,
-            )
-            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",
-                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,
-            )
-            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,
-            )
+
+            ## 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_display_id = gr.Dropdown(
-                label="Select Model from the list",
+            model_id_dropdown = gr.Dropdown(
+                label="Select known model id or type your own below",
                 choices=[],
                 visible=False,
             )
@@ -618,31 +633,24 @@ with (
                 visible=False,
                 interactive=True,
             )
-            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,
-                    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,
+
+            # 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."
+                    ),
                 )
-                if x
-                else model_id_textbox.value,
-                inputs=[model_display_id, model_provider],
-                outputs=[model_id_textbox],
-            )
+                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 specifications
             temperature = gr.Slider(
                 label="Temperature",
@@ -653,44 +661,157 @@ with (
             )
             max_tokens = gr.Slider(
                 label="Max Tokens",
-                minimum=64,
-                maximum=4096,
-                value=512,
+                minimum=128,
+                maximum=8192,
+                value=2048,
                 step=64,
             )
 
-        connect_btn = gr.Button("🔌  Connect to Model", variant="primary")
-        status_textbox = gr.Textbox(label="Connection Status", interactive=False)
-
-        connect_btn.click(
-            update_connection_status,
-            inputs=[
-                model_provider,
-                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,
-            ],
-            outputs=[status_textbox],
+        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,
+        )
+
+        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="This is a simple agent that uses MCP tools.",
+            title="👩‍💻 TDAgent 👨‍💻",
+            description="A simple threat analyst agent with MCP tools.",
         )
 
+    ## UI Events ##
+
+    def _toggle_model_choices_ui(
+        provider: str,
+    ) -> dict[str, Any]:
+        if provider in MODEL_OPTIONS:
+            model_choices = list(MODEL_OPTIONS[provider].keys())
+            return gr.update(
+                choices=model_choices,
+                value=model_choices[0],
+                visible=True,
+                interactive=True,
+            )
+
+        return gr.update(choices=[], visible=False)
+
+    def _toggle_model_aws_bedrock_conf_ui(
+        provider: str,
+    ) -> tuple[dict[str, Any], ...]:
+        is_aws = provider == "AWS Bedrock"
+        return gr.update(visible=is_aws), gr.update(visible=is_aws)
+
+    def _toggle_model_hf_conf_ui(
+        provider: str,
+    ) -> tuple[dict[str, Any], ...]:
+        is_hf = provider == "HuggingFace"
+        return gr.update(visible=is_hf), gr.update(visible=is_hf)
+
+    def _toggle_model_azure_conf_ui(
+        provider: str,
+    ) -> tuple[dict[str, Any], ...]:
+        is_azure = provider == "Azure OpenAI"
+        return gr.update(visible=is_azure), gr.update(visible=is_azure)
+
+    ## Connect Event Listeners ##
+
+    model_provider.change(
+        _toggle_model_choices_ui,
+        inputs=[model_provider],
+        outputs=[model_id_dropdown],
+    )
+    model_provider.change(
+        _toggle_model_aws_bedrock_conf_ui,
+        inputs=[model_provider],
+        outputs=[aws_bedrock_conf_group, connect_aws_bedrock_btn],
+    )
+    model_provider.change(
+        _toggle_model_hf_conf_ui,
+        inputs=[model_provider],
+        outputs=[hf_conf_group, connect_hf_btn],
+    )
+    model_provider.change(
+        _toggle_model_azure_conf_ui,
+        inputs=[model_provider],
+        outputs=[azure_conf_group, connect_azure_btn],
+    )
+
+    connect_aws_bedrock_btn.click(
+        gr_connect_to_bedrock,
+        inputs=[
+            model_id_textbox,
+            aws_access_key_textbox,
+            aws_secret_key_textbox,
+            aws_session_token_textbox,
+            aws_region_dropdown,
+            mcp_list.state,
+            agent_system_message_radio,
+            agent_trace_tools_checkbox,
+            temperature,
+            max_tokens,
+        ],
+        outputs=[status_textbox],
+    )
+
+    connect_hf_btn.click(
+        gr_connect_to_hf,
+        inputs=[
+            model_id_textbox,
+            hf_token,
+            mcp_list.state,
+            agent_system_message_radio,
+            agent_trace_tools_checkbox,
+            temperature,
+            max_tokens,
+        ],
+        outputs=[status_textbox],
+    )
+
+    connect_azure_btn.click(
+        gr_connect_to_azure,
+        inputs=[
+            model_id_textbox,
+            azure_endpoint,
+            azure_api_token,
+            azure_api_version,
+            mcp_list.state,
+            agent_system_message_radio,
+            agent_trace_tools_checkbox,
+            temperature,
+            max_tokens,
+        ],
+        outputs=[status_textbox],
+    )
+
+    model_id_dropdown.change(
+        lambda x, y: (
+            gr.update(
+                value=MODEL_OPTIONS.get(y, {}).get(x),
+                visible=True,
+            )
+            if x
+            else model_id_textbox.value
+        ),
+        inputs=[model_id_dropdown, model_provider],
+        outputs=[model_id_textbox],
+    )
+
+## Entry Point ##
 
 if __name__ == "__main__":
     gr_app.launch()