main.py

from openai import OpenAI
import os
import json
from pydantic import BaseModel
from concurrent.futures import ThreadPoolExecutor 
import base64

gpt_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
class ImageAdEssentials(BaseModel):
    phsychologyTriggers: str
    angles: list[str]
    concepts: list[str]

class ImageAdEssentialsOutput(BaseModel):
    output: list[ImageAdEssentials]

class Text(BaseModel):
    textToBeWrittern: str
    color: str
    placement: str
    
class CreativeStrategies(BaseModel):
    phsychologyTrigger: str
    angle: str
    concept: str
    text: Text
    cta: str
    visualDirection: str
    titleIdeas: str
    captionIdeas: str
    bodyIdeas: str

class CreativeStrategiesOutput(BaseModel):
    output: list[CreativeStrategies]

class AdImagePrompt(BaseModel):
    prompt: str

class CopyWriterOutput(BaseModel):
    title: str
    body: str
    description: str

def researcher(target_audience, prompt, niche, product_dec):

    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": """You are the researcher for the e-commerce brand company which does research on trending angles, concepts and psychology trigers based on the user input.
                    The e-commerce brand company name is Amalfa which is a contemporary jewellery brand known in the Indian market for its demi-fine and fashion jewellery collections.
                    Amalfa aims to be a style-forward, expressive brand for today's youth and modern women, blending trend-driven design with accessible pricing.
                    A psychology trigger is an emotional or cognitive stimulus that pushes someone toward action—clicking, signing up, or buying—before logic kicks in.
                    An ad angle is the reason someone should care right now. Same product → different reasons to click → different angles.
                    An ad concept is the creative execution style or storyline you use to deliver an angle.
                    
                    Keeping in mind all this, make sure you provide different angles and concepts we can try based on the phsychology triggers for the image ads for the given input based on e-commerce brand.
                    User will provide you the category on which he needs to run the ads, his requirement, product description and what is target audience."""
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": f"""Following are the inputs:
                    Niche: {niche}
                    Target Audience: {target_audience}
                    User Prompt: {prompt}
                    Product Description: {product_des}
                    
                    Provide the different phsychology triggers, angles and concept based on the given input."""
                }
            ]
        }
    ]
    
    completion = gpt_client.beta.chat.completions.parse(
        model="gpt-4o",
        messages=messages,
        response_format=ImageAdEssentialsOutput,
    )

    response = completion.choices[0].message
    if response.parsed:
        return response.parsed.output
    else:
        return response.refusal

target_aud = "Corporate woman"
user_prompt = "Want to show how they glow up their looks in office"
category = "Ring"
product_des = "The Tigress Claw Knuckle Ring is a bold expression of strength and modern elegance. Inspired by the fierce beauty of a tigress, this multi-finger knuckle ring features sleek claw-shaped gold bands paired with dazzling zircon stone detailing.. Its modern silhouette makes it perfect for party wear, special occasions, and fashion-forward everyday styling."

researcher_output = researcher(target_aud, user_prompt, category, product_des)
print(json.dumps([item.model_dump() for item in researcher_output], indent=2))