feat: sync with duplicated space

README.md

@@ -14,49 +14,71 @@ short_description: AI-driven TDAgent to automate threat analysis with MCP tools
 
 ---
 
-# Hackathon Participation: Cybersecurity AI Agents
+# TDAgentTools & TDAgent: Empowering Cybersecurity with Agentic AI
 
-This project is our contribution to Tracks 1 and 3 of the [Agents-MCP-Hackathon](https://huggingface.co/Agents-MCP-Hackathon), focused on applying AI technologies in the cybersecurity domain. Our aim is to develop solutions that improve the operational efficiency in cybersecurity through automation and data-driven insights.
+Welcome to TDAgentTools & TDAgent, our innovative proof of concept (PoC) crafted for the Agents-MCP Hackathon. Our initiatives focus on leveraging Agentic AI to enhance cybersecurity threat analysis, providing robust tools for data enrichment and strategic advice for incident handling.
 
-## Team Overview
+## Team Introduction
 
-Our team is part of the AI division in our company's cybersecurity department. We focus on implementing AI-based solutions to assist cybersecurity operations. Our team members include:
+We are an AI-focused team within a company, dedicated to empowering other teams by implementing AI solutions. Our expertise lies in automating processes to enhance productivity and tackle complex tasks that AI excels in. Our hackathon team members include:
 
-- **Pedro Completo Bento**
-- **Josep Pon Farreny**
-- **Sofia Jeronimo dos Santos**
-- **Rodrigo Dominguez Sanz**
-- **Miguel Rodin**
+- Pedro Completo Bento
+- Josep Pon Farreny
+- Sofia Jeronimo dos Santos
+- Rodrigo Dominguez Sanz
+- Miguel Rodin
 
-## Project Goals
+## Project Overview
 
-We are exploring the application of AI agents to aid cybersecurity analysts in threat data enrichment and threat analysis. Our main goals are:
+### Track 1: MCP Tool - TDAgentTools
 
-1. To experiment with agentic technologies like Gradio and MCP.
-2. To explore how AI can improve data enrichment capabilities in threat analysis.
-3. To develop autonomous agents capable of API interaction, data enrichment, and threat evaluation.
+TDAgentTools serves as an MCP server built using Gradio, offering a wide array of cybersecurity intelligence tools. These tools enable users to augment their LLMs' capabilities by integrating with various publicly available cybersecurity intel resources. Our TDAgentTools are accessible via the following link: [TDAgentTools Space](https://huggingface.co/spaces/Agents-MCP-Hackathon/TDAgentTools).
 
-## Track 1: MCP Tool / Server
+#### Available Tools:
+1. **TDAgentTools_get_url_http_content**: Retrieve URL content through an HTTP GET request.
+2. **TDAgentTools_query_abuseipdb**: Query AbuseIPDB to check if an IP is reported for abusive behavior.
+3. **TDAgentTools_query_rdap**: Gather information about internet resources such as domain names and IP addresses.
+4. **TDAgentTools_get_virus_total_url_info**: Fetch URL information using VirusTotal URL Scanner.
+5. **TDAgentTools_get_geolocation**: Obtain location details from an IP address.
+6. **TDAgentTools_enumerate_dns**: Access DNS configuration details for a given domain.
+7. **TDAgentTools_scrap_subdomains_for_domain**: Retrieve subdomains related to a domain.
+8. **TDAgentTools_retrieve_ioc_from_threatfox**: Get potential IoC information from ThreatFox.
+9. **TDAgentTools_get_stix_object_of_attack_id**: Access a STIX object using an ATT&CK ID.
+10. **TDAgentTools_lookup_user**: Seek user details from the Company User Lookup System.
+11. **TDAgentTools_lookup_cloud_account**: Investigate cloud account information.
+12. **TDAgentTools_send_email**: Simulate emailing from cert@company.com.
 
-In Track 1, we developed **TDAgentTools**, a Gradio-powered MCP server offering a set of public cybersecurity intelligence tools. This tool is designed to assist cybersecurity professionals in their threat analysis and response tasks.
+> **Note:** TDAgentTools rely on publicly provided APIs and some of which require API keys. If any of these API keys are revoked, certain tools may not function as intended.
 
-Access TDAgentTools here: [TDAgentTools Space](https://huggingface.co/spaces/Agents-MCP-Hackathon/TDAgentTools)
+### Track 3: Agentic Demo Showcase - TDAgent
 
-## Track 3: Agentic Demo Showcase
+TDAgent is an adaptive and interactive AI agent. This agent facilitates a dynamic AI experience, allowing users to switch the LLM used and adjust the system prompt to refine the agent’s behavior and objectives. It uses TDAgentTools to enrich threat data. Explore it here: [TDAgent Space](https://huggingface.co/spaces/Agents-MCP-Hackathon/TDAgent).
 
-For Track 3, we created **TDAgent**, an AI agent with a chat interface that connects to MCPs, defaulting to TDAgent MCP. The agent utilizes **TDAgentTools** or other MCP servers to gather additional threat intelligence, providing enriched data for more comprehensive threat evaluations.
+#### Key Features:
+- **Intelligent API Interactions**: The agent autonomously interacts with APIs for data enrichment and analysis without explicit user guidance.
+- **Enhanced Data Enrichment**: Automatically enriches initial incident data, providing deeper insights.
+- **Actionable Intelligence**: Suggests actions based on enriched data and analysis, displaying concise outputs for clearer communication.
+- **Versatile Adaptability**: Capable of switching LLMs for varied results and enhanced debugging.
 
-## Usage and Purpose
+## Motivation and Goals
 
-- **TDAgentTools**: Provides cybersecurity professionals with essential analysis tools via a user-friendly interface.
-- **TDAgent**: Facilitates interactive AI-supported threat analysis, enhancing efficiency, by leveraging data from MCP servers for improved insights.
+Our primary motivation is to explore Agentic AI applications in the cybersecurity realm, focusing on AI agent support for:
+1. Enriching reported threat data.
+2. Assisting analysts in threat analysis.
 
-Our work aims to reduce the manual effort involved in threat analysis, allowing cybersecurity teams to focus on strategic activities by utilizing AI for operational tasks.
+We aimed to:
+- Explore Agentic AI technologies like Gradio and MCP.
+- Enhance AI agent data enrichment with custom tools.
+- Enable agent autonomy in API interaction and threat assessment.
+- Equip the agent to propose specific incident response actions.
 
-## Conclusion
+## Insights & Conclusions
 
-This project seeks to demonstrate the practical applications of AI agents in cybersecurity, providing tools and frameworks to improve security operations.
+- **Agent's Autonomy**: Demonstrated autonomous API interactions and data enrichment capabilities.
+- **Enhanced Decision-Making**: The agent suggests data-driven insights beyond API outputs.
+- **Future Improvements**: Plan to fine-tune threat escalation logic and introduce additional decision layers for enhanced threat management.
 
+Our projects successfully demonstrated rapid prototyping with Gradio and Hugging Face Spaces, achieving all intended objectives while providing an engaging and rewarding experience for our team. This PoC shows the potential for future expansions and refinements in the realm of cybersecurity AI support!
 
 
 # TDA Agent
@@ -71,8 +93,3 @@ To start, sync all the dependencies with `uv sync --all-groups`.
 Then, install the pre-commit hooks (`uv run pre-commit install`) to
 ensure that future commits comply with the bare minimum to keep
 code _readable_.
-
-
-## Old content
-
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).

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()