models.py

from typing import List, Optional, Union
from pydantic import BaseModel, Field
from instructor import Field as InstructorField

class Location(BaseModel):
  name: str
  lat: float
  lon: float
  relevance: str = Field(description="Why this location matters in the historical context")

class ZoomRange(BaseModel):
  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):
  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 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
  
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)")