File size: 2,880 Bytes
2e82565 eea2f4b 2e82565 eea2f4b 8c01ffb eea2f4b 8c01ffb eea2f4b 8c01ffb eea2f4b 8c01ffb eea2f4b 8fe992b eea2f4b 2e82565 5a945fb eea2f4b 2e82565 eea2f4b 2e82565 eea2f4b 2e82565 eea2f4b 2e82565 eea2f4b 5a945fb eea2f4b 2e82565 5a945fb eea2f4b 5a945fb eea2f4b 5a945fb eea2f4b 5a945fb eea2f4b 5a945fb eea2f4b 9b5b26a 2e82565 eea2f4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import json
import yaml
from dotenv import load_dotenv
import gradio as gr
from smolagents import CodeAgent
from smolagents.models import HfApiModel
from tools.final_answer import FinalAnswerTool
from tools.web_search import DuckDuckGoSearchTool
from tools.visit_webpage import VisitWebpageTool
from tools.vuln_search import VulnerabilitySearchTool
# Load environment variables
load_dotenv()
def load_agent_config():
"""Load agent configuration from agent.json"""
with open('agent.json', 'r') as f:
return json.load(f)
def load_prompts():
"""Load prompt templates from prompts.yaml"""
with open('prompts.yaml', 'r') as f:
return yaml.safe_load(f)
def initialize_tools():
"""Initialize agent tools"""
tools = {
'final_answer': FinalAnswerTool(),
'web_search': DuckDuckGoSearchTool(),
'visit_webpage': VisitWebpageTool(),
'vuln_search': VulnerabilitySearchTool()
}
return tools
def create_agent():
"""Create and configure the vulnerability agent"""
config = load_agent_config()
prompts = load_prompts()
# Configure model
model_config = config['agent_config']['model']
model = HfApiModel(
model_id=model_config['model_id'],
max_tokens=model_config['max_tokens'],
temperature=model_config['temperature']
)
# Initialize tools
tools = initialize_tools()
# Create agent
agent = CodeAgent(
model=model,
tools=tools,
max_steps=config['agent_config']['max_steps'],
verbosity_level=config['agent_config']['verbosity_level']
)
return agent, prompts
def process_query(message, history):
"""Process a user query in chat format"""
agent, prompts = create_agent()
# Format the prompt
template = prompts['user_prompt']
formatted_prompt = template.format(query=message)
# Execute agent
system_prompt = prompts['system_prompt']
result = agent.run(formatted_prompt, system_prompt=system_prompt)
return result
# Gradio Chat Interface
def create_interface():
"""Create the Gradio chat interface"""
with gr.Blocks(title="Vulnerability Intelligence Agent") as interface:
gr.Markdown("# Vulnerability Intelligence Agent (VIA)")
chatbot = gr.Chatbot(
[],
elem_id="chatbot",
bubble_full_width=False,
avatar_images=(None, "🤖"),
height=600,
)
txt = gr.Textbox(
show_label=False,
placeholder="Enter your security query...",
container=False
)
txt.submit(
process_query,
[txt, chatbot],
[chatbot]
)
return interface
if __name__ == "__main__":
interface = create_interface()
interface.launch() |