app.py

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