src/visualization_schema.py

from typing import List, Dict, Any, Optional, Union, Literal
from pydantic import BaseModel, Field
from enum import Enum

class ChartType(str, Enum):
    BAR = "bar"
    LINE = "line"
    PIE = "pie"
    SCATTER = "scatter"
    HISTOGRAM = "histogram"
    AREA = "area"
    BOX = "box"
    HEATMAP = "heatmap"
    TIMELINE = "timeline"

class AxisConfig(BaseModel):
    type: str = Field(description="Axis type: 'category', 'value', 'time', or 'log'")
    name: Optional[str] = Field(None, description="Axis name/label")
    data: Optional[List[Union[str, int, float]]] = Field(None, description="Category data for category axis")
    min: Optional[Union[int, float]] = Field(None, description="Minimum value")
    max: Optional[Union[int, float]] = Field(None, description="Maximum value")

class SeriesData(BaseModel):
    name: str = Field(description="Series name")
    type: ChartType = Field(description="Chart type for this series")
    data: List[Union[Dict[str, Any], List[Union[str, int, float]], str, int, float]] = Field(description="Data points for the series")
    color: Optional[str] = Field(None, description="Series color")
    smooth: Optional[bool] = Field(False, description="Whether to smooth the line (for line charts)")

class EChartsConfig(BaseModel):
    title: Optional[Dict[str, Any]] = Field(None, description="Chart title configuration")
    tooltip: Optional[Dict[str, Any]] = Field(None, description="Tooltip configuration")
    legend: Optional[Dict[str, Any]] = Field(None, description="Legend configuration")
    grid: Optional[Dict[str, Any]] = Field(None, description="Grid configuration")
    xAxis: Optional[Union[AxisConfig, List[AxisConfig]]] = Field(None, description="X-axis configuration")
    yAxis: Optional[Union[AxisConfig, List[AxisConfig]]] = Field(None, description="Y-axis configuration")
    series: List[SeriesData] = Field(description="Chart series data")
    backgroundColor: Optional[str] = Field("transparent", description="Background color")
    textStyle: Optional[Dict[str, Any]] = Field(None, description="Global text style")

class VisualizationResult(BaseModel):
    """Structured output for data visualization generation"""
    success: bool = Field(description="Whether the visualization was generated successfully")
    chart_title: str = Field(description="Title of the generated chart")
    chart_type: str = Field(description="Type of chart generated")
    data: Dict[str, Any] = Field(description="Simple chart data with labels and values")
    data_summary: Dict[str, Any] = Field(description="Summary of the processed data")
    reasoning: str = Field(description="Explanation of the visualization choices made")
    error: Optional[str] = Field(default=None, description="Error message if generation failed")
    usage: Optional[Dict[str, int]] = Field(default=None, description="Token usage statistics (input, output, total)")

class FieldInfo(BaseModel):
    """Information about form fields"""
    id: str = Field(description="Field ID")
    label: str = Field(description="Field label")
    field_type: str = Field(description="Field type (select, radio, checkbox, number, etc.)")
    options: Optional[List[str]] = Field(None, description="Available options for choice fields")

class VisualizationRequest(BaseModel):
    """Request for data visualization generation"""
    query: str = Field(description="Natural language query for chart generation")
    csv_data: str = Field(description="CSV data as string")
    field_info: List[FieldInfo] = Field(description="Information about form fields")
    form_title: str = Field(description="Title of the form")

class VisualizationResponse(BaseModel):
    """Response from visualization generation"""
    success: bool
    result: Optional[VisualizationResult] = None
    error: Optional[str] = None
    processing_time: float = 0.0
    usage: Optional[Dict[str, int]] = None