Update models.py

models.py

@@ -1,72 +1,52 @@
-from typing import List, Optional, Literal
+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
+  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")
-    explanation: str
+  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
-    options: List[str] = Field(..., description="Available style options")
-    explanation: str
-
-class TemporalZoom(BaseModel):
-    level: Literal["century", "decade", "year"]
-    range: List[int] = Field(..., description="Start and end years")
-    explanation: str
+  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):
-    level: Literal["continent", "country", "city"]
-    locations: List[Location]
-    explanation: str = ""
-
-class StyleZoom(BaseModel):
-    level: Literal["movement", "sub_movement", "specific"]
-    options: List[str]
-    explanation: str
+  locations: List[Location]
+  explanation: str = Field(description="Significance of these locations")
 
 class AxisImpact(BaseModel):
-    geographical: Optional[str] = None
-    temporal: Optional[str] = None
-    style: Optional[str] = None
+  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: Literal["st.slider"]
-    current_zoom: TemporalZoom
-    available_zooms: dict[Literal["in", "out"], ZoomRange]
-    impacted_by: AxisImpact
+  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: Literal["st.map"]
-    current_zoom: GeographicalZoom
-    available_zooms: dict[Literal["in", "out"], dict[Literal["locations", "explanation"], str | List[Location]]]
-    impacted_by: AxisImpact
+  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: Literal["st.multiselect", "st.selectbox"]
-    current_zoom: StyleZoom
-
+  component: str = Field(default="st.multiselect")
+  current_zoom: StyleConfig
+  
 class Analysis(BaseModel):
-    query_focus: str
-    historical_context: str
-
-class ArtResponse(BaseModel):
-    """Art history exploration response with contextual zoom configurations"""
-    analysis: Analysis
-    axis_configurations: dict[
-        Literal["temporal", "geographical", "style"],
-        TemporalConfig | GeographicalConfig | StyleConfig
-    ]
+  query_focus: str = Field(description="Main subject of the query")
+  historical_context: str = Field(description="Brief historical context")
 
-    @classmethod
-    def from_llm(cls, query: str, current_zoom_states: dict) -> "ArtResponse":
-        """Factory method to create response from LLM output"""
-        # This would be implemented with instructor
-        pass
+class ArtHistoryResponse(BaseModel):
+  analysis: Analysis
+  axis_configurations: dict[str, Union[TemporalConfig, GeographicalConfig, StyleConfig]] = \
+      Field(description="Configuration for each axis (temporal, geographical, style)")