Try to get structured output from AI

cloud/safety_checker.py

@@ -1,27 +1,20 @@
 import json
+from pydantic import BaseModel, Field
 from APIs.aichat import generate
 
-DISH_SAFETY_SCHEMA = {
-    "type": "object",
-    "properties": {
-        "is_safe": {
-            "type": "boolean",
-            "description": "True if no microwave hazards (metal, foil, sealed packaging) are present."
-        },
-        "warning_message": {
-            "type": "string",
-            "description": "Explanation of any hazard found, or an empty string if safe."
-        },
-        "detected_hazards": {
-            "type": "array",
-            "items": {
-                "type": "string"
-            },
-            "description": "List of specific hazard items detected."
-        }
-    },
-    "required": ["is_safe", "warning_message", "detected_hazards"]
-}
+
+# 1. Define the output schema as a Pydantic model
+class DishSafetyResult(BaseModel):
+    is_safe: bool = Field(
+        description="True if no microwave hazards (metal, foil, sealed packaging) are present."
+    )
+    warning_message: str = Field(
+        description="Explanation of any hazard found, or an empty string if safe."
+    )
+    detected_hazards: list[str] = Field(
+        description="List of specific hazard items detected."
+    )
+
 
 def check_dish_safety(image_path: str) -> dict:
     prompt = (
@@ -29,15 +22,28 @@ def check_dish_safety(image_path: str) -> dict:
         "Inspect the area for metal utensils, aluminum foil, metallic dish patterns, or unvented plastic wraps."
     )
 
-    # Pass the schema directly to the generation call
+    # 2. Pass the Pydantic class directly to generate()
     response_raw = generate(
         prompt=prompt,
         images=[image_path],
-        output_format=json.dumps(DISH_SAFETY_SCHEMA)
+        output_format=DishSafetyResult,
+        should_think=False
     )
     print(f"[Debug] Raw response from AI generator: {response_raw}")
 
-    # 1. If response_raw is already a dict, return it directly
+    # 3. Handle response parsing
+    if isinstance(response_raw, str):
+        # Parse and validate the JSON string into the Pydantic model, then return as a dict
+        try:
+            validated_result = DishSafetyResult.model_validate_json(response_raw)
+            return validated_result.model_dump()
+        except Exception:
+            # Fallback to standard json.loads if raw parsing is needed
+            return json.loads(response_raw)
+
+    if isinstance(response_raw, DishSafetyResult):
+        return response_raw.model_dump()
+
     if isinstance(response_raw, dict):
         return response_raw