o1 model Python quickstart using the OpenAI API

%env OPENAI_API_KEY=KEY

!pip install openai weave

import os
import weave
from openai import OpenAI
import re

# Initialize Weave and W&B
weave.init('o1-preview-setup')

# Initialize OpenAI client
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

o1_assistant_prompt = "You are a " + input ("Who should I be, as I answer your prompt?") 
o1_user_prompt = input ("What prompt do you want me to do?") 
o1_prompt = o1_assistant_prompt, o1_user_prompt
print(o1_prompt)

# Function to parse steps_text into a list of steps
def parse_steps(steps_text):
    steps = re.split(r'\n\s*\d+\.\s+|\n\s*-\s+|\n\s*\*\s+', steps_text.strip())
    steps = [step.strip() for step in steps if step.strip()]
    return steps

def generate_steps(o1_assistant_prompt: str, o1_user_prompt: str) -> list:
    steps_prompt = f"You want to {o1_user_prompt}. Before generating the strategy, please outline the reasoning steps you would take to develop the strategy. The steps should be given as a simple list, avoiding commentary or unnecessary detail. Make your best guess at what you are being asked to create a list of reasoning steps for. List the steps numerically."
    
    steps_messages = [
        {"role": "assistant", "content": o1_assistant_prompt},
        {"role": "user", "content": steps_prompt}
    ]
    
    # API call to get steps
    steps_response = client.chat.completions.create(
        model="o1-preview",
        messages=steps_messages
    )
    
    steps_text = steps_response.choices[0].message.content
    steps_list = steps_text.split("\n")  # Assuming steps are returned as a list-like format
    
    return steps_list


@weave.op()
def generate_content_with_steps(o1_assistant_prompt: str, o1_user_prompt: str, steps_list: list) -> str:
    # Step 2: Use the steps as input to generate the final content strategy
    steps_text = "\n".join(f"{i+1}. {step}" for i, step in enumerate(steps_list))

    strategy_prompt = (
        f"Using the following steps, please generate a content strategy to {o1_user_prompt}:\n{steps_text}\n"
        f"Please proceed directly without requesting additional information.")
    
    strategy_messages = [
        {"role": "assistant", "content": o1_assistant_prompt},
        {"role": "user", "content": strategy_prompt}
    ]
    
    # API call to generate the strategy
    strategy_response = client.chat.completions.create(
        model="gpt-4o",
        messages=strategy_messages,
        temperature=0.2,
        max_tokens=1000,
        frequency_penalty=0.0
    )
    
    strategy_text = strategy_response.choices[0].message.content
    
    # Only return the strategy text
    return strategy_text


# First, get the steps
steps = generate_steps("Initial assistant prompt", "desired user prompt")

# Then, pass the steps to generate the content strategy
strategy = generate_content_with_steps(o1_assistant_prompt, o1_user_prompt, steps)

print(strategy)