Assistants Deep Dive - creating a png from a csv

We’re learning from the Assistants Deep Dive tutorial but have been unsuccessful in creating an image url or png file displaying the data of an uploaded .csv file. It appears the messages exchanged in the thread were retrieved successfully but the 2nd element of the file_ID is “None.”

content=[ImageFileContentBlock(image_file=ImageFile(file_id=‘file-UNbdcMSePCvuHE02jCyDsZsz’, detail=None), type=‘image_file’)

Would be grateful for advice how to revise our python code:

from openai import OpenAI, OpenAIError
import os
from dotenv import load_dotenv
import logging
import time

# Clear the environment variable
if 'OPENAI_API_KEY' in os.environ:
    del os.environ['OPENAI_API_KEY']

# Load environment variables
load_dotenv()

# Retrieve the OpenAI API key from environment variables
api_key = os.getenv('OPENAI_API_KEY')

if not api_key:
    logging.error("API key not found in environment variables.")
    exit(1)

# Initialize the OpenAI client
client = OpenAI(api_key=api_key)

# Configure logging
logging.basicConfig(level=logging.INFO)

def upload_file(file_path):
    """Upload a file to the OpenAI API and return the file ID."""
    try:
        with open(file_path, "rb") as f:
            file_data = client.files.create(file=f, purpose='assistants')
        logging.info(f"File uploaded successfully with ID: {file_data.id}")
        return file_data.id
    except OpenAIError as e:
        logging.error(f"Failed to upload file: {e}")
        raise

def create_assistant(file_id):
    """Create an assistant using the uploaded file and return the assistant ID."""
    try:
        assistant_data = client.beta.assistants.create(
            name="Data visualizer GC 0802",
            description=(
                "You are great at creating beautiful data visualizations. "
                "You analyze data present in .csv files, understand trends, "
                "and come up with data visualizations relevant to those trends. "
                "You also share a brief text summary of the trends observed."
            ),
            model="gpt-4o-mini",
            tools=[{"type": "code_interpreter"}],
            tool_resources={"code_interpreter": {"file_ids": [file_id]}}
        )
        logging.info(f"Assistant created successfully with ID: {assistant_data.id}")
        return assistant_data.id
    
    except OpenAIError as e:
        logging.error(f"Failed to create assistant: {e}")
        raise

def create_thread(file_id):
    """Create a thread for data visualization and return the thread ID."""
    try:
        thread_data = client.beta.threads.create(
            messages=[
                {
                    "role": "user",
                    "content": "Create a data visualization based on the trends in this file.",
                    "attachments": [{"file_id": file_id, "tools": [{"type": "code_interpreter"}]}]
                }
            ]
        )
        logging.info(f"Thread created successfully with ID: {thread_data.id}")
        return thread_data.id
    except OpenAIError as e:
        logging.error(f"Failed to create thread: {e}")
        raise

def run_thread(thread_id, assistant_id):
    """Run the thread and return the run ID."""
    try:
        run_data = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=assistant_id)
        logging.info(f"Thread run started successfully with ID: {run_data.id}")
        return run_data.id
    except OpenAIError as e:
        logging.error(f"Failed to run thread: {e}")
        raise

def check_run_status(run_id, thread_id):
    """Check the status of the run and return the response."""
    try:
        response = client.beta.threads.runs.retrieve(run_id=run_id, thread_id=thread_id)
        status = response.status
        logging.info(f"Run status: {status}")
        return response
    except OpenAIError as e:
        logging.error(f"Failed to retrieve run status: {e}")
        return None

def get_thread_messages(thread_id):
    """Retrieve all messages from the thread."""
    try:
        messages = client.beta.threads.messages.list(thread_id=thread_id)
        logging.info(f"Messages retrieved: {messages}")
        return messages
    except OpenAIError as e:
        logging.error(f"Failed to retrieve thread messages: {e}")
        return None

