1  """
 2  This is an example for leveraging MLflow's auto tracing capabilities for AutoGen.
 3  
 4  For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
 5  """
 6  
 7  import os
 8  from typing import Annotated, Literal
 9  
10  from autogen import ConversableAgent
11  
12  import mlflow
13  
14  # Turn on auto tracing for AutoGen by calling mlflow.autogen.autolog()
15  mlflow.autogen.autolog()
16  
17  
18  config_list = [
19      {
20          "model": "gpt-4o-mini",
21          # Please set your OpenAI API Key to the OPENAI_API_KEY env var before running this example
22          "api_key": os.environ.get("OPENAI_API_KEY"),
23      }
24  ]
25  
26  Operator = Literal["+", "-", "*", "/"]
27  
28  
29  def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
30      if operator == "+":
31          return a + b
32      elif operator == "-":
33          return a - b
34      elif operator == "*":
35          return a * b
36      elif operator == "/":
37          return int(a / b)
38      else:
39          raise ValueError("Invalid operator")
40  
41  
42  # First define the assistant agent that suggests tool calls.
43  assistant = ConversableAgent(
44      name="Assistant",
45      system_message="You are a helpful AI assistant. "
46      "You can help with simple calculations. "
47      "Return 'TERMINATE' when the task is done.",
48      llm_config={"config_list": config_list},
49  )
50  
51  # The user proxy agent is used for interacting with the assistant agent
52  # and executes tool calls.
53  user_proxy = ConversableAgent(
54      name="ToolAgent",
55      llm_config=False,
56      is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
57      human_input_mode="NEVER",
58  )
59  
60  # Register the tool signature with the assistant agent.
61  assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)
62  user_proxy.register_for_execution(name="calculator")(calculator)
63  response = user_proxy.initiate_chat(assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?")