pedrobento988 commited on
Commit
25f5cbc
·
verified ·
1 Parent(s): 1bb1418

ui-improvements (#17)

Browse files

- feat: improves ui (1cf2398913610e9b0967eb0764fe71e54969d765)

Files changed (6) hide show
  1. README.md +28 -27
  2. pyproject.toml +1 -0
  3. requirements-dev.txt +1 -0
  4. requirements.txt +1 -0
  5. tdagent/grchat.py +229 -146
  6. uv.lock +11 -0
README.md CHANGED
@@ -14,48 +14,48 @@ short_description: AI-driven TDAgent to automate threat analysis with MCP tools
14
 
15
  ---
16
 
17
- # TDAgentTools & TDAgent: Empowering Cybersecurity with Agentic AI
18
 
19
- 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.
20
 
21
  ## Team Introduction
22
 
23
  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:
24
 
25
- - Pedro Completo Bento
26
- - Josep Pon Farreny
27
- - Sofia Jeronimo dos Santos
28
- - Rodrigo Dominguez Sanz
29
- - Miguel Rodin
30
 
31
  ## Project Overview
32
 
33
- ### Track 1: MCP Tool - TDAgentTools
34
 
35
- 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).
36
 
37
  #### Available Tools:
38
- 1. **TDAgentTools_get_url_http_content**: Retrieve URL content through an HTTP GET request.
39
- 2. **TDAgentTools_query_abuseipdb**: Query AbuseIPDB to check if an IP is reported for abusive behavior.
40
- 3. **TDAgentTools_query_rdap**: Gather information about internet resources such as domain names and IP addresses.
41
- 4. **TDAgentTools_get_virus_total_url_info**: Fetch URL information using VirusTotal URL Scanner.
42
- 5. **TDAgentTools_get_geolocation**: Obtain location details from an IP address.
43
- 6. **TDAgentTools_enumerate_dns**: Access DNS configuration details for a given domain.
44
- 7. **TDAgentTools_scrap_subdomains_for_domain**: Retrieve subdomains related to a domain.
45
- 8. **TDAgentTools_retrieve_ioc_from_threatfox**: Get potential IoC information from ThreatFox.
46
- 9. **TDAgentTools_get_stix_object_of_attack_id**: Access a STIX object using an ATT&CK ID.
47
- 10. **TDAgentTools_lookup_user**: Seek user details from the Company User Lookup System.
48
- 11. **TDAgentTools_lookup_cloud_account**: Investigate cloud account information.
49
- 12. **TDAgentTools_send_email**: Simulate emailing from [email protected].
50
-
51
- > **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.
52
 
53
  [Track1 Demo link](https://youtu.be/c7Yg_jOD6J0)
54
 
55
 
56
- ### Track 3: Agentic Demo Showcase - TDAgent
57
 
58
- 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).
59
 
60
  #### Key Features:
61
  - **Intelligent API Interactions**: The agent autonomously interacts with APIs for data enrichment and analysis without explicit user guidance.
@@ -85,14 +85,15 @@ We aimed to:
85
 
86
  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!
87
 
 
88
 
89
  # TDA Agent
90
 
91
- # Development setup
92
 
93
  To start developing you need the following tools:
94
 
95
- * [uv](https://docs.astral.sh/uv/)
96
 
97
  To start, sync all the dependencies with `uv sync --all-groups`.
98
  Then, install the pre-commit hooks (`uv run pre-commit install`) to
 
14
 
15
  ---
16
 
17
+ # Welcome to **TDAgentTools & TDAgent**
18
 
19
+ 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.
20
 
21
  ## Team Introduction
22
 
23
  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:
24
 
25
+ - **Pedro Completo Bento**
26
+ - **Josep Pon Farreny**
27
+ - **Sofia Jeronimo dos Santos**
28
+ - **Rodrigo Dominguez Sanz**
29
+ - **Miguel Rodin**
30
 
31
  ## Project Overview
32
 
33
+ ### Track 1: MCP Tool - **TDAgentTools**
34
 
35
+ **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).
36
 
37
  #### Available Tools:
