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
    group_by: Optional[List[str]] = None
    agg: Optional[List[str]] = None
    pivot: Optional[Dict] = None
    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