utils/export_handler.py

"""Export handler for comparison results."""

import csv
import json
from datetime import datetime
from io import StringIO
from typing import Any, Dict, List


class ExportHandler:
    """Handles exporting comparison results to various formats."""

    @staticmethod
    def export_to_json(
        prompt: str,
        tools_enabled: List[str],
        comparisons: List[Dict[str, Any]],
    ) -> str:
        """Export comparison results to JSON format.

        Args:
            prompt: User prompt
            tools_enabled: List of enabled tool names
            comparisons: List of comparison results

        Returns:
            JSON string
        """
        export_data = {
            "timestamp": datetime.now().isoformat(),
            "prompt": prompt,
            "tools_enabled": tools_enabled,
            "comparisons": [],
        }

        for comp in comparisons:
            metrics = comp.get("metrics")
            export_data["comparisons"].append(
                {
                    "model": comp.get("model_name", "Unknown"),
                    "output": comp.get("output", ""),
                    "tools_called": comp.get("tools_used", []),
                    "metrics": metrics.to_dict() if hasattr(metrics, "to_dict") else {},
                }
            )

        return json.dumps(export_data, indent=2)

    @staticmethod
    def export_to_csv(
        prompt: str,
        tools_enabled: List[str],
        comparisons: List[Dict[str, Any]],
    ) -> str:
        """Export comparison results to CSV format.

        Args:
            prompt: User prompt
            tools_enabled: List of enabled tool names
            comparisons: List of comparison results

        Returns:
            CSV string
        """
        output = StringIO()
        writer = csv.writer(output)

        # Write header
        writer.writerow(
            [
                "Timestamp",
                "Prompt",
                "Tools Enabled",
                "Model",
                "Total Time (s)",
                "Tokens/sec",
                "Tokens Generated",
                "Tool Calls",
                "Tool Exec Time (s)",
                "Success",
                "Output Preview",
            ]
        )

        # Write data rows
        timestamp = datetime.now().isoformat()
        tools_str = ", ".join(tools_enabled)

        for comp in comparisons:
            metrics = comp.get("metrics")
            if hasattr(metrics, "to_dict"):
                metrics_dict = metrics.to_dict()
            else:
                metrics_dict = {}

            output_preview = comp.get("output", "")[:100] + "..." if len(comp.get("output", "")) > 100 else comp.get("output", "")

            writer.writerow(
                [
                    timestamp,
                    prompt,
                    tools_str,
                    comp.get("model_name", "Unknown"),
                    metrics_dict.get("total_time", 0),
                    metrics_dict.get("tokens_per_second", 0),
                    metrics_dict.get("tokens_generated", 0),
                    metrics_dict.get("tool_calls_count", 0),
                    metrics_dict.get("tool_execution_time", 0),
                    metrics_dict.get("success", False),
                    output_preview,
                ]
            )

        return output.getvalue()

    @staticmethod
    def create_download_filename(format_type: str) -> str:
        """Create filename for download.

        Args:
            format_type: 'json' or 'csv'

        Returns:
            Filename with timestamp
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"model_comparison_{timestamp}.{format_type}"