src/utils/visualization.py

import plotly.express as px
import plotly.graph_objects as go

def create_radar_chart(metrics_data, title):
    """
    Create a Plotly radar chart for a set of metrics
    
    Args:
        metrics_data (dict): Dictionary of metric names and values
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Radar chart
    """
    categories = list(metrics_data.keys())
    values = list(metrics_data.values())
    
    fig = go.Figure()
    
    fig.add_trace(go.Scatterpolar(
        r=values,
        theta=categories,
        fill='toself',
        name=title
    ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 5]  # Assuming metrics are on a 1-5 scale
            )
        ),
        title=title,
        showlegend=True
    )
    
    return fig

def create_bar_chart(metrics_data, title):
    """
    Create a Plotly bar chart for a set of metrics
    
    Args:
        metrics_data (dict): Dictionary of metric names and values
        title (str): Chart title
        
    Returns:
        plotly.express.Figure: Bar chart
    """
    fig = px.bar(
        x=list(metrics_data.keys()),
        y=list(metrics_data.values()),
        labels={'x': 'Metric', 'y': 'Score'},
        title=title
    )
    
    fig.update_layout(
        xaxis_title="Metric",
        yaxis_title="Score",
        yaxis=dict(range=[0, 5]),  # Assuming metrics are on a 1-5 scale
        xaxis_tickangle=-45
    )
    
    return fig

def visualize_eval_data(eval_results, chart_type='radar'):
    """
    Generate Plotly visualizations for evaluation results
    
    Args:
        eval_results (dict): Evaluation results
        chart_type (str): Type of chart to create ('radar' or 'bar')
        
    Returns:
        list: List of tuples (title, chart)
    """
    charts = []
    
    if not eval_results or "message" in eval_results:
        return []
    
    # Handle different result structures
    if "journeyEvals" in eval_results and "aiScriptEvals" in eval_results:
        # Query 1 or ungrouped Query 2 result
        for eval_name, metrics in eval_results["journeyEvals"].items():
            if chart_type == 'radar':
                chart = create_radar_chart(metrics, f"Journey Eval: {eval_name}")
            else:
                chart = create_bar_chart(metrics, f"Journey Eval: {eval_name}")
            charts.append((f"Journey Eval: {eval_name}", chart))
        
        for eval_name, metrics in eval_results["aiScriptEvals"].items():
            if chart_type == 'radar':
                chart = create_radar_chart(metrics, f"AI Script Eval: {eval_name}")
            else:
                chart = create_bar_chart(metrics, f"AI Script Eval: {eval_name}")
            charts.append((f"AI Script Eval: {eval_name}", chart))
    
    elif isinstance(eval_results, dict):
        # Query 2 with group_by_plots=True or Query 3 with group_by_plots=True
        for plot_name, plot_data in eval_results.items():
            
            if isinstance(plot_data, dict) and "journeyEvals" in plot_data and "aiScriptEvals" in plot_data:
                # Query 2 result grouped by plots
                for eval_name, metrics in plot_data["journeyEvals"].items():
                    if chart_type == 'radar':
                        chart = create_radar_chart(metrics, f"{plot_name} - Journey Eval: {eval_name}")
                    else:
                        chart = create_bar_chart(metrics, f"{plot_name} - Journey Eval: {eval_name}")
                    charts.append((f"{plot_name} - Journey Eval: {eval_name}", chart))
                
                for eval_name, metrics in plot_data["aiScriptEvals"].items():
                    if chart_type == 'radar':
                        chart = create_radar_chart(metrics, f"{plot_name} - AI Script Eval: {eval_name}")
                    else:
                        chart = create_bar_chart(metrics, f"{plot_name} - AI Script Eval: {eval_name}")
                    charts.append((f"{plot_name} - AI Script Eval: {eval_name}", chart))
            
            elif isinstance(plot_data, dict):
                # Query 3 result grouped by plots
                for eval_name, metrics in plot_data.items():
                    if chart_type == 'radar':
                        chart = create_radar_chart(metrics, f"{plot_name} - Eval: {eval_name}")
                    else:
                        chart = create_bar_chart(metrics, f"{plot_name} - Eval: {eval_name}")
                    charts.append((f"{plot_name} - Eval: {eval_name}", chart))
    
    return charts