app.py

"""Gradio interface for Model Comparison Playground."""

import os
from typing import List, Tuple

import gradio as gr

from models.model_config import AVAILABLE_MODELS, get_model_choices, get_model_key_by_name
from models.model_runner import ModelRunner
from tools.tool_definitions import TOOLS, get_tool_names, get_tools_subset
from utils.export_handler import ExportHandler


# Global state
hf_token = os.getenv("HF_TOKEN")
export_handler = ExportHandler()
last_comparison_results = None


def run_comparison(
    prompt: str,
    model1_name: str,
    model2_name: str,
    tool_calculator: bool,
    tool_weather: bool,
    tool_search: bool,
    tool_time: bool,
    tool_convert: bool,
) -> Tuple[str, str, str, str, gr.update, gr.update]:
    """Run comparison between two models.

    Args:
        prompt: User prompt
        model1_name: First model name
        model2_name: Second model name
        tool_calculator: Enable calculator tool
        tool_weather: Enable weather tool
        tool_search: Enable web search tool
        tool_time: Enable time tool
        tool_convert: Enable unit conversion tool

    Returns:
        Tuple of (output1, metrics1, output2, metrics2, json_button, csv_button)
    """
    global last_comparison_results

    if not prompt:
        return (
            "Please enter a prompt.",
            "",
            "Please enter a prompt.",
            "",
            gr.update(interactive=False),
            gr.update(interactive=False),
        )

    # Determine which tools are enabled
    enabled_tools = []
    if tool_calculator:
        enabled_tools.append("calculator")
    if tool_weather:
        enabled_tools.append("get_weather")
    if tool_search:
        enabled_tools.append("web_search")
    if tool_time:
        enabled_tools.append("get_current_time")
    if tool_convert:
        enabled_tools.append("convert_units")

    # Get tools subset
    tools = get_tools_subset(enabled_tools) if enabled_tools else []

    try:
        # Get model keys
        model1_key = get_model_key_by_name(model1_name)
        model2_key = get_model_key_by_name(model2_name)

        model1_info = AVAILABLE_MODELS[model1_key]
        model2_info = AVAILABLE_MODELS[model2_key]

        # Initialize runners
        runner1 = ModelRunner(model1_info.id, model1_info.name, hf_token)
        runner2 = ModelRunner(model2_info.id, model2_info.name, hf_token)

        # Run both models
        result1 = runner1.run(prompt, tools)
        result2 = runner2.run(prompt, tools)

        # Format outputs
        output1 = f"**{model1_info.name}**\n\n{result1['output']}\n\n**Tools Used:** {', '.join(result1['tools_used']) if result1['tools_used'] else 'None'}"
        output2 = f"**{model2_info.name}**\n\n{result2['output']}\n\n**Tools Used:** {', '.join(result2['tools_used']) if result2['tools_used'] else 'None'}"

        metrics1 = result1["metrics"].format_for_display()
        metrics2 = result2["metrics"].format_for_display()

        # Store results for export
        last_comparison_results = {
            "prompt": prompt,
            "tools_enabled": enabled_tools,
            "comparisons": [
                {
                    "model_name": model1_info.name,
                    "output": result1["output"],
                    "tools_used": result1["tools_used"],
                    "metrics": result1["metrics"],
                },
                {
                    "model_name": model2_info.name,
                    "output": result2["output"],
                    "tools_used": result2["tools_used"],
                    "metrics": result2["metrics"],
                },
            ],
        }

        return (
            output1,
            metrics1,
            output2,
            metrics2,
            gr.update(interactive=True),
            gr.update(interactive=True),
        )

    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return (
            error_msg,
            "",
            error_msg,
            "",
            gr.update(interactive=False),
            gr.update(interactive=False),
        )


def export_json() -> Tuple[str, str]:
    """Export results to JSON.

    Returns:
        Tuple of (file_path, filename)
    """
    if not last_comparison_results:
        return None, "No results to export"

    json_content = export_handler.export_to_json(
        last_comparison_results["prompt"],
        last_comparison_results["tools_enabled"],
        last_comparison_results["comparisons"],
    )

    filename = export_handler.create_download_filename("json")
    filepath = f"/tmp/{filename}"

    with open(filepath, "w") as f:
        f.write(json_content)

    return filepath