def get_run_result(run_id, thread_id):
    """Retrieve the result of the run, including any generated PNG file."""
    try:
        run_result = client.beta.threads.runs.retrieve(run_id=run_id, thread_id=thread_id)
        logging.info(f"Run result: {run_result}")

        messages = get_thread_messages(thread_id)
        if messages:
            for message in messages:
                logging.info(f"Message from {message.role}: {message.content}")
                if 'attachments' in message:
                    for attachment in message['attachments']:
                        if attachment['mime_type'] == 'image/png':
                            logging.info("PNG file found in the message attachments.")
                            return attachment['file_id']

        logging.info("No PNG file found in the thread messages.")
        return None
    except OpenAIError as e:
        logging.error(f"Failed to retrieve run result: {e}")
        return None

def download_image(file_id, download_path):
    """Download the image file from the OpenAI API."""
    try:
        image_data = client.files.content(file_id)
        image_data_bytes = image_data.read()
        
        with open(download_path, "wb") as file:
            file.write(image_data_bytes)
        
        logging.info(f"Image downloaded successfully: {download_path}")
    except OpenAIError as e:
        logging.error(f"Failed to download image: {e}")

def create_folder_if_not_exists(folder_path):
    """Create a folder if it does not exist."""
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)
        logging.info(f"Created folder: {folder_path}")

def main():
    file_path = "revenue-forecast.csv"  # Set the path to your CSV file here
    logging.info("Starting file upload...")
    file_id = upload_file(file_path)

    logging.info("Creating assistant...")
    assistant_id = create_assistant(file_id)

    logging.info("Creating thread...")
    thread_id = create_thread(file_id)

    logging.info("Running thread...")
    run_id = run_thread(thread_id, assistant_id)

    # Check the run status and retrieve results
    while True:
        logging.info("Checking run status...")
        result = check_run_status(run_id, thread_id)
        if result:
            status = result.status
            logging.info(f"Current run status: {status}")
            if status == 'completed':
                break
            elif status == 'failed':
                logging.error("Run failed.")
                break
        else:
            logging.error("Failed to get run status.")
        time.sleep(10)  # Wait for 10 seconds before checking again

    # Retrieve the PNG file
    logging.info("Retrieving run result...")
    file_id = get_run_result(run_id, thread_id)

    if file_id:
        logging.info(f"Download PNG file with ID: {file_id}")
        folder_path = "downloaded_images"
        create_folder_if_not_exists(folder_path)
        download_path = os.path.join(folder_path, "data_visualization.png")
        download_image(file_id, download_path)
    else:
        logging.info("No PNG file was generated.")

if __name__ == "__main__":
    main()

Solved it. Here’s the full code with comments for those who encounter the same issue:


def upload_file(file_path):
    """Upload a file to the OpenAI API and return the file ID."""
    try:
        with open(file_path, "rb") as f:
            file_data = client.files.create(file=f, purpose='assistants')
            logging.info(f"File uploaded successfully with ID: {file_data.id}")
        return file_data.id
    except OpenAIError as e:
        logging.error(f"Failed to upload file: {e}")
        raise

def create_assistant(file_id):
    """Create an assistant using the uploaded file and return the assistant ID."""
    try:
        assistant_data = client.beta.assistants.create(
            name="Data visualizer GC 0804",
            description=(
                "You are great at creating beautiful data visualizations. "
                "You analyze data present in .csv files, understand trends, "
                "and come up with data visualizations relevant to those trends. "
                "You also share a brief text summary of the trends observed."
            ),
            model="gpt-4o-mini",
            tools=[{"type": "code_interpreter"}],
            tool_resources={"code_interpreter": {"file_ids": [file_id]}}
        )
        logging.info(f"Assistant created successfully with ID: {assistant_data.id}")
        return assistant_data.id
    except OpenAIError as e:
        logging.error(f"Failed to create assistant: {e}")
        raise

def create_thread(file_id):
    """Create a thread for data visualization and return the thread ID."""
    try:
        thread_data = client.beta.threads.create(
            messages=[
                {
                    "role": "user",
                    "content": "Create a data visualization based on the trends in this file.",
                    "attachments": [{"file_id": file_id, "tools": [{"type": "code_interpreter"}]}]
                }
            ]
        )
        logging.info(f"Thread created successfully with ID: {thread_data.id}")
        return thread_data.id
    except OpenAIError as e:
        logging.error(f"Failed to create thread: {e}")
        raise

