Create schemas.py

schemas.py

@@ -0,0 +1,39 @@
+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