38
+ 1. ***TDAgentTools_get_url_http_content***: Retrieve URL content through an HTTP GET request.
39
+ 2. ***TDAgentTools_query_abuseipdb***: Query AbuseIPDB to check if an IP is reported for abusive behavior.
40
+ 3. ***TDAgentTools_query_rdap***: Gather information about internet resources such as domain names and IP addresses.
41
+ 4. ***TDAgentTools_get_virus_total_url_info***: Fetch URL information using VirusTotal URL Scanner.
42
+ 5. ***TDAgentTools_get_geolocation***: Obtain location details from an IP address.
43
+ 6. ***TDAgentTools_enumerate_dns***: Access DNS configuration details for a given domain.
44
+ 7. ***TDAgentTools_scrap_subdomains_for_domain***: Retrieve subdomains related to a domain.
45
+ 8. ***TDAgentTools_retrieve_ioc_from_threatfox***: Get potential IoC information from ThreatFox.
46
+ 9. ***TDAgentTools_get_stix_object_of_attack_id***: Access a STIX object using an ATT&CK ID.
47
+ 10. ***TDAgentTools_lookup_user***: Seek user details from the Company User Lookup System.
48
+ 11. ***TDAgentTools_lookup_cloud_account***: Investigate cloud account information.
49
+ 12. ***TDAgentTools_send_email***: Simulate emailing from [email protected].
50
+
51
+ > **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.
52
 
53
  [Track1 Demo link](https://youtu.be/c7Yg_jOD6J0)
54
 
55
 
56
+ ### Track 3: Agentic Demo Showcase - **TDAgent**
57
 
58
+ **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).
59
 
60
  #### Key Features:
61
  - **Intelligent API Interactions**: The agent autonomously interacts with APIs for data enrichment and analysis without explicit user guidance.
 
85
 
86
  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!
87
 
88
+ ---
89
 
90
  # TDA Agent
91
 
92
+ ## Development setup
93
 
94
  To start developing you need the following tools:
95
 
96
+ - [uv](https://docs.astral.sh/uv/)
97
 
98
  To start, sync all the dependencies with `uv sync --all-groups`.
99
  Then, install the pre-commit hooks (`uv run pre-commit install`) to
pyproject.toml CHANGED
@@ -22,6 +22,7 @@ dependencies = [
22
  "langchain-mcp-adapters>=0.1.1",
23
  "langchain-openai>=0.3.19",
24
  "langgraph>=0.4.7",
 
25
  "openai>=1.84.0",
26
  ]
27
 
 
22
  "langchain-mcp-adapters>=0.1.1",
23
  "langchain-openai>=0.3.19",
24
  "langgraph>=0.4.7",
25
+ "markdown>=3.8",
26
  "openai>=1.84.0",
27
  ]
28
 
requirements-dev.txt CHANGED
@@ -59,6 +59,7 @@ langgraph-prebuilt==0.2.2
59
  langgraph-sdk==0.1.70
60
  langsmith==0.3.43
61
  license-expression==30.4.1
 
62
  markdown-it-py==3.0.0
63
  markupsafe==3.0.2
64
  mcp==1.9.0
 
59
  langgraph-sdk==0.1.70
60
  langsmith==0.3.43
61
  license-expression==30.4.1
62
+ markdown==3.8
63
  markdown-it-py==3.0.0
64
  markupsafe==3.0.2
65
  mcp==1.9.0
requirements.txt CHANGED
@@ -51,6 +51,7 @@ langgraph-checkpoint==2.0.26
51
  langgraph-prebuilt==0.2.2
52
  langgraph-sdk==0.1.70
53
  langsmith==0.3.43
 
54
  markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
55
  markupsafe==3.0.2
56
  mcp==1.9.0
 
51
  langgraph-prebuilt==0.2.2
52
  langgraph-sdk==0.1.70
53
  langsmith==0.3.43
54
+ markdown==3.8
55
  markdown-it-py==3.0.0 ; sys_platform != 'emscripten'
56
  markupsafe==3.0.2
57
  mcp==1.9.0
tdagent/grchat.py CHANGED
@@ -5,6 +5,7 @@ import enum
5
  import os
6
  from collections import OrderedDict
7
  from collections.abc import Mapping, Sequence
 
8
  from types import MappingProxyType
9
  from typing import TYPE_CHECKING, Any
10
 
@@ -13,6 +14,7 @@ import botocore
13
  import botocore.exceptions
14
  import gradio as gr
15
  import gradio.themes as gr_themes
 
16
  from langchain_aws import ChatBedrock
17
  from langchain_core.callbacks import BaseCallbackHandler
18
  from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
