adds-azure-provider

README.md

@@ -14,13 +14,13 @@ short_description: AI-driven TDAgent to automate threat analysis with MCP tools
 
 ---
 
-# Welcome to **TDAgentTools & TDAgent**
+# Hackathon Participation: Cybersecurity AI Agents
 
-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.
+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.
 
-## Team Introduction
+## Team Overview
 
-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:
+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:
 
 - **Pedro Completo Bento**
 - **Josep Pon Farreny**
@@ -28,66 +28,51 @@ We are an AI-focused team within a company, dedicated to empowering other teams
 - **Rodrigo Dominguez Sanz**
 - **Miguel Rodin**
 
-## Project Overview
+## Project Goals
 
-### Track 1: MCP Tool - **TDAgentTools**
+We are exploring the application of AI agents to aid cybersecurity analysts in threat data enrichment and threat analysis. Our main goals are:
 
-**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).
+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.
 
-#### 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.
+## Track 1: MCP Tool / Server
 
-> **Note:** TDAgentTools rely on publicly provided APIs, and some of these require API keys. If any of these API keys are revoked, certain tools may not function as intended.
+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.
 
-### Track 3: Agentic Demo Showcase - **TDAgent**
+Access TDAgentTools here: [TDAgentTools Space](https://huggingface.co/spaces/Agents-MCP-Hackathon/TDAgentTools)
 
-**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).
+## Track 3: Agentic Demo Showcase
 
-#### 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.
+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.
 
-## Motivation and Goals
+## Usage and Purpose
 
-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.
+- **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.
 
-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.
+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.
 
-## Insights & Conclusions
+## Conclusion
 
-- **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.
+This project seeks to demonstrate the practical applications of AI agents in cybersecurity, providing tools and frameworks to improve security operations.
 
-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
 
-## Development setup
+# Development setup
 
 To start developing you need the following tools:
 
-- [uv](https://docs.astral.sh/uv/)
+ * [uv](https://docs.astral.sh/uv/)
+
+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
 
-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_.
+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

@@ -22,7 +22,6 @@ dependencies = [
     "langchain-mcp-adapters>=0.1.1",
     "langchain-openai>=0.3.19",
     "langgraph>=0.4.7",
-    "markdown>=3.8",
     "openai>=1.84.0",
 ]
 
@@ -147,5 +146,4 @@ 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"]

requirements-dev.txt

@@ -59,7 +59,6 @@ langgraph-prebuilt==0.2.2
 langgraph-sdk==0.1.70
 langsmith==0.3.43
 license-expression==30.4.1
-markdown==3.8
 markdown-it-py==3.0.0
 markupsafe==3.0.2
 mcp==1.9.0

requirements.txt

@@ -51,7 +51,6 @@ langgraph-checkpoint==2.0.26
 langgraph-prebuilt==0.2.2
 langgraph-sdk==0.1.70
 langsmith==0.3.43
-markdown==3.8
 markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
 markupsafe==3.0.2
 mcp==1.9.0

tdagent/grchat.py

@@ -1,11 +1,8 @@
 from __future__ import annotations
 
-import dataclasses
-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
 
@@ -14,11 +11,8 @@ 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
-from langchain_core.tools import BaseTool
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langchain_openai import AzureChatOpenAI
@@ -35,46 +29,22 @@ if TYPE_CHECKING:
 
 #### Constants ####
 
-
-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,
-            """
+SYSTEM_MESSAGE = SystemMessage(
+    """
 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, 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.
+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.
 
+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(),
-        ),
-    ),
 )
 
 
@@ -85,7 +55,6 @@ GRADIO_ROLE_TO_LG_MESSAGE_TYPE = MappingProxyType(
     },
 )
 
