sketch_design/sk_render_ai.py

"""
This module provides the DesignAssistant class to interact with OpenAI's API
for creating and managing threads and messages related to design assistance.
"""
import os
from dotenv import load_dotenv
from openai import OpenAI

# Load environment variables from .env file
load_dotenv()

class DesignAssistant:
    def __init__(self):
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        self.assistant = self.client.beta.assistants.create(
            name="Extraction Assistant",
            # instructions="You are a professional Bedroom renovator engine, you get an information and you respond with a dict each time with the following keys: 'design_style', 'room_type'. Respond with the dict only and no need to encapsulate between quotes.",
            instructions= "You are a professional dictionary creator.\n\nFollowing are some important rules;\nYou get some information and you respond with a dict each time with the following keys: 'design_style', 'room_type' only and always.  Respond with the dict only and no need to encapsulate between quotes.\n\nThe room_type and design_style has to be choosen from the following, incase of confusion choose the closest one;\nroom_types = [\"Living Room\", \"Bedroom\", \"Bathroom\", \"Kitchen\", \"Dining Room\", \"Attic\", \"Family Room\", \"Kids Room\", \"Study Room\", \"Home Office\", \"Meeting Room\", \"Coffee Shop\", \"Office\", \"Coworking Space\", \"Working Space\", \"Library\", \"Science Laboratory\", \"Photo Studio\", \"Exhibition Space\", \"Art Studio\", \"Multimedia Room\", \"Medical Exam Room\", \"Reception Area\", \"Formal Dining Room\", \"Home Theater\", \"Gaming Room\", \"Fitness Gym\", \"Restaurant\", \"Hotel Lobby\", \"Hotel Room\", \"Wine Cellar\", \"Home Spa\", \"Dressing Room\", \"Home Bar\", \"Music Room\", \"Workshop\", \"Craft Room\", \"Yoga Studio\", \"Auditorium\", \"Art Studio\", \"Balcony\", \"Rooftop Terrace\", \"Mudroom\", \"Guest Bedroom\", \"Wedding Room\", \"Outdoor Kitchen\", \"Utility Room\", \"Laundry Room\", \"Pet Room\", \"Walk-in Closet\", \"Lounge\", \"Reading Nook\", \"Foyer\", \"Open Kitchen Living Room\", \"Prayer Room\", \"Conservatory\", \"Playroom\", \"Man Cave\", \"She Shed\", \"Greenhouse\", \"Sauna\", \"Nursery\", \"Home Gym\"]\n\ndesign_styles = [\"No Style\", \"Eclectic\", \"Modern\", \"Contemporary\", \"Transitional\", \"Scandinavian\", \"Mediterranean\", \"Ikea\", \"Industrial\", \"Kids Room\", \"Shabby Chic\", \"Coastal\", \"Bauhaus\", \"Bohemian\", \"Traditional\", \"Rustic\", \"Minimalism\", \"Japandi\", \"Japanese Design\", \"Modern Arabic\", \"Traditional Arabic\", \"Bali\", \"Tropical\", \"Asian Decor\", \"Zen\", \"Hollywood Regency\", \"Hollywood Glam\", \"Minimalist\", \"Christmas\", \"Futuristic\", \"Luxurious\", \"Midcentury Modern\", \"Biophilic\", \"Cottage Core\", \"French Country\", \"Art Deco\", \"Art Nouveau\", \"South Western\", \"Modern Farm House\", \"Moroccan\", \"Gothic\", \"Victorian\", \"Steampunk\", \"Urban Modern\", \"Desert Modernism\", \"Colonial\", \"Brutalist\", \"Nordic Noir\", \"Postmodern\", \"Psychedelic\", \"Cosmic Chic\", \"Mexican Hacienda\", \"Coastal Modern\", \"Eco Friendly\", \"Pop Art\", \"Vintage Glam\", \"Candy Land\", \"Airbnb\", \"Glam Rock\", \"Barbie\", \"Doodle\", \"Sketch\", \"Maximalist\", \"Professional\", \"Halloween\", \"Retro\", \"Romantic\", \"Safari\", \"Tuscan\", \"Nautical\", \"Craftsman\", \"Farmhouse Chic\", \"Prairie\", \"Cubism\", \"Quiet Luxury\"]\n\nThe information can be in italian aswell, so choose the closest values from the above lists for the dict's values as the response will be always in english.",
            tools=[],
            model="gpt-4o",
        )

    def create_thread(self):
        """
        Creates a new thread for communication.
        
        Returns:
            Thread object containing the details of the created thread.
        """
        return self.client.beta.threads.create()

    def create_message(self, thread_id, content):
        """
        Sends a message in a thread.

        Args:
            thread_id (str): The ID of the thread.
            content (str): The content of the message.

        Returns:
            Message object containing the details of the created message.
        """
        return self.client.beta.threads.messages.create(
            thread_id=thread_id,
            role="user",
            content=content
        )

    def run_and_poll(self, thread_id):
        """
        Runs a process and polls for its completion.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            Run object containing the status and details of the run.
        """
        return self.client.beta.threads.runs.create_and_poll(
            thread_id=thread_id,
            assistant_id=self.assistant.id,
            instructions="Respond with a dict each time with the following keys: 'design_style', 'room_type'"
        )

    def get_messages(self, thread_id):
        """
        Retrieves messages from a thread.

        Args:
            thread_id (str): The ID of the thread.

        Returns:
            List of Message objects containing the messages in the thread.
        """
        return self.client.beta.threads.messages.list(
            thread_id=thread_id
        )

if __name__ == "__main__":
    assistant = DesignAssistant()
    thread = assistant.create_thread()
    message = assistant.create_message(thread.id, "I want to generate a design for a scandanavian bedroom")
    run = assistant.run_and_poll(thread.id)

    if run.status == 'completed': 
        messages = assistant.get_messages(thread.id)
        print(messages.data[0].content[0].text.value)
    else:
        print(run.status)