schema.py

from typing import List, Optional, Dict
from pydantic import BaseModel, Field, validator

class TaskPlan(BaseModel):
    title: str
    format: str = Field(description="table|list|comparison|map|narrative|chart")
    data_key: Optional[str] = None
    filter: Optional[str] = None
    derive: Optional[List[str]] = None
    joins: Optional[List[Dict]] = None  # [{"right_key":"...", "left_on":"...", "right_on":"...", "how":"left"}]
    group_by: Optional[List[str]] = None
    agg: Optional[List[str]] = None     # ["avg(x)", "median(y)", "p90(z)", "count(*)"]
    pivot: Optional[Dict] = None        # {"index":["a","b"], "columns":"c", "values":"v"}
    sort_by: Optional[str] = None
    sort_dir: Optional[str] = "desc"
    top: Optional[int] = None
    fields: Optional[List[str]] = None
    chart: Optional[str] = None
    encodings: Optional[Dict[str, str]] = None
    number_format: Optional[Dict[str, str]] = None
    title_override: Optional[str] = None

    @validator("format")
    def _fmt(cls, v):
        allowed = {"table","list","comparison","map","narrative","chart"}
        if v not in allowed: raise ValueError(f"format must be one of {allowed}")
        return v

class ScenarioPlan(BaseModel):
    tasks: List[TaskPlan]
    notes: Optional[str] = None