-
 MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
     (
         (
@@ -113,60 +82,16 @@ MODEL_OPTIONS = OrderedDict(  # Initialize with tuples to preserve options order
         (
             "Azure OpenAI",
             {
-                "GPT-4o": ("ggpt-4o-global-standard"),
-                "GPT-4o Mini": ("o4-mini"),
-                "GPT-4.5 Preview": ("gpt-4.5-preview"),
+                "GPT-3.5 Turbo": ("gpt-35-turbo"),
+                "GPT-4o": ("gpt-4o"),
             },
         ),
     ),
 )
 
-CONNECT_STATE_DEFAULT = gr.State()
-
-
-@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 ####
 
@@ -232,8 +157,6 @@ def create_hf_llm(
 
 
 ## OpenAI LLM creation ##
-
-
 def create_openai_llm(
     model_id: str,
     token_id: str,
@@ -267,16 +190,12 @@ def create_azure_llm(
     try:
         os.environ["AZURE_OPENAI_ENDPOINT"] = endpoint
         os.environ["AZURE_OPENAI_API_KEY"] = token_id
-        if "o4-mini" in model_id:
-            kwargs = {"max_completion_tokens": max_tokens}
-        else:
-            kwargs = {"max_tokens": max_tokens}
         llm = AzureChatOpenAI(
             azure_deployment=model_id,
             api_key=token_id,
             api_version=api_version,
             temperature=temperature,
-            **kwargs,
+            max_tokens=max_tokens,
         )
     except Exception as e:  # noqa: BLE001
         return None, str(e)
@@ -284,56 +203,6 @@ 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,
@@ -341,15 +210,11 @@ 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
-    CONNECT_STATE_DEFAULT.value = True
-
     if not access_key or not secret_key:
         return "❌ Please provide both Access Key ID and Secret Access Key"
 
@@ -366,13 +231,32 @@ 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=await gr_fetch_mcp_tools(
-            mcp_servers,
-            trace_tools=trace_tool_calls,
-        ),
-        prompt=gr_make_system_message(agent_type=agent_type),
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
     )
 
     return "✅ Successfully connected to AWS Bedrock!"
@@ -382,14 +266,12 @@ 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:
     """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,
@@ -399,33 +281,39 @@ 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=await gr_fetch_mcp_tools(
-            mcp_servers,
-            trace_tools=trace_tool_calls,
-        ),
-        prompt=gr_make_system_message(agent_type=agent_type),
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
     )
 
     return "✅ Successfully connected to Hugging Face!"
 
 
-async def gr_connect_to_azure(  # noqa: PLR0913
+async def gr_connect_to_azure(
     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:
     """Initialize Hugging Face agent."""
     global llm_agent  # noqa: PLW0603
-    CONNECT_STATE_DEFAULT.value = True
 
     llm, error = create_azure_llm(
         model_id,
@@ -438,48 +326,59 @@ async def gr_connect_to_azure(  # noqa: PLR0913
 
     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=await gr_fetch_mcp_tools(mcp_servers, trace_tools=trace_tool_calls),
-        prompt=gr_make_system_message(agent_type=agent_type),
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
     )
 
-    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
-#     connected_state.value = True
-
-#     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!"
+    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
@@ -497,17 +396,12 @@ 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 _add_tools_trace_to_message(
-            llm_response["messages"][-1].content,
-        )
+        return llm_response["messages"][-1].content
     except Exception as err:
         raise gr.Error(
             f"We encountered an error while invoking the model:\n{err}",
@@ -515,50 +409,111 @@ async def gr_chat_function(  # noqa: D103
         ) from err
 
 
-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)
+## UI components ##
 
-    return f"{message}\n\n# Tools Trace\n\n" + "\n".join(traces)
 
+# 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),
+    )
 
-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
+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,
+    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 pretty_model:
+        return "❌ Please select a provider and model."
+
+    model_id = MODEL_OPTIONS.get(provider, {}).get(pretty_model)
+    connection = "❌ Invalid provider"
+    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,
+                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)
 
-    markdown_body = "".join(lines).strip()
-    return markdown.markdown(markdown_body)
+    return connection
 
 
-## UI components ##
-custom_css = """
-.main-header {
-    background: linear-gradient(135deg, #00a388 0%, #ffae00 100%);
-    padding: 30px;
-    border-radius: 5px;
-    margin-bottom: 20px;
-    text-align: center;
-}
-"""
 with (
     gr.Blocks(
         theme=gr_themes.Origin(
@@ -567,334 +522,156 @@ with (
             font="sans-serif",
         ),
         title="TDAgent",
-        fill_height=True,
-        fill_width=True,
-        css=custom_css,
     ) as gr_app,
+    gr.Row(),
 ):
-    gr.HTML(
-        """
-    <div class="main-header">
-        <h1>👩‍💻 TDAgentTools & TDAgent 👨‍💻</h1>
-        <p style="font-size: 1.2em; margin: 10px 0 0 0;">
-            Empowering Cybersecurity with Agentic AI
-        </p>
-    </div>
-    """,
-    )
-    with gr.Tabs():
-        with gr.TabItem("About"), gr.Row():
-            html_content = _read_markdown_body_as_html("README.md")
-            gr.Markdown(html_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",
-                            ),
-                        ],
-                        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",
-                            ],
-                            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 the model ID",
-                        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"
-                            ),
-                        )
-
-                    # 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,
-                    )
-                    max_tokens = gr.Slider(
-                        label="Max Tokens",
-                        minimum=128,
-                        maximum=8192,
-                        value=2048,
-                        step=64,
-                    )
-
-                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
-                    description="A simple threat analyst agent with MCP tools.",
-                )
-        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.
-            """,
-            )
-            gr.HTML(
-                """<iframe width="560" height="315" src="https://youtu.be/C6Z9EOW-3lE?feature=shared" frameborder="0" allowfullscreen></iframe>""",  # noqa: E501
+    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",
             )
 
-    ## 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,
+        with gr.Accordion("⚙️  Provider Configuration", open=True):
+            model_provider = gr.Dropdown(
+                choices=list(MODEL_OPTIONS.keys()),
+                value=None,
+                label="Select Model Provider",
             )
-
-        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)
-
-    # Initialize a flag to check if connected
-
-    def _on_change_model_configuration(*args: str) -> Any:  # noqa: ARG001
-        # 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,
+            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,
             )
-        return gr.update()
-
-    ## 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],
-    )
+        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,
+                    azure_endpoint,
+                    azure_api_token,
+                    azure_api_version,
+                ],
+            )
+            # 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_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],
-    )
+        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_display_id,
+                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],
+        )
 
