Multi-Chain Prompt with Traces

import os
import time
from freeplay import Freeplay, CallInfo, RecordPayload
from openai import OpenAI
import re

# Configure environment variables for API access
project_id = os.environ.get("FREEPLAY_PROJECT_ID")
freeplay_api_key = os.environ.get("FREEPLAY_API_KEY")
freeplay_url = os.environ.get("FREEPLAY_URL") # ie "https://app.freeplay.ai/api"
openai_api_key = os.environ.get("OPENAI_API_KEY")

# Initialize OpenAI client
openai = OpenAI(api_key=openai_api_key)

# Initialize Freeplay client with development API endpoint
fp_client = Freeplay(
    freeplay_api_key=freeplay_api_key,
    api_base=freeplay_url
)

# Create a new Freeplay session to group related completions
session = fp_client.sessions.create({})

user_input = "My favorite artist is Taylor Swift"

# Create a trace to combine multiple prompts into a single workflow
trace_info = session.create_trace(
    input=user_input, # Commonly user question but any str input to a trace
    agent_name="musicAgent", # Optionally pass an agent name
    custom_metadata={
        "version": "1.0.8"
    }
)

# =============================================================================
# Prompt 1: Generate Album Title
# =============================================================================

# Define variables for the album title generation prompt
prompt_vars_a = {"pop_star": "Taylor Swift"}

# Fetch and format the album generation prompt template
formatted_prompt_a = fp_client.prompts.get_formatted(
    project_id=project_id,           # Freeplay project identifier
    template_name="album_bot",       # Name of the prompt template
    environment="latest",            # Environment tag for prompt versioning
    variables=prompt_vars_a          # Variables to interpolate into the prompt
)

# Execute the OpenAI completion for album title generation
start = time.time()
chat_completion_a = openai.chat.completions.create(
    messages=formatted_prompt_a.llm_prompt,
    model=formatted_prompt_a.prompt_info.model,
    **formatted_prompt_a.prompt_info.model_parameters
)
end = time.time()

# Extract the generated album name from the response
album_name = chat_completion_a.choices[0].message.content
print(f"Album Name: {album_name}")

# Record the first completion to Freeplay for tracking and analysis
fp_client.recordings.create(
    RecordPayload(
        all_messages=formatted_prompt_a.all_messages(new_message={
            "role": chat_completion_a.choices[0].message.role,
            "content": chat_completion_a.choices[0].message.content
        }),
        inputs=prompt_vars_a,
        session_info=session.session_info,
        prompt_info=formatted_prompt_a.prompt_info,
        call_info=CallInfo.from_prompt_info(formatted_prompt_a.prompt_info, start, end),
        trace_info=trace_info
    ),
)

# =============================================================================
# Prompt 2: Generate Song List for Album
# =============================================================================

# Define variables for the song list generation prompt (includes generated album name)
prompt_vars_b = {"album_name": album_name, "pop_star": "Taylor Swift"}

# Fetch and format the song list generation prompt template
formatted_prompt_b = fp_client.prompts.get_formatted(
    project_id=project_id,
    template_name="song_bot",
    environment="latest",
    variables=prompt_vars_b
)

# Execute the OpenAI completion for song list generation
start = time.time()
chat_completion_b = openai.chat.completions.create(
    messages=formatted_prompt_b.llm_prompt,
    model=formatted_prompt_b.prompt_info.model,
    **formatted_prompt_b.prompt_info.model_parameters
)
end = time.time()

song_track = chat_completion_b.choices[0].message.content

# Record the second completion to Freeplay for tracking and analysis
fp_client.recordings.create(
  RecordPayload(
				project_id=project_id,
        all_messages=formatted_prompt_b.all_messages(new_message={
            "role": chat_completion_b.choices[0].message.role,
            "content": song_track
        }),
        inputs=prompt_vars_b,
        session_info=session.session_info,
        prompt_version_info=formatted_prompt_b.prompt_info,
        call_info=CallInfo.from_prompt_info(formatted_prompt_b.prompt_info, start, end),
        trace_info=trace_info
    )
)

# Record the final output to complete the trace
trace_info.record_output(
    project_id=project_id,
    output=song_track, # Final trace output (str)
    # Optional code evals logged to the trace
    eval_results={
      "sentiment": 0.7,
      "songTrackLength": len(re.findall(r'\n', song_track)) + 1, # count number of lines
    }
)