Update art_explorer.py

art_explorer.py

@@ -1,19 +1,27 @@
 from typing import Dict, Any, List, Optional
 from openai import OpenAI
 import json
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+from prompts import SYSTEM_PROMPT, format_exploration_prompt, DEFAULT_RESPONSE
 
 class ExplorationNode(BaseModel):
-  title: str
-  description: str
-  connections: List[str]
-  depth: int
+  id: Optional[str] = None
+  node_id: Optional[str] = None
+  title: Optional[str] = Field(None, alias='content')
+  description: str = ""
+  connections: List[Dict[str, Any]] = Field(default_factory=list)
+  depth: Optional[int] = 0
+  content: Optional[str] = None
 
 class ExplorationPath(BaseModel):
   nodes: List[ExplorationNode]
   query: str
   domain: Optional[str] = None
 
+  class Config:
+      populate_by_name = True
+      arbitrary_types_allowed = True
+
 class ExplorationPathGenerator:
   def __init__(self, api_key: str):
       self.client = OpenAI(
@@ -23,62 +31,44 @@ class ExplorationPathGenerator:
       
   def generate_exploration_path(
       self,
-      query: str,  # Changed from user_query to query to match app.py
+      query: str,
       selected_path: List[Dict[str, Any]] = None,
       exploration_parameters: Dict[str, Any] = None
   ) -> Dict[str, Any]:
-      """
-      Generate an exploration path based on the query and parameters
-      
-      Args:
-          query: The user's exploration query
-          selected_path: Previously selected exploration path nodes
-          exploration_parameters: Additional parameters for exploration
-      
-      Returns:
-          Dict containing the exploration path and metadata
-      """
+      """Generate an exploration path based on the query and parameters"""
       try:
-          # Set defaults
           selected_path = selected_path or []
           exploration_parameters = exploration_parameters or {}
           
-          # Extract parameters
-          depth = exploration_parameters.get("depth", 5)
-          domain = exploration_parameters.get("domain")
-          
-          # Prepare system prompt
-          system_prompt = """You are an expert art historian assistant. Generate structured exploration paths 
-          that help users discover connections and insights in art history. Format your response as a JSON object 
-          with 'nodes' containing exploration points."""
-
-          # Prepare user message
-          user_message = {
-              "query": query,
-              "depth": depth,
-              "domain": domain,
-              "previous_explorations": selected_path
-          }
+          # Format the exploration prompt using the helper function
+          formatted_prompt = format_exploration_prompt(
+              user_query=query,
+              selected_path=selected_path,
+              exploration_parameters=exploration_parameters
+          )
 
-          # Make API call
           response = self.client.chat.completions.create(
               model="mixtral-8x7b-32768",
               messages=[
-                  {"role": "system", "content": system_prompt},
-                  {"role": "user", "content": json.dumps(user_message)}
+                  {"role": "system", "content": SYSTEM_PROMPT},
+                  {"role": "user", "content": formatted_prompt}
               ],
               temperature=0.7,
               max_tokens=2000
           )
 
-          # Parse and validate response
-          result = json.loads(response.choices[0].message.content)
-          
-          # Validate with Pydantic
+          # Parse the response
+          try:
+              result = json.loads(response.choices[0].message.content)
+          except json.JSONDecodeError:
+              print("Failed to parse API response as JSON, using default response")
+              result = DEFAULT_RESPONSE
+
+          # Convert to ExplorationPath model
           exploration_path = ExplorationPath(
-              nodes=[ExplorationNode(**node) for node in result["nodes"]],
+              nodes=[ExplorationNode(**node) for node in result.get("nodes", [])],
               query=query,
-              domain=domain
+              domain=exploration_parameters.get("domain")
           )
           
           return exploration_path.model_dump()
@@ -88,5 +78,6 @@ class ExplorationPathGenerator:
           return {
               "error": str(e),
               "status": "failed",
-              "message": "Failed to generate exploration path"
+              "message": "Failed to generate exploration path",
+              "default_response": DEFAULT_RESPONSE
           }