---
title: Tracing AutoGen | Databricks on AWS
source: https://docs.databricks.com/aws/en/mlflow3/genai/tracing/integrations/autogen
ingestedAt: "2026-06-18T08:16:57.169Z"
---

    import osfrom typing import Annotated, Literalfrom autogen import ConversableAgentimport mlflow# Ensure your OPENAI_API_KEY (or other LLM provider keys) is set in your environment# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured# Turn on auto tracing for AutoGenmlflow.autogen.autolog()# Set up MLflow tracking on Databricksmlflow.set_tracking_uri("databricks")mlflow.set_experiment("/Shared/autogen-tracing-demo")# Define a simple multi-agent workflow using AutoGenconfig_list = [    {        "model": "gpt-4o-mini",        # Please set your OpenAI API Key to the OPENAI_API_KEY env var before running this example        "api_key": os.environ.get("OPENAI_API_KEY"),    }]Operator = Literal["+", "-", "*", "/"]def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:    if operator == "+":        return a + b    elif operator == "-":        return a - b    elif operator == "*":        return a * b    elif operator == "/":        return int(a / b)    else:        raise ValueError("Invalid operator")# First define the assistant agent that suggests tool calls.assistant = ConversableAgent(    name="Assistant",    system_message="You are a helpful AI assistant. "    "You can help with simple calculations. "    "Return 'TERMINATE' when the task is done.",    llm_config={"config_list": config_list},)# The user proxy agent is used for interacting with the assistant agent# and executes tool calls.user_proxy = ConversableAgent(    name="Tool Agent",    llm_config=False,    is_termination_msg=lambda msg: msg.get("content") is not None    and "TERMINATE" in msg["content"],    human_input_mode="NEVER",)# Register the tool signature with the assistant agent.assistant.register_for_llm(name="calculator", description="A simple calculator")(    calculator)user_proxy.register_for_execution(name="calculator")(calculator)response = user_proxy.initiate_chat(    assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?")
