Greeting Agent¶

Introduction¶

This notebook will focus on using trace to optimize a sales agent that greets a customer.

Setup and Installation¶

Let's start by importing the necessary libraries.

In [ ]:

%pip install trace-opt

%pip install trace-opt

In [ ]:

%pip install openai httpx pywidgets

%pip install openai httpx pywidgets

In [ ]:

from opto import trace
from opto.trace import node, bundle, model, ExecutionError
from opto.optimizers import OptoPrime

from opto import trace from opto.trace import node, bundle, model, ExecutionError from opto.optimizers import OptoPrime

Add API keys for LLM calls. Run the code below:

In [1]:

import os
import ipywidgets as widgets
from IPython.display import display

# Function to save the environment variable and API key
def save_env_variable(env_name, api_key):
    # Validate inputs
    if not env_name.strip():
        print("⚠️ Environment variable name cannot be empty.")
        return
    if not api_key.strip():
        print("⚠️ API key cannot be empty.")
        return
    
    # Store the API key as an environment variable
    os.environ[env_name] = api_key
    globals()[env_name] = api_key  # Set it as a global variable
    print(f"✅ API key has been set for environment variable: {env_name}")

# Create the input widgets
env_name_input = widgets.Text(
    value="OPENAI_API_KEY",  # Default value
    description="Env Name:",
    placeholder="Enter env variable name (e.g., MY_API_KEY)",
)

api_key_input = widgets.Password(
    description="API Key:",
    placeholder="Enter your API key",
)

# Create the button to submit the inputs
submit_button = widgets.Button(description="Set API Key")

# Display the widgets
display(env_name_input, api_key_input, submit_button)

# Callback function for the button click
def on_button_click(b):
    env_name = env_name_input.value
    api_key = api_key_input.value
    save_env_variable(env_name, api_key)

# Attach the callback to the button
submit_button.on_click(on_button_click)

import os import ipywidgets as widgets from IPython.display import display # Function to save the environment variable and API key def save_env_variable(env_name, api_key): # Validate inputs if not env_name.strip(): print("⚠️ Environment variable name cannot be empty.") return if not api_key.strip(): print("⚠️ API key cannot be empty.") return # Store the API key as an environment variable os.environ[env_name] = api_key globals()[env_name] = api_key # Set it as a global variable print(f"✅ API key has been set for environment variable: {env_name}") # Create the input widgets env_name_input = widgets.Text( value="OPENAI_API_KEY", # Default value description="Env Name:", placeholder="Enter env variable name (e.g., MY_API_KEY)", ) api_key_input = widgets.Password( description="API Key:", placeholder="Enter your API key", ) # Create the button to submit the inputs submit_button = widgets.Button(description="Set API Key") # Display the widgets display(env_name_input, api_key_input, submit_button) # Callback function for the button click def on_button_click(b): env_name = env_name_input.value api_key = api_key_input.value save_env_variable(env_name, api_key) # Attach the callback to the button submit_button.on_click(on_button_click)

Text(value='OPENAI_API_KEY', description='Env Name:', placeholder='Enter env variable name (e.g., MY_API_KEY)'…

Password(description='API Key:', placeholder='Enter your API key')

Button(description='Set API Key', style=ButtonStyle())

Define an Agent¶

In here, we use @trace.bundle to wrap functions so that they show up in TraceGraph. We use trace.node to wrap system prompts. @trace.model does not do much, except to provide us some convenience to grab all the trainable parameters. |

In [ ]:

@trace.model
class Agent:

    def __init__(self, system_prompt):
        self.system_prompt = system_prompt
        self.instruct1 = trace.node("Decide the language", trainable=True)
        self.instruct2 = trace.node("Extract name if it's there", trainable=True)

    def __call__(self, user_query):
        response = trace.operators.call_llm(self.system_prompt,
                                            self.instruct1, user_query)
        en_or_es = self.decide_lang(response)

        user_name = trace.operators.call_llm(self.system_prompt,
                                             self.instruct2, user_query)
        greeting = self.greet(en_or_es, user_name)

        return greeting

    @trace.bundle(trainable=True)
    def decide_lang(self, response):
        """Map the language into a variable"""
        return

    @trace.bundle(trainable=True)
    def greet(self, lang, user_name):
        """Produce a greeting based on the language"""
        greeting = "Hola"
        return f"{greeting}, {user_name}!"

@trace.model class Agent: def __init__(self, system_prompt): self.system_prompt = system_prompt self.instruct1 = trace.node("Decide the language", trainable=True) self.instruct2 = trace.node("Extract name if it's there", trainable=True) def __call__(self, user_query): response = trace.operators.call_llm(self.system_prompt, self.instruct1, user_query) en_or_es = self.decide_lang(response) user_name = trace.operators.call_llm(self.system_prompt, self.instruct2, user_query) greeting = self.greet(en_or_es, user_name) return greeting @trace.bundle(trainable=True) def decide_lang(self, response): """Map the language into a variable""" return @trace.bundle(trainable=True) def greet(self, lang, user_name): """Produce a greeting based on the language""" greeting = "Hola" return f"{greeting}, {user_name}!"

Define Feedback and Training¶

In [ ]:

def feedback_fn(generated_response, gold_label='en'):
    if  gold_label == 'en' and 'Hello' in generated_response:
        return "Correct"
    elif gold_label == 'es' and 'Hola' in generated_response:
        return "Correct"
    else:
        return "Incorrect"


def train():
    epoch = 3
    agent = Agent("You are a sales assistant.")
    optimizer = OptoPrime(agent.parameters())

    for i in range(epoch):
        print(f"Training Epoch {i}")
        try:
            greeting = agent("Hello, I'm John.")
            feedback = feedback_fn(greeting.data, 'en')
        except ExecutionError as e:
            greeting = e.exception_node
            feedback, terminal, reward = greeting.data, False, 0

        optimizer.zero_feedback()
        optimizer.backward(greeting, feedback)
        optimizer.step(verbose=True)

        if feedback == 'Correct':
            break

    return agent

def feedback_fn(generated_response, gold_label='en'): if gold_label == 'en' and 'Hello' in generated_response: return "Correct" elif gold_label == 'es' and 'Hola' in generated_response: return "Correct" else: return "Incorrect" def train(): epoch = 3 agent = Agent("You are a sales assistant.") optimizer = OptoPrime(agent.parameters()) for i in range(epoch): print(f"Training Epoch {i}") try: greeting = agent("Hello, I'm John.") feedback = feedback_fn(greeting.data, 'en') except ExecutionError as e: greeting = e.exception_node feedback, terminal, reward = greeting.data, False, 0 optimizer.zero_feedback() optimizer.backward(greeting, feedback) optimizer.step(verbose=True) if feedback == 'Correct': break return agent

In [ ]:

agent = train()

agent = train()