Update app.py

app.py

@@ -2,9 +2,15 @@ import os
 import openai
 import gradio as gr
 import json
+from typing import Optional, Dict, Any, List
 import plotly.graph_objects as go
+import instructor
+from models import ArtHistoryResponse, Location
 from variables import CONTEXTUAL_ZOOM_PROMPT, CONTEXTUAL_ZOOM_default_response
 
+# Enable instructor
+instructor.patch()
+
 class ArtExplorer:
   def __init__(self):
       self.client = openai.OpenAI(
@@ -16,17 +22,22 @@ class ArtExplorer:
           "selections": {}
       }
 
-  def create_map(self, locations):
+  def create_map(self, locations: List[Location]) -> go.Figure:
       """Create a Plotly map figure from location data"""
       if not locations:
-          locations = [{"name": "Paris", "lat": 48.8566, "lon": 2.3522}]
+          locations = [Location(
+              name="Paris",
+              lat=48.8566,
+              lon=2.3522,
+              relevance="Default location"
+          )]
 
       fig = go.Figure(go.Scattermapbox(
-          lat=[loc.get('lat') for loc in locations],
-          lon=[loc.get('lon') for loc in locations],
+          lat=[loc.lat for loc in locations],
+          lon=[loc.lon for loc in locations],
           mode='markers',
           marker=go.scattermapbox.Marker(size=10),
-          text=[loc.get('name') for loc in locations]
+          text=[loc.name for loc in locations]
       ))
 
       fig.update_layout(
@@ -39,29 +50,8 @@ class ArtExplorer:
       )
       return fig
 
-  def clean_json_string(self, json_str: str) -> str:
-      """Clean and prepare JSON string for parsing"""
-      # Remove any leading/trailing whitespace and newlines
-      json_str = json_str.strip()
-      
-      # Remove any BOM or special characters at the start
-      if json_str.startswith('\ufeff'):
-          json_str = json_str[1:]
-          
-      # Ensure it starts with {
-      if not json_str.startswith('{'):
-          start_idx = json_str.find('{')
-          if start_idx != -1:
-              json_str = json_str[start_idx:]
-              
-      # Ensure it ends with }
-      end_idx = json_str.rfind('}')
-      if end_idx != -1:
-          json_str = json_str[:end_idx+1]
-          
-      return json_str
-
-  def get_llm_response(self, query: str, zoom_context: dict = None) -> dict:
+  def get_llm_response(self, query: str, zoom_context: Optional[Dict[str, Any]] = None) -> ArtHistoryResponse:
+      """Get response from LLM with proper validation"""
       try:
           print("\n=== Starting LLM Request ===")
           print(f"Input query: {query}")
@@ -87,39 +77,17 @@ class ArtExplorer:
               )}
           ]
 
-          print("\nPrepared messages for LLM:")
-          print(json.dumps(messages, indent=2))
-
           print("\nSending request to LLM...")
           response = self.client.chat.completions.create(
               model="mixtral-8x7b-32768",
               messages=messages,
               temperature=0.1,
-              max_tokens=2048
+              max_tokens=2048,
+              response_model=ArtHistoryResponse
           )
-          print("\nReceived raw response from LLM:")
-          print(response)
-          
-          # Get the response content and clean it
-          content = response.choices[0].message.content
-          print("\nExtracted content from response:")
-          print(content)
-          
-          # Clean the JSON string
-          cleaned_content = self.clean_json_string(content)
-          print("\nCleaned JSON string:")
-          print(cleaned_content)
-          
-          try:
-              result = json.loads(cleaned_content)
-              print("\nSuccessfully parsed JSON:")
-              print(json.dumps(result, indent=2))
-              return result
-          except json.JSONDecodeError as e:
-              print(f"\nJSON parsing error: {str(e)}")
-              print(f"Failed content: {cleaned_content}")
-              return self.get_default_response()
-              
+          print("\nReceived validated response from LLM")
+          return response
+
       except Exception as e:
           print(f"\nError in LLM response: {str(e)}")
           print(f"Full error details: {e.__class__.__name__}")