@@ -119,6 +121,8 @@ MODEL_OPTIONS = OrderedDict( # Initialize with tuples to preserve options order
119
  ),
120
  )
121
 
 
 
122
 
123
  @dataclasses.dataclass
124
  class ToolInvocationInfo:
@@ -344,6 +348,7 @@ async def gr_connect_to_bedrock( # noqa: PLR0913
344
  ) -> str:
345
  """Initialize Bedrock agent."""
346
  global llm_agent # noqa: PLW0603
 
347
 
348
  if not access_key or not secret_key:
349
  return "❌ Please provide both Access Key ID and Secret Access Key"
@@ -384,7 +389,7 @@ async def gr_connect_to_hf(
384
  ) -> str:
385
  """Initialize Hugging Face agent."""
386
  global llm_agent # noqa: PLW0603
387
-
388
  llm, error = create_hf_llm(
389
  model_id,
390
  hf_access_token_textbox,
@@ -420,6 +425,7 @@ async def gr_connect_to_azure( # noqa: PLR0913
420
  ) -> str:
421
  """Initialize Hugging Face agent."""
422
  global llm_agent # noqa: PLW0603
 
423
 
424
  llm, error = create_azure_llm(
425
  model_id,
@@ -449,6 +455,7 @@ async def gr_connect_to_azure( # noqa: PLR0913
449
  # ) -> str:
450
  # """Initialize Hugging Face agent."""
451
  # global llm_agent
 
452
 
453
  # llm, error = create_openai_llm(model_id, nebius_access_token_textbox)
454
 
@@ -527,9 +534,31 @@ def _add_tools_trace_to_message(message: str) -> str:
527
  return f"{message}\n\n# Tools Trace\n\n" + "\n".join(traces)
528
 
529
 
530
- ## UI components ##
 
 
 
 
 
 
 
 
 
 
 
 
531
 
532
 
 
 
 
 
 
 
 
 
 
 
533
  with (
534
  gr.Blocks(
535
  theme=gr_themes.Origin(
@@ -538,162 +567,196 @@ with (
538
  font="sans-serif",
539
  ),
540
  title="TDAgent",
 
 
 
541
  ) as gr_app,
542
- gr.Row(),
543
  ):
544
- with gr.Column(scale=1):
545
- with gr.Accordion("🔌 MCP Servers", open=False):
546
- mcp_list = MutableCheckBoxGroup(
547
- values=[
548
- MutableCheckBoxGroupEntry(
549
- name="TDAgent tools",
550
- value="https://agents-mcp-hackathon-tdagenttools.hf.space/gradio_api/mcp/sse",
551
- ),
552
- ],
553
- label="MCP Servers",
554
- new_value_label="MCP endpoint",
555
- new_name_label="MCP endpoint name",
556
- new_value_placeholder="https://my-cool-mcp-server.com/mcp/sse",
557
- new_name_placeholder="Swiss army knife of MCPs",
558
- )
559
-
560
- with gr.Accordion("⚙️ Provider Configuration", open=True):
561
- model_provider = gr.Dropdown(
562
- choices=list(MODEL_OPTIONS.keys()),
563
- value=None,
564
- label="Select Model Provider",
565
- )
566
-
567
- ## Amazon Bedrock Configuration ##
568
- with gr.Group(visible=False) as aws_bedrock_conf_group:
569
- aws_access_key_textbox = gr.Textbox(
570
- label="AWS Access Key ID",
571
- type="password",
572
- placeholder="Enter your AWS Access Key ID",
573
- )
574
- aws_secret_key_textbox = gr.Textbox(
575
- label="AWS Secret Access Key",
576
- type="password",
577
- placeholder="Enter your AWS Secret Access Key",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
578
  )
579
- aws_region_dropdown = gr.Dropdown(
580
- label="AWS Region",
581
- choices=[
582
- "us-east-1",
583
- "us-west-2",
584
- "eu-west-1",
585
- "eu-central-1",
586
- "ap-southeast-1",
587
- ],
588
- value="eu-west-1",
589
  )
590
- aws_session_token_textbox = gr.Textbox(
591
- label="AWS Session Token",
592
- type="password",
593
- placeholder="Enter your AWS session token",
594
  )
595
 
596
- ## Huggingface Configuration ##
597
- with gr.Group(visible=False) as hf_conf_group:
598
- hf_token = gr.Textbox(
599
- label="HuggingFace Token",
600
- type="password",
601
- placeholder="Enter your Hugging Face Access Token",
602
- )
603
-
604
- ## Azure Configuration ##
605
- with gr.Group(visible=False) as azure_conf_group:
606
- azure_endpoint = gr.Textbox(
607
- label="Azure OpenAI Endpoint",
608
- type="text",
609
- placeholder="Enter your Azure OpenAI Endpoint",
610
- )
611
- azure_api_token = gr.Textbox(
612
- label="Azure Access Token",
613
- type="password",
614
- placeholder="Enter your Azure OpenAI Access Token",
615
  )
616
- azure_api_version = gr.Textbox(
617
- label="Azure OpenAI API Version",
618
- type="text",
619
- placeholder="Enter your Azure OpenAI API Version",
620
- value="2024-12-01-preview",
621
- )
622
-
623
- with gr.Accordion("🧠 Model Configuration", open=True):
624
- model_id_dropdown = gr.Dropdown(
625
- label="Select known model id or type your own below",
626
- choices=[],
627
- visible=False,
628
- )
629
- model_id_textbox = gr.Textbox(
630
- label="Model ID",
631
- type="text",
632
- placeholder="Enter the model ID",
633
- visible=False,
634
- interactive=True,
635
- )
636
 
637
- # Agent configuration options
638
- with gr.Group():
639
- agent_system_message_radio = gr.Radio(
640
- choices=list(AGENT_SYSTEM_MESSAGES.keys()),
641
- value=next(iter(AGENT_SYSTEM_MESSAGES.keys())),
642
- label="Agent type",
643
- info=(
644
- "Changes the system message to pre-condition the agent"
645
- " to act in a desired way."
646
- ),
647
- )
648
- agent_trace_tools_checkbox = gr.Checkbox(
649
- value=False,
650
- label="Trace tool calls",
651
- info="Add the invoked tools trace at the end of the message",
652
  )
653
-
654
- # Initialize the temperature and max tokens based on model specifications
655
- temperature = gr.Slider(
656
- label="Temperature",
657
- minimum=0.0,
658
- maximum=1.0,
659
- value=0.8,
660
- step=0.1,
661
  )
662
- max_tokens = gr.Slider(
663
- label="Max Tokens",
664
- minimum=128,
665
- maximum=8192,
666
- value=2048,
667
- step=64,
668
  )
669
 
670
- connect_aws_bedrock_btn = gr.Button(
671
- "🔌 Connect to Bedrock",
672
- variant="primary",
673
- visible=False,
674
- )
675
- connect_hf_btn = gr.Button(
676
- "🔌 Connect to Huggingface 🤗",
677
- variant="primary",
678
- visible=False,
679
- )
680
- connect_azure_btn = gr.Button(
681
- "🔌 Connect to Azure",
682
- variant="primary",
683
- visible=False,
684
- )
685
-
686
- status_textbox = gr.Textbox(label="Connection Status", interactive=False)
687
-
688
- with gr.Column(scale=2):
689
- chat_interface = gr.ChatInterface(
690
- fn=gr_chat_function,
691
- type="messages",
692
- examples=[], # Add examples if needed
693
- title="👩‍💻 TDAgent 👨‍💻",
694
- description="A simple threat analyst agent with MCP tools.",
695
- )
696
-
697
  ## UI Events ##
698
 
699
  def _toggle_model_choices_ui(
@@ -728,6 +791,18 @@ with (
728
  is_azure = provider == "Azure OpenAI"
729
  return gr.update(visible=is_azure), gr.update(visible=is_azure)
730
 
 
 
 
 
 
 
 
 
 
 
 
 
731
  ## Connect Event Listeners ##
732
 
733
  model_provider.change(
@@ -810,6 +885,14 @@ with (
810
  inputs=[model_id_dropdown, model_provider],
811
  outputs=[model_id_textbox],
812
  )
 
 
 
 
 
 
 
 
813
 
814
  ## Entry Point ##
815
 
 
5
  import os
6
  from collections import OrderedDict
7
  from collections.abc import Mapping, Sequence
8
+ from pathlib import Path
9
  from types import MappingProxyType
10
  from typing import TYPE_CHECKING, Any
11
 
 
14
  import botocore.exceptions
15
  import gradio as gr
16
  import gradio.themes as gr_themes
17
+ import markdown
18
  from langchain_aws import ChatBedrock
19
  from langchain_core.callbacks import BaseCallbackHandler
20
  from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
121
  ),
122
  )
123
 
124
+ CONNECT_STATE_DEFAULT = gr.State()
125
+
126
 
127
  @dataclasses.dataclass
128
  class ToolInvocationInfo:
 
348
  ) -> str:
349
  """Initialize Bedrock agent."""
350
  global llm_agent # noqa: PLW0603
351
+ CONNECT_STATE_DEFAULT.value = True
352
 
353
  if not access_key or not secret_key:
354
  return "❌ Please provide both Access Key ID and Secret Access Key"
 
389
  ) -> str:
390
  """Initialize Hugging Face agent."""
391
  global llm_agent # noqa: PLW0603
392
+ CONNECT_STATE_DEFAULT.value = True
393
  llm, error = create_hf_llm(
394
  model_id,
395
  hf_access_token_textbox,
 
425
  ) -> str:
426
  """Initialize Hugging Face agent."""
427
  global llm_agent # noqa: PLW0603
428
+ CONNECT_STATE_DEFAULT.value = True
429
 
430
  llm, error = create_azure_llm(
431
  model_id,
 
455
  # ) -> str:
456
  # """Initialize Hugging Face agent."""
457
  # global llm_agent
458
+ # connected_state.value = True
459
 
460
  # llm, error = create_openai_llm(model_id, nebius_access_token_textbox)
461
 
 
534
  return f"{message}\n\n# Tools Trace\n\n" + "\n".join(traces)
535
 
536
 
537
+ def _read_markdown_body_as_html(path: str = "README.md") -> str:
538
+ with Path(path).open(encoding="utf-8") as f: # Default mode is "r"
539
+ lines = f.readlines()
540
+
541
+ # Skip YAML front matter if present
542
+ if lines and lines[0].strip() == "---":
543
+ for i in range(1, len(lines)):
544
+ if lines[i].strip() == "---":
545
+ lines = lines[i + 1 :] # skip metadata block
546
+ break
547
+
548
+ markdown_body = "".join(lines).strip()
549
+ return markdown.markdown(markdown_body)
550
 
551
 
552
+ ## UI components ##
553
+ custom_css = """
554
+ .main-header {
555
+ background: linear-gradient(135deg, #00a388 0%, #ffae00 100%);
556
+ padding: 30px;
557
+ border-radius: 5px;
558
+ margin-bottom: 20px;
559
+ text-align: center;
560
+ }
561
+ """
562
  with (
563
  gr.Blocks(
564
  theme=gr_themes.Origin(
 
567
  font="sans-serif",
568
  ),
569
  title="TDAgent",
570
+ fill_height=True,
571
+ fill_width=True,
572
+ css=custom_css,
573
  ) as gr_app,
 
574
  ):
575
+ gr.HTML(
576
+ """
577
+ <div class="main-header">
578
+ <h1>👩‍💻 TDAgentTools & TDAgent 👨‍💻</h1>
579
+ <p style="font-size: 1.2em; margin: 10px 0 0 0;">
580
+ Empowering Cybersecurity with Agentic AI
581
+ </p>
582
+ </div>
583
+ """,
584
+ )
585
+ with gr.Tabs():
586
+ with gr.TabItem("About"), gr.Row():
587
+ html_content = _read_markdown_body_as_html("README.md")
588
+ gr.Markdown(html_content)
589
+
590
+ with gr.TabItem("TDAgent"), gr.Row():
591
+ with gr.Column(scale=1):
592
+ with gr.Accordion("🔌 MCP Servers", open=False):
593
+ mcp_list = MutableCheckBoxGroup(
594
+ values=[
595
+ MutableCheckBoxGroupEntry(
596
+ name="TDAgent tools",
597
+ value="https://agents-mcp-hackathon-tdagenttools.hf.space/gradio_api/mcp/sse",
598
+ ),
599
+ ],
600
+ label="MCP Servers",
601
+ new_value_label="MCP endpoint",
602
+ new_name_label="MCP endpoint name",
603
+ new_value_placeholder="https://my-cool-mcp-server.com/mcp/sse",
604
+ new_name_placeholder="Swiss army knife of MCPs",
605
+ )
606
+
607
+ with gr.Accordion("⚙️ Provider Configuration", open=True):
608
+ model_provider = gr.Dropdown(
609
+ choices=list(MODEL_OPTIONS.keys()),
610
+ value=None,
611
+ label="Select Model Provider",
612
+ )
613
+
614
+ ## Amazon Bedrock Configuration ##
615
+ with gr.Group(visible=False) as aws_bedrock_conf_group:
616
+ aws_access_key_textbox = gr.Textbox(
617
+ label="AWS Access Key ID",
618
+ type="password",
619
+ placeholder="Enter your AWS Access Key ID",
620
+ )
621
+ aws_secret_key_textbox = gr.Textbox(
622
+ label="AWS Secret Access Key",
623
+ type="password",
624
+ placeholder="Enter your AWS Secret Access Key",
625
+ )
626
+ aws_region_dropdown = gr.Dropdown(
627
+ label="AWS Region",
628
+ choices=[
629
+ "us-east-1",
630
+ "us-west-2",
631
+ "eu-west-1",
632
+ "eu-central-1",
633
+ "ap-southeast-1",
634
+ ],
635
+ value="eu-west-1",
636
+ )
637
+ aws_session_token_textbox = gr.Textbox(
638
+ label="AWS Session Token",
639
+ type="password",
640
+ placeholder="Enter your AWS session token",
641
+ )
642
+
643
+ ## Huggingface Configuration ##
644
+ with gr.Group(visible=False) as hf_conf_group:
645
+ hf_token = gr.Textbox(
646
+ label="HuggingFace Token",
647
+ type="password",
648
+ placeholder="Enter your Hugging Face Access Token",
649
+ )
650
+
651
+ ## Azure Configuration ##
652
+ with gr.Group(visible=False) as azure_conf_group:
653
+ azure_endpoint = gr.Textbox(
654
+ label="Azure OpenAI Endpoint",
655
+ type="text",
656
+ placeholder="Enter your Azure OpenAI Endpoint",
657
+ )
658
+ azure_api_token = gr.Textbox(
659
+ label="Azure Access Token",
660
+ type="password",
661
+ placeholder="Enter your Azure OpenAI Access Token",
662
+ )
663
+ azure_api_version = gr.Textbox(
664
+ label="Azure OpenAI API Version",
665
+ type="text",
666
+ placeholder="Enter your Azure OpenAI API Version",
667
+ value="2024-12-01-preview",
668
+ )
669
+
670
+ with gr.Accordion("🧠 Model Configuration", open=True):
671
+ model_id_dropdown = gr.Dropdown(
672
+ label="Select known model id or type your own below",
673
+ choices=[],
674
+ visible=False,
675
+ )
676
+ model_id_textbox = gr.Textbox(
677
+ label="Model ID",
678
+ type="text",
679
+ placeholder="Enter the model ID",
680
+ visible=False,
681
+ interactive=True,
682
+ )
683
+
684
+ # Agent configuration options
685
+ with gr.Group():
686
+ agent_system_message_radio = gr.Radio(
687
+ choices=list(AGENT_SYSTEM_MESSAGES.keys()),
688
+ value=next(iter(AGENT_SYSTEM_MESSAGES.keys())),
689
+ label="Agent type",
690
+ info=(
691
+ "Changes the system message to pre-condition the agent"
692
+ " to act in a desired way."
693
+ ),
694
+ )
695
+ agent_trace_tools_checkbox = gr.Checkbox(
696
+ value=False,
697
+ label="Trace tool calls",
698
+ info=(
699
+ "Add the invoked tools trace at the end of the"
700
+ " message"
701
+ ),
702
+ )
703
+
704
+ # Initialize the temperature and max tokens based on model specs
705
+ temperature = gr.Slider(
706
+ label="Temperature",
707
+ minimum=0.0,
708
+ maximum=1.0,
709
+ value=0.8,
710
+ step=0.1,
711
+ )
712
+ max_tokens = gr.Slider(
713
+ label="Max Tokens",
714
+ minimum=128,
715
+ maximum=8192,
716
+ value=2048,
717
+ step=64,
718
+ )
719
+
720
+ connect_aws_bedrock_btn = gr.Button(
721
+ "🔌 Connect to Bedrock",
722
+ variant="primary",
723
+ visible=False,
724
  )
725
+ connect_hf_btn = gr.Button(
726
+ "🔌 Connect to Huggingface 🤗",
727
+ variant="primary",
728
+ visible=False,
 
 
 
 
 
 
729
  )
730
+ connect_azure_btn = gr.Button(
731
+ "🔌 Connect to Azure",
732
+ variant="primary",
733
+ visible=False,
734
  )
735
 
736
+ status_textbox = gr.Textbox(
737
+ label="Connection Status",
738
+ interactive=False,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
739
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
740
 
741
+ with gr.Column(scale=2):
742
+ chat_interface = gr.ChatInterface(
743
+ fn=gr_chat_function,
744
+ type="messages",
745
+ examples=[], # Add examples if needed
746
+ description="A simple threat analyst agent with MCP tools.",
 
 
 
 
 
 
 
 
 
747
  )
748
+ with gr.TabItem("Demo"):
749
+ gr.Markdown(
750
+ """
751
+ This is a demo of TDAgent, a simple threat analyst agent with MCP tools.
752
+ You can configure the agent to use different LLM providers and connect to
753
+ various MCP servers to access tools.
754
+ """,
 
755
  )
756
+ gr.HTML(
757
+ """<iframe width="560" height="315" src="https://youtu.be/C6Z9EOW-3lE?feature=shared" frameborder="0" allowfullscreen></iframe>""", # noqa: E501
 
 
 
 
758
  )
759
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
760
  ## UI Events ##
761
 
762
  def _toggle_model_choices_ui(
 
791
  is_azure = provider == "Azure OpenAI"
792
  return gr.update(visible=is_azure), gr.update(visible=is_azure)
793
 
794
+ # Initialize a flag to check if connected
795
+
796
+ def _on_change_model_configuration(*args: str) -> Any: # noqa: ARG001
797
+ # If model configuration changes after connecting, issue a warning
798
+ if CONNECT_STATE_DEFAULT.value:
799
+ CONNECT_STATE_DEFAULT.value = False # Reset the state
800
+ return gr.Warning(
801
+ "When changing model configuration, you need to reconnect.",
802
+ duration=5,
803
+ )
804
+ return gr.update()
805
+
806
  ## Connect Event Listeners ##
807
 
808
  model_provider.change(
 
885
  inputs=[model_id_dropdown, model_provider],
886
  outputs=[model_id_textbox],
887
  )
888
+ model_provider.change(
889
+ _on_change_model_configuration,
890
+ inputs=[model_provider],
891
+ )
892
+ model_id_dropdown.change(
893
+ _on_change_model_configuration,
894
+ inputs=[model_id_dropdown, model_provider],
895
+ )
896
 
897
  ## Entry Point ##
898
 
uv.lock CHANGED
@@ -1148,6 +1148,15 @@ wheels = [
1148
  { 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" },
1149
  ]
1150
 
 
 
 
 
 
 
 
 
 
1151
  [[package]]
1152
  name = "markdown-it-py"
1153
  version = "3.0.0"
@@ -2869,6 +2878,7 @@ dependencies = [
2869
  { name = "langchain-mcp-adapters" },
2870
  { name = "langchain-openai" },
2871
  { name = "langgraph" },
 
2872
  { name = "openai" },
2873
  ]
2874
 
@@ -2897,6 +2907,7 @@ requires-dist = [
2897
  { name = "langchain-mcp-adapters", specifier = ">=0.1.1" },
2898
  { name = "langchain-openai", specifier = ">=0.3.19" },
2899
  { name = "langgraph", specifier = ">=0.4.7" },
 
2900
  { name = "openai", specifier = ">=1.84.0" },
2901
  ]
2902
 
 
1148
  { 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" },
1149
  ]
1150
 
1151
+ [[package]]
1152
+ name = "markdown"
1153
+ version = "3.8"
1154
+ source = { registry = "https://pypi.org/simple" }
1155
+ 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" }
1156
+ wheels = [
1157
+ { 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" },
1158
+ ]
1159
+
1160
  [[package]]
1161
  name = "markdown-it-py"
1162
  version = "3.0.0"
 
2878
  { name = "langchain-mcp-adapters" },
2879
  { name = "langchain-openai" },
2880
  { name = "langgraph" },
2881
+ { name = "markdown" },
2882
  { name = "openai" },
2883
  ]
2884
 
 
2907
  { name = "langchain-mcp-adapters", specifier = ">=0.1.1" },
2908
  { name = "langchain-openai", specifier = ">=0.3.19" },
2909
  { name = "langgraph", specifier = ">=0.4.7" },
2910
+ { name = "markdown", specifier = ">=3.8" },
2911
  { name = "openai", specifier = ">=1.84.0" },
2912
  ]
2913