schemas.py

from pydantic import BaseModel, Field, ConfigDict
from typing import List, Optional, Literal, Dict

FormatType = Literal["table", "list", "comparison", "map", "narrative", "chart"]
SortDir = Literal["asc", "desc"]
JoinHow = Literal["left", "inner", "right", "outer"]
ChartType = Literal["bar", "line", "area", "point", "tick", "rule"]

class JoinSpec(BaseModel):
    right_key: str
    left_on: str
    right_on: str
    how: JoinHow = "left"

class TaskSpec(BaseModel):
    title: str
    data_key: Optional[str] = None
    format: FormatType = "table"
    filter: Optional[str] = None
    derive: Optional[List[str]] = None
    group_by: Optional[List[str]] = None
    agg: Optional[List[str]] = None
    pivot: Optional[Dict[str, str]] = None
    join: Optional[List[JoinSpec]] = None
    sort_by: Optional[str] = None
    sort_dir: SortDir = "desc"
    top: Optional[int] = None
    fields: Optional[List[str]] = None
    chart: Optional[ChartType] = None
    x: Optional[str] = None
    y: Optional[str] = None
    color: Optional[str] = None
    column: Optional[str] = None

class ScenarioPlan(BaseModel):
    model_config = ConfigDict(extra="forbid")
    tasks: List[TaskSpec] = Field(default_factory=list)
    narrative_required: bool = True
    notes: Optional[str] = None