app.py

#!/usr/bin/env python3
"""
Gradio UI for the Vulnerability Intelligence Agent (VIA).
This provides a chat interface to interact with the VIA using natural language.
"""
import os
import sys
import argparse
import logging
from typing import Dict, List, Any, Optional

import gradio as gr
from smolagents import CodeAgent, HfApiModel, GradioUI
from smolagents.tools import load_tool, tool

# Asegurarse de que el directorio actual esté en sys.path para que los imports funcionen
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from agents.coordinator_agent import search_vulnerabilities_for_software, get_vulnerability_details
from tools import utils

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = utils.setup_logger("gradio_ui")

# Cargar las herramientas básicas usando las que ya existen en smolagents
final_answer = load_tool("smolagents/final_answer", trust_remote_code=True)

def get_agent_description():
    """
    Get the description for the agent.
    """
    return """
    # 🔐 Vulnerability Intelligence Agent (VIA)
    
    I am an intelligent agent designed to help you find vulnerabilities in software and systems.
    
    ## What I can do:
    - Search for known vulnerabilities in software by name and version
    - Provide detailed information about specific vulnerabilities (CVE, CWE, etc.)
    - Generate reports about vulnerabilities
    
    ## How to use me:
    - Ask about vulnerabilities in specific software, e.g., "Find vulnerabilities in OpenSSL 1.1.1k"
    - Ask about a specific vulnerability, e.g., "Tell me about CVE-2021-44228"
    - Use natural language to describe what you're looking for
    
    ## Examples:
    - "What vulnerabilities exist in Apache 2.4.54?"
    - "Are there any critical vulnerabilities in log4j 2.14.1?"
    - "Give me details about CVE-2021-44228"
    - "What security issues should I be aware of in OpenSSL 1.1.1k?"
    """

def create_parser():
    """Create command line argument parser."""
    parser = argparse.ArgumentParser(description="Vulnerability Intelligence Agent (VIA) UI")
    parser.add_argument("--port", type=int, default=7860, help="Port to run the Gradio app on")
    parser.add_argument("--host", type=str, default="127.0.0.1", help="Host to run the Gradio app on")
    parser.add_argument("--model", type=str, default="Qwen/Qwen2.5-Coder-32B-Instruct", 
                        help="HuggingFace model ID to use")
    parser.add_argument("--share", action="store_true", help="Create a public link")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    
    return parser

def main():
    """Main entry point for the Gradio UI."""
    args = create_parser().parse_args()
    
    # Configure logging level
    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=log_level)
    
    # Initialize the model
    model = HfApiModel(
        max_tokens=2096,
        temperature=0.5,
        model_id=args.model,
        custom_role_conversions=None,
    )
    
    # Initialize the agent con las herramientas ya existentes y las que hemos creado
    agent = CodeAgent(
        model=model,
        tools=[search_vulnerabilities_for_software, get_vulnerability_details, final_answer],
        max_steps=10,
        verbosity_level=2 if args.verbose else 1,
    )
    
    # Create Gradio UI
    ui = GradioUI(agent)
    
    # Launch the UI
    ui.launch(
        share=args.share,
        server_name=args.host,
        server_port=args.port,
        show_api=False,
        favicon_path=None,
        allowed_paths=[],
        app_kwargs={
            "title": "🔐 Vulnerability Intelligence Agent (VIA)",
            "description": get_agent_description(),
            "theme": gr.themes.Base(),
        },
    )

if __name__ == "__main__":
    main()