update model

app.py

@@ -36,7 +36,7 @@ async def onClick(vacationType,familyType,duration,purpose,interests):
     "Visual Arts: Imagery from H.R. Giger (Alien) or Moebius (The Incal).",
     "Graphic Novels: Saga, Transmetropolitan."
     ),
-    #result_type = customModel.Planet
+    result_type = customModel.TravelAgency
     )
     result_1 = agent.run_sync(user_prompt="for space travel which planet is the best")
     result = result_1

planet_model.py

@@ -1,16 +1,44 @@
 from pydantic import BaseModel
 
+from dataclasses import dataclass, field
+from typing import List, Dict
+
+@dataclass
 class Planet(BaseModel):
-    def __init__(self, name, atmosphere, species, signature_dish, activities):
-        self.name = name
-        self.atmosphere = atmosphere
-        self.species = species
-        self.signature_dish = signature_dish
-        self.activities = activities
-
-    def describe(self):
-        return f"Planet: {self.name}\n" \
-               f"Atmosphere: {self.atmosphere}\n" \
-               f"Species: {self.species}\n" \
-               f"Signature Dish: {self.signature_dish}\n" \
-               f"Activities: {', '.join(self.activities)}"
+    name: str
+    description: str
+    accommodation: str
+    activities: List[str]
+    photography: str
+
+@dataclass
+class Category(BaseModel):
+    budget_range: str
+    planets: List[Planet] = field(default_factory=list)
+
+@dataclass
+class TravelAgency(BaseModel):
+    name: str
+    categories: Dict[str, Category] = field(default_factory=dict)
+
+    def add_planet(self, category_name: str, budget_range: str, planet: Planet):
+        # Ensure category exists, create if not
+        if category_name not in self.categories:
+            self.categories[category_name] = Category(budget_range=budget_range)
+        self.categories[category_name].planets.append(planet)
+
+
+# class Planet(BaseModel):
+#     def __init__(self, name, atmosphere, species, signature_dish, activities):
+#         self.name = name
+#         self.atmosphere = atmosphere
+#         self.species = species
+#         self.signature_dish = signature_dish
+#         self.activities = activities
+
+#     def describe(self):
+#         return f"Planet: {self.name}\n" \
+#                f"Atmosphere: {self.atmosphere}\n" \
+#                f"Species: {self.species}\n" \
+#                f"Signature Dish: {self.signature_dish}\n" \
+#                f"Activities: {', '.join(self.activities)}"