End-to-end Generative Optimization for AI Agents

PyPI PyPI - Python Version

[This repository accomponanies the Trace paper. It is a fully functional implementation of the platform for generative optimization described in the paper, and contains code necessary to reproduce the experiments reported. This library was implemented and maintained by the authors while they were at Microsoft.]

Trace is a new AutoDiff-like tool for training AI systems end-to-end with general feedback (like numerical rewards or losses, natural language text, compiler errors, etc.). Trace generalizes the back-propagation algorithm by capturing and propagating an AI system's execution trace. Trace is implemented as a PyTorch-like Python library. Users write Python code directly and can use Trace primitives to optimize certain parts, just like training neural networks!

Setup

Simply run

pip install trace-opt

Or for development, clone the repo and run the following.

pip install -e .

The library requires Python >= 3.9. By default (starting with v0.1.3.5), we use LiteLLM as the backend of LLMs. For backward compatibility, we provide backend-support with AutoGen; when installing, users can add [autogen] tag to install a compatible AutoGen version (e.g., pip install trace-opt[autogen]). You may require Git Large File Storage if git is unable to clone the repository.

For questions or reporting bugs, please use Github Issues or post on our Discord channel. We actively check these channels.

Updates

2025.2.7 Trace was featured in the G-Research NeurIPS highlight by the Science Director Hugh Salimbeni.
2024.12.10 Trace was demoed in person at NeurIPS 2024 Expo.
2024.11.05 Ching-An Cheng gave a talk at UW Robotics Colloquium on Trace: video.
2024.10.21 New paper by Nvidia, Stanford, Visa, & Intel applies Trace to optimize for mapper code of parallel programming (for scientific computing and matrix multiplication). Trace (OptoPrime) learns code achieving 1.3X speed up under 10 minutes, compared to the code optimized by a system engineer expert.
2024.9.30 Ching-An Cheng gave a talk to the AutoGen community: link.
2024.9.25 Trace Paper is accepted to NeurIPS 2024.
2024.9.14 TextGrad is available as an optimizer in Trace.
2024.8.18 Allen Nie gave a talk to Pasteur Labs & Institute for Simulation Intelligence.

QuickStart

Trace has two primitives: node and bundle. node is a primitive to define a node in the computation graph. bundle is a primitive to define a function that can be optimized.

from opto.trace import node

x = node(1, trainable=True)
y = node(3)
z = x / y
z2 = x / 3  # the int 3 would be converted to a node automatically

list_of_nodes = [x, node(2), node(3)]
node_of_list = node([1, 2, 3])

node_of_list.append(3)

# easy built-in computation graph visualization
z.backward("maximize z", visualize=True, print_limit=25)

Once a node is declared, all the following operations on the node object will be automatically traced. We built many magic functions to make a node object act like a normal Python object. By marking trainable=True, we tell our optimizer that this node's content can be changed by the optimizer.

For functions, Trace uses decorators like @bundle to wrap over Python functions. A bundled function behaves like any other Python function.

from opto.trace import node, bundle


@bundle(trainable=True)
def strange_sort_list(lst):
    '''
    Given list of integers, return list in strange order.
    Strange sorting, is when you start with the minimum value,
    then maximum of the remaining integers, then minimum and so on.
    '''
    lst = sorted(lst)
    return lst


test_input = [1, 2, 3, 4]
test_output = strange_sort_list(test_input)
print(test_output)

Now, after declaring what is trainable and what isn't, and use node and bundle to define the computation graph, we can use the optimizer to optimize the computation graph.

from opto.optimizers import OptoPrime


# we first declare a feedback function
# think of this as the reward function (or loss function)
def get_feedback(predict, target):
    if predict == target:
        return "test case passed!"
    else:
        return "test case failed!"


test_ground_truth = [1, 4, 2, 3]
test_input = [1, 2, 3, 4]

epoch = 2

optimizer = OptoPrime(strange_sort_list.parameters())

for i in range(epoch):
    print(f"Training Epoch {i}")
    test_output = strange_sort_list(test_input)
    correctness = test_output.eq(test_ground_truth)
    feedback = get_feedback(test_output, test_ground_truth)

    if correctness:
        break

    optimizer.zero_feedback()
    optimizer.backward(correctness, feedback)
    optimizer.step()

Then, we can use the familiar PyTorch-like syntax to conduct the optimization.

Here is another example of a simple sales agent:

from opto import trace

@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}!"

Imagine we have a feedback function (like a reward function) that tells us how well the agent is doing. We can then optimize this agent online:

from opto.optimizers import OptoPrime

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("Hola, soy Juan.")
            feedback = feedback_fn(greeting.data, 'es')
        except trace.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

agent = train()

Defining and training an agent through Trace will give you more flexibility and control over what the agent learns.

Tutorials

| Level | Tutorial | Run in Colab | Description | | --- |-------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | Beginner | Getting Started | | Introduces basic primitives like node and bundle. Showcases a code optimi

Trace

Install / Use

README

End-to-end Generative Optimization for AI Agents

Setup

Updates

QuickStart

Tutorials