Update models.py

models.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict
 from pydantic import BaseModel, Field
 
 class Location(BaseModel):
@@ -7,45 +7,40 @@ class Location(BaseModel):
   lon: float
   relevance: str = Field(description="Why this location matters in the historical context")
 
-class ZoomRange(BaseModel):
+class TemporalZoom(BaseModel):
+  level: str = Field(description="Level of temporal detail (century/decade/year)")
   range: List[int] = Field(description="Start and end years for the time period")
   explanation: str = Field(description="Historical significance of this time period")
 
-class StyleConfig(BaseModel):
+class GeographicalZoom(BaseModel):
+  level: str = Field(description="Level of geographical detail (continent/country/city)")
+  locations: List[Location]
+  explanation: str = Field(description="Significance of these locations")
+
+class StyleZoom(BaseModel):
   level: str = Field(description="Level of style detail (movement/sub_movement/specific)")
   options: List[str] = Field(description="Available artistic styles for selection")
   explanation: str = Field(description="Historical context for the artistic styles")
 
-class GeographicalZoom(BaseModel):
-  locations: List[Location]
-  explanation: str = Field(description="Significance of these locations")
+class ZoomConfiguration(BaseModel):
+  in_zoom: Optional[Union[TemporalZoom, GeographicalZoom, StyleZoom]]
+  out_zoom: Optional[Union[TemporalZoom, GeographicalZoom, StyleZoom]]
 
 class AxisImpact(BaseModel):
-  temporal: Optional[str] = Field(description="How time period affects this axis")
-  geographical: Optional[str] = Field(description="How location affects this axis")
-  style: Optional[str] = Field(description="How artistic style affects this axis")
-
-class TemporalConfig(BaseModel):
-  component: str = Field(default="st.slider")
-  current_zoom: ZoomRange
-  available_zooms: dict[str, ZoomRange] = Field(description="Available zoom levels (in/out)")
-  impacted_by: AxisImpact
-
-class GeographicalConfig(BaseModel):
-  component: str = Field(default="st.map")
-  current_zoom: GeographicalZoom
-  available_zooms: dict[str, GeographicalZoom] = Field(description="Available zoom levels (in/out)")
-  impacted_by: AxisImpact
-
-class StyleConfig(BaseModel):
-  component: str = Field(default="st.multiselect")
-  current_zoom: StyleConfig
-  
+  temporal: Optional[str] = Field(default=None, description="How time period affects this axis")
+  geographical: Optional[str] = Field(default=None, description="How location affects this axis")
+  style: Optional[str] = Field(default=None, description="How artistic style affects this axis")
+
+class AxisConfiguration(BaseModel):
+  component: str
+  current_zoom: Union[TemporalZoom, GeographicalZoom, StyleZoom]
+  available_zooms: Optional[ZoomConfiguration] = None
+  impacted_by: Optional[AxisImpact] = None
+
 class Analysis(BaseModel):
   query_focus: str = Field(description="Main subject of the query")
   historical_context: str = Field(description="Brief historical context")
 
 class ArtHistoryResponse(BaseModel):
   analysis: Analysis
-  axis_configurations: dict[str, Union[TemporalConfig, GeographicalConfig, StyleConfig]] = \
-      Field(description="Configuration for each axis (temporal, geographical, style)")
+  axis_configurations: Dict[str, AxisConfiguration]