utils/planner.py

# utils/planner.py

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

# ----------------------------
# 🔐 Load Environment & Client
# ----------------------------
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# ----------------------------
# 🧠 Scene Plan Extractor
# ----------------------------
SCENE_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": SCENE_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)

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

# ----------------------------
# 🧠 Prompt Variation Generator
# ----------------------------
def generate_prompt_variations_from_scene(scene_plan: dict, base_prompt: str, n: int = 3) -> list:
    try:
        system_msg = f"""
You are a creative prompt variation generator for an AI image generation system.
Given a base user prompt and its structured scene plan, generate {n} diverse image generation prompts.
Each prompt should:
- Be visually rich and descriptive
- Include stylistic or contextual variation
- Reference the same product and environment
- Stay faithful to the base prompt and extracted plan
Respond ONLY with a JSON array of strings. No explanations.
"""

        response = client.chat.completions.create(
            model="gpt-4o-mini-2024-07-18",
            messages=[
                {"role": "system", "content": system_msg},
                {"role": "user", "content": json.dumps({
                    "base_prompt": base_prompt,
                    "scene_plan": scene_plan
                })}
            ],
            temperature=0.7,
            max_tokens=600
        )

        content = response.choices[0].message.content
        print("🧠 Prompt Variations (Raw):", content)
        return json.loads(content)

    except Exception as e:
        print("❌ generate_prompt_variations_from_scene() Error:", e)
        return [base_prompt]  # fallback to original