def export_csv() -> Tuple[str, str]:
    """Export results to CSV.

    Returns:
        Tuple of (file_path, filename)
    """
    if not last_comparison_results:
        return None, "No results to export"

    csv_content = export_handler.export_to_csv(
        last_comparison_results["prompt"],
        last_comparison_results["tools_enabled"],
        last_comparison_results["comparisons"],
    )

    filename = export_handler.create_download_filename("csv")
    filepath = f"/tmp/{filename}"

    with open(filepath, "w") as f:
        f.write(csv_content)

    return filepath


# Build Gradio interface
with gr.Blocks(title="Agent Model Comparison Playground", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Agent Model Comparison Playground")
    gr.Markdown(
        "Compare small agent-capable models side-by-side. Test how different models handle the same prompt with tool-calling capabilities."
    )

    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Your Prompt",
                placeholder="Example: What's the weather in Paris and what's 15% tip on €45?",
                lines=3,
            )

    gr.Markdown("### 🔧 Available Tools")
    with gr.Row():
        tool_calculator = gr.Checkbox(label="Calculator", value=True)
        tool_weather = gr.Checkbox(label="Weather", value=True)
        tool_search = gr.Checkbox(label="Web Search", value=False)
        tool_time = gr.Checkbox(label="Date/Time", value=False)
        tool_convert = gr.Checkbox(label="Unit Converter", value=False)

    gr.Markdown("### 🎯 Model Selection")
    with gr.Row():
        model1_dropdown = gr.Dropdown(
            choices=get_model_choices(),
            value=get_model_choices()[0],
            label="Model 1",
        )
        model2_dropdown = gr.Dropdown(
            choices=get_model_choices(),
            value=get_model_choices()[1] if len(get_model_choices()) > 1 else get_model_choices()[0],
            label="Model 2",
        )

    run_button = gr.Button("🚀 Run Comparison", variant="primary", size="lg")

    gr.Markdown("### 📊 Results")
    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Model 1 Output")
            output1 = gr.Markdown()
            gr.Markdown("#### Model 1 Metrics")
            metrics1 = gr.Textbox(label="Performance Metrics", lines=6)

        with gr.Column():
            gr.Markdown("#### Model 2 Output")
            output2 = gr.Markdown()
            gr.Markdown("#### Model 2 Metrics")
            metrics2 = gr.Textbox(label="Performance Metrics", lines=6)

    gr.Markdown("### 💾 Export Results")
    with gr.Row():
        json_button = gr.DownloadButton("Export as JSON", interactive=False)
        csv_button = gr.DownloadButton("Export as CSV", interactive=False)

    # Event handlers
    run_button.click(
        fn=run_comparison,
        inputs=[
            prompt_input,
            model1_dropdown,
            model2_dropdown,
            tool_calculator,
            tool_weather,
            tool_search,
            tool_time,
            tool_convert,
        ],
        outputs=[output1, metrics1, output2, metrics2, json_button, csv_button],
    )

    json_button.click(fn=export_json, outputs=json_button)
    csv_button.click(fn=export_csv, outputs=csv_button)

    gr.Markdown("---")
    gr.Markdown(
        """
    ### 📚 About

    This playground allows you to compare small agent-capable language models on the same task with tool-calling capabilities.

    **Available Models:**
    - **Zephyr 7B Beta**: Fine-tuned Mistral variant optimized for helpfulness
    - **Mistral 7B Instruct v0.2**: Efficient 7B parameter model
    - **Phi-3 Mini**: Compact 3.8B parameter model
    - **Qwen 2.5 7B Instruct**: Advanced model with strong tool calling

    **Available Tools:**
    - **Calculator**: Perform mathematical calculations
    - **Weather**: Get current weather (simulated)
    - **Web Search**: Search for information (simulated)
    - **Date/Time**: Get current time in different timezones
    - **Unit Converter**: Convert between units

    **Metrics Tracked:**
    - Total inference time
    - Tokens per second
    - Number of tool calls
    - Tool execution time
    """
    )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)