utils/planner.py

# utils/planner.py

import openai
import os
from openai import OpenAI
from dotenv import load_dotenv
import json
load_dotenv()

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

SYSTEM_INSTRUCTIONS = """
You are a scene planning assistant for an AI image generation system.
Your job is to take the user's prompt and return a structured JSON with:
- scene (environment, setting)
- subject (main actor)
- objects (main product or items)
- layout (foreground/background elements and their placement)
- rules (validation rules to ensure visual correctness)
Respond ONLY in raw JSON format. Do NOT include explanations.
"""

def extract_scene_plan(prompt: str) -> dict:
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini-2024-07-18",
            messages=[
                {"role": "system", "content": SYSTEM_INSTRUCTIONS},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=500
        )
        json_output = response.choices[0].message.content
        print("🧠 Scene Plan (Raw):", json_output)
        return json.loads(json_output) # Be cautious: Use `json.loads()` if possible

    except Exception as e:
        print("❌ extract_scene_plan() Error:", e)
        return {
            "scene": None,
            "subject": None,
            "objects": [],
            "layout": {},
            "rules": {}
        }