app.py

"""
Gradio Web Interface for CoDA.

Provides a user-friendly web UI for the CoDA visualization system,
designed for deployment on Hugging Face Spaces.
"""

import logging
import os
import tempfile
from pathlib import Path
from typing import Optional

import gradio as gr

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def create_coda_interface():
    """Create the Gradio interface for CoDA."""
    
    def process_visualization(
        query: str,
        data_file,
        progress=gr.Progress()
    ) -> tuple[Optional[str], str, str]:
        """
        Process a visualization request.
        
        Args:
            query: Natural language visualization query
            data_file: Uploaded data file
            progress: Gradio progress tracker
            
        Returns:
            Tuple of (image_path, status_message, details)
        """
        if not query.strip():
            return None, "❌ Error", "Please enter a visualization query."
        
        if data_file is None:
            return None, "❌ Error", "Please upload a data file."
        
        try:
            from coda.config import Config
            from coda.orchestrator import CodaOrchestrator
        except ImportError as e:
            return None, "❌ Import Error", f"Failed to import CoDA: {e}"
        
        groq_api_key = os.getenv("GROQ_API_KEY", "")
        if not groq_api_key:
            return (
                None,
                "❌ Configuration Error",
                "GROQ_API_KEY environment variable is not set. "
                "Please add your API key in the Spaces settings."
            )
        
        with tempfile.TemporaryDirectory() as temp_dir:
            data_path = Path(temp_dir) / Path(data_file.name).name
            
            with open(data_file.name, 'rb') as src:
                with open(data_path, 'wb') as dst:
                    dst.write(src.read())
            
            def update_progress(status: str, pct: float):
                progress(pct, desc=status)
            
            try:
                config = Config(
                    groq_api_key=groq_api_key,
                )
                
                orchestrator = CodaOrchestrator(
                    config=config,
                    progress_callback=update_progress,
                )
                
                result = orchestrator.run(
                    query=query,
                    data_paths=[str(data_path)],
                )
                
                if result.success and result.output_file:
                    scores = result.scores or {}
                    details = format_results(result, scores)
                    return result.output_file, "✅ Success", details
                else:
                    error_msg = result.error or "Unknown error occurred"
                    return None, "❌ Failed", f"Visualization failed: {error_msg}"
                    
            except Exception as e:
                logger.exception("Pipeline error")
                return None, "❌ Error", f"An error occurred: {str(e)}"
    
    def format_results(result, scores: dict) -> str:
        """Format the results for display."""
        lines = [
            f"**Iterations:** {result.iterations}",
            "",
            "### Quality Scores",
        ]
        
        if scores:
            for key, value in scores.items():
                emoji = "🟢" if value >= 7 else "🟡" if value >= 5 else "🔴"
                lines.append(f"- {key.title()}: {emoji} {value:.1f}/10")
        
        if result.evaluation:
            if result.evaluation.strengths:
                lines.extend(["", "### Strengths"])
                for s in result.evaluation.strengths[:3]:
                    lines.append(f"- {s}")
            
            if result.evaluation.recommendations:
                lines.extend(["", "### Recommendations"])
                for r in result.evaluation.recommendations[:3]:
                    lines.append(f"- {r}")
        
        return "\n".join(lines)
    
    with gr.Blocks(
        title="CoDA - Collaborative Data Visualization",
        theme=gr.themes.Soft(),
        css="""
        .main-title {
            text-align: center;
            margin-bottom: 1rem;
        }
        .status-box {
            padding: 1rem;
            border-radius: 8px;
            margin-top: 1rem;
        }
        """
    ) as interface:
        gr.Markdown(
            """
            # 🎨 CoDA: Collaborative Data Visualization Agents
            
            Transform your data into beautiful visualizations using natural language.
            Simply upload your data and describe what you want to see!
            """,
            elem_classes=["main-title"]
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                query_input = gr.Textbox(
                    label="Visualization Query",
                    placeholder="e.g., 'Create a line chart showing sales trends over time'",
                    lines=3,
                )
                
                file_input = gr.File(
                    label="Upload Data File",
                    file_types=[".csv", ".json", ".xlsx", ".xls", ".parquet"],
                )
                
                submit_btn = gr.Button(
                    "🚀 Generate Visualization",
                    variant="primary",
                    size="lg",
                )
                
                gr.Markdown(
                    """
                    ### Supported Formats
                    - CSV, JSON, Excel (.xlsx, .xls), Parquet
                    
                    ### Example Queries
                    - "Show me a bar chart of sales by category"
                    - "Create a scatter plot of price vs quantity"
                    - "Plot the distribution of ages as a histogram"
                    """
                )
            
            with gr.Column(scale=2):
                output_image = gr.Image(
                    label="Generated Visualization",
                    type="filepath",
                )
                
                with gr.Row():
                    status_output = gr.Textbox(
                        label="Status",
                        interactive=False,
                    )
                
                details_output = gr.Markdown(
                    label="Details",
                )
        
        gr.Examples(
            examples=[
                ["Create a bar chart showing the top 10 values", None],
                ["Plot a line chart of trends over time", None],
                ["Show a scatter plot with correlation", None],
                ["Create a pie chart of category distribution", None],
            ],
            inputs=[query_input, file_input],
        )
        
        submit_btn.click(
            fn=process_visualization,
            inputs=[query_input, file_input],
            outputs=[output_image, status_output, details_output],
        )
    
    return interface


app = create_coda_interface()

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