@@ -127,10 +95,12 @@ class ArtExplorer:
           print(traceback.format_exc())
           return self.get_default_response()
 
-  def get_default_response(self):
-      return CONTEXTUAL_ZOOM_default_response
+  def get_default_response(self) -> ArtHistoryResponse:
+      """Return default response when LLM fails"""
+      return ArtHistoryResponse(**CONTEXTUAL_ZOOM_default_response)
 
-  def create_interface(self):
+  def create_interface(self) -> gr.Blocks:
+      """Create the Gradio interface"""
       with gr.Blocks() as demo:
           gr.Markdown("# Art History Explorer")
           
@@ -160,6 +130,7 @@ class ArtExplorer:
                   geo_zoom = gr.Button("🔍 Zoom Geography")
 
           with gr.Row():
+              # Style axis
               style_select = gr.Dropdown(
                   choices=["Classical", "Modern"],
                   multiselect=True,
@@ -169,48 +140,50 @@ class ArtExplorer:
               style_explanation = gr.Markdown()
               style_zoom = gr.Button("🔍 Zoom Styles")
 
-          def initial_search(query):
+          def initial_search(query: str) -> tuple:
               """Handle the initial search query"""
-              config = self.get_llm_response(query)
-              temporal_config = config["axis_configurations"]["temporal"]["current_zoom"]
-              geographical_config = config["axis_configurations"]["geographical"]["current_zoom"]
-              style_config = config["axis_configurations"]["style"]["current_zoom"]
+              response = self.get_llm_response(query)
+              
+              temporal_config = response.axis_configurations["temporal"].current_zoom
+              geographical_config = response.axis_configurations["geographical"].current_zoom
+              style_config = response.axis_configurations["style"].current_zoom
               
-              map_fig = self.create_map(geographical_config["locations"])
+              map_fig = self.create_map(geographical_config.locations)
               
               return (
-                  temporal_config["range"],
+                  temporal_config.range,
                   map_fig,
-                  style_config["options"],
-                  temporal_config["explanation"],
-                  geographical_config.get("explanation", ""),
-                  style_config["explanation"]
+                  style_config.options,
+                  temporal_config.explanation,
+                  geographical_config.explanation,
+                  style_config.explanation
               )
 
-          def zoom_axis(query, axis_name, current_value):
+          def zoom_axis(query: str, axis_name: str, current_value: Any) -> tuple:
               """Handle zoom events for any axis"""
               self.current_state["zoom_level"] += 1
-              config = self.get_llm_response(
+              response = self.get_llm_response(
                   query,
                   zoom_context={axis_name: current_value}
               )
-              axis_config = config["axis_configurations"][axis_name]["current_zoom"]
+              
+              axis_config = response.axis_configurations[axis_name].current_zoom
               
               if axis_name == "temporal":
                   return (
-                      axis_config["range"],
-                      axis_config["explanation"]
+                      axis_config.range,
+                      axis_config.explanation
                   )
               elif axis_name == "geographical":
-                  map_fig = self.create_map(axis_config["locations"])
+                  map_fig = self.create_map(axis_config.locations)
                   return (
                       map_fig,
-                      axis_config.get("explanation", "")
+                      axis_config.explanation
                   )
               else:  # style
                   return (
-                      axis_config["options"],
-                      axis_config["explanation"]
+                      axis_config.options,
+                      axis_config.explanation
                   )
 
           # Connect event handlers
@@ -248,7 +221,8 @@ class ArtExplorer:
           return demo
 
 def main():
-  print("Starting initialization...")
+  """Main entry point"""
+  print("Starting Art History Explorer...")
   explorer = ArtExplorer()
   print("Created ArtExplorer instance")
   demo = explorer.create_interface()