50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import json
|
|
from pydantic import BaseModel, Field
|
|
from APIs.aichat import generate
|
|
|
|
|
|
# 1. Define the output schema as a Pydantic model
|
|
class DishSafetyResult(BaseModel):
|
|
is_safe: bool = Field(
|
|
description="True if no microwave hazards presents."
|
|
)
|
|
warning_message: str = Field(
|
|
description="Short explanation of any hazard found (no more than one sentence), 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 = (
|
|
"Analyze this top-down photo of a dish prepared for microwave cooking. "
|
|
"Inspect the area for metal utensils, aluminum foil, metallic dish patterns, or unvented plastic wraps that could be dangerous for the cooking process or pose a safety risk."
|
|
)
|
|
|
|
# 2. Pass the Pydantic class directly to generate()
|
|
response_raw = generate(
|
|
prompt=prompt,
|
|
images=[image_path],
|
|
output_format=DishSafetyResult,
|
|
should_think=False
|
|
)
|
|
print(f"[Debug] Raw response from AI generator: {response_raw}")
|
|
|
|
# 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
|
|
|
|
raise ValueError(f"Unexpected response type from AI generator: {type(response_raw)}") |