-    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],
-    )
-    model_provider.change(
-        _on_change_model_configuration,
-        inputs=[model_provider],
-    )
-    model_id_dropdown.change(
-        _on_change_model_configuration,
-        inputs=[model_id_dropdown, model_provider],
-    )
+    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.",
+        )
 
-## Entry Point ##
 
 if __name__ == "__main__":
     gr_app.launch()

uv.lock

@@ -1148,15 +1148,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/53/84/8a89614b2e7eeeaf0a68a4046d6cfaea4544c8619ea02595ebeec9b2bae3/license_expression-30.4.1-py3-none-any.whl", hash = "sha256:679646bc3261a17690494a3e1cada446e5ee342dbd87dcfa4a0c24cc5dce13ee", size = 111457, upload-time = "2025-01-14T05:11:38.658Z" },
 ]
 
-[[package]]
-name = "markdown"
-version = "3.8"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/2f/15/222b423b0b88689c266d9eac4e61396fe2cc53464459d6a37618ac863b24/markdown-3.8.tar.gz", hash = "sha256:7df81e63f0df5c4b24b7d156eb81e4690595239b7d70937d0409f1b0de319c6f", size = 360906, upload-time = "2025-04-11T14:42:50.928Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/51/3f/afe76f8e2246ffbc867440cbcf90525264df0e658f8a5ca1f872b3f6192a/markdown-3.8-py3-none-any.whl", hash = "sha256:794a929b79c5af141ef5ab0f2f642d0f7b1872981250230e72682346f7cc90dc", size = 106210, upload-time = "2025-04-11T14:42:49.178Z" },
-]
-
 [[package]]
 name = "markdown-it-py"
 version = "3.0.0"
@@ -2878,7 +2869,6 @@ dependencies = [
     { name = "langchain-mcp-adapters" },
     { name = "langchain-openai" },
     { name = "langgraph" },
-    { name = "markdown" },
     { name = "openai" },
 ]
 
@@ -2907,7 +2897,6 @@ requires-dist = [
     { name = "langchain-mcp-adapters", specifier = ">=0.1.1" },
     { name = "langchain-openai", specifier = ">=0.3.19" },
     { name = "langgraph", specifier = ">=0.4.7" },
-    { name = "markdown", specifier = ">=3.8" },
     { name = "openai", specifier = ">=1.84.0" },
 ]