def run_thread(thread_id, assistant_id):
    """Run the thread and return the run ID."""
    try:
        run_data = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=assistant_id)
        logging.info(f"Thread run started successfully with ID: {run_data.id}")
        return run_data.id
    except OpenAIError as e:
        logging.error(f"Failed to run thread: {e}")
        raise

def check_run_status(run_id, thread_id):
    """Check the status of the run and return the response."""
    try:
        response = client.beta.threads.runs.retrieve(run_id=run_id, thread_id=thread_id)
        status = response.status
        logging.info(f"Run status: {status}")
        return response
    except OpenAIError as e:
        logging.error(f"Failed to retrieve run status: {e}")
        return None

def get_thread_messages(thread_id):
    """Retrieve all messages from the thread."""
    try:
        messages = client.beta.threads.messages.list(thread_id=thread_id)
        logging.info(f"Messages retrieved: {messages}")
        return messages
    except OpenAIError as e:
        logging.error(f"Failed to retrieve thread messages: {e}")
        return None

def get_run_result(run_id, thread_id):
    """Retrieve the result of the run, including any generated PNG file."""
    try:
        run_result = client.beta.threads.runs.retrieve(run_id=run_id, thread_id=thread_id)
        logging.info(f"Run result: {run_result}")
        messages = get_thread_messages(thread_id)
        if messages:
            for message in messages.data:
                logging.info(f"Message from {message.role}: {message.content}")
                for content_block in message.content:
                    if content_block.type == 'image_file':
                        image_file = content_block.image_file
                        logging.info("Image file found in the message content blocks.")
                        return image_file.file_id
        logging.info("No PNG file found in the thread messages.")
        return None
    except OpenAIError as e:
        logging.error(f"Failed to retrieve run result: {e}")
        return None

def download_image(file_id, download_path, headers):
    """Download an image file from the OpenAI API and save it to the specified path."""
    try:
        url = f"https://api.openai.com/v1/files/{file_id}/content"
        response = httpx.get(url, headers=headers)
        response.raise_for_status()  # Raise an exception for HTTP errors
        with open(download_path, "wb") as f:
            f.write(response.content)
        logging.info(f"Image downloaded successfully and saved to {download_path}")
    except httpx.HTTPStatusError as e:
        logging.error(f"HTTP error occurred while downloading image: {e.response.status_code} {e.response.text}")
    except Exception as e:
        logging.error(f"An error occurred while downloading image: {e}")

def main():
    file_path = "revenue-forecast.csv"  # Set the path to your CSV file here
    logging.info("Starting file upload...")
    file_id = upload_file(file_path)

    logging.info("Creating assistant...")
    assistant_id = create_assistant(file_id)

    logging.info("Creating thread...")
    thread_id = create_thread(file_id)
    logging.info(f"Generated thread_id: {thread_id}")
    if thread_id is None:
        logging.error("Failed to create thread, exiting...")
        return

    logging.info("Running thread...")
    run_id = run_thread(thread_id, assistant_id)

    # Check the run status and retrieve results
    while True:
        logging.info("Checking run status...")
        result = check_run_status(run_id, thread_id)
        if result:
            status = result.status
            logging.info(f"Current run status: {status}")
            if status == 'completed':
                break
            elif status == 'failed':
                logging.error("Run failed.")
                return
        else:
            logging.error("Failed to get run status.")
        time.sleep(10)  # Wait for 10 seconds before checking again

    # Retrieve the PNG file
    logging.info("Retrieving run result...")
    file_id = get_run_result(run_id, thread_id)
    if file_id:
        logging.info(f"Download PNG file with ID: {file_id}")
        folder_path = "downloaded_images"
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        download_path = os.path.join(folder_path, "data_visualization.png")
        download_image(file_id, download_path, headers={"Authorization": f"Bearer {api_key}"})
    else:
        logging.info("No PNG file was generated.")

if __name__ == "__main__":
    main()