app.py

import gradio as gr
import requests
import os

# API configuration
API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8000")
API_KEY = os.getenv("API_KEY", "")

def analyze_code(code: str, language: str, model: str) -> str:
    """Send code to the vulnerability analysis API and return results."""
    if not code.strip():
        return "Please enter some code to analyze."

    api_url = f"{API_BASE_URL}/internal/analyze"

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }

    payload = {
        "code": code,
        "model": model,
        "language": language
    }

    try:
        response = requests.post(api_url, json=payload, headers=headers, timeout=60)

        if response.status_code == 403:
            error_detail = response.json().get("detail", "Token limit exceeded.")
            return f"Error: {error_detail}"

        if response.status_code != 200:
            error_detail = response.json().get("detail", f"API error: {response.status_code}")
            return f"Error: {error_detail}"

        result = response.json()

        # Format the result
        status = result.get("result", {}).get("status", "unknown")
        cwe_type = result.get("result", {}).get("cweType", "N/A")
        model_used = result.get("result", {}).get("model", model)
        explanation = result.get("result", {}).get("response", "")
        if "## Final Answer" in explanation:
            explanation = explanation.split("## Final Answer")[0].strip()
        if "nopolicy" in model_used:
            model_used = "VirtueGuard Code"
        if status == "yes":
            output = f"⚠️ **Vulnerability Detected**\n\n"
            output += f"**CWE Type:** {cwe_type}\n"
            output += f"**Model:** {model_used}\n\n"
            output += f"**Analysis:**\n{explanation}"
        else:
            output = f"✅ **No Vulnerability Detected**\n\n"
            output += f"**Model:** {model_used}\n\n"
            output += f"**Analysis:**\n{explanation}"

        return output

    except requests.exceptions.Timeout:
        return "Error: Request timed out. Please try again."
    except requests.exceptions.ConnectionError:
        return f"Error: Could not connect to API at {API_BASE_URL}"
    except Exception as e:
        return f"Error: {str(e)}"

# Language options
LANGUAGES = [
    "python", "javascript", "typescript", "java", "c", "cpp",
    "csharp", "go", "rust", "php", "ruby", "swift", "kotlin"
]

# Model options
MODELS = ["virtueguard-code", "claude-4-sonnet", "gpt-4.1"]

# Create Gradio interface
with gr.Blocks(title="VulnLLM-R Demo") as demo:
    gr.Markdown("# VulnLLM-R Demo")
    gr.Markdown("Analyze your code for potential security vulnerabilities using VulnLLM-R.")

    with gr.Row():
        with gr.Column(scale=2):
            code_input = gr.Code(
                label="Code to Analyze",
                language="python",
                lines=15
            )
        with gr.Column(scale=1):
            language_dropdown = gr.Dropdown(
                choices=LANGUAGES,
                value="python",
                label="Programming Language"
            )
            model_dropdown = gr.Dropdown(
                choices=MODELS,
                value="virtueguard-code",
                label="Model"
            )
            analyze_btn = gr.Button("🔍 Analyze Code", variant="primary")

    result_output = gr.Markdown(label="Analysis Result")

    analyze_btn.click(
        fn=analyze_code,
        inputs=[code_input, language_dropdown, model_dropdown],
        outputs=result_output
    )

if __name__ == "__main__":
    demo.launch()