Autogen

BRAINTRUST_API_KEY=your-api-key
BRAINTRUST_PARENT=project_name:autogen-demo
OPENAI_API_KEY=your-openai-key

import asyncio
import os
from typing import Optional
 
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.ui import Console
from autogen_core import SingleThreadedAgentRuntime
from autogen_ext.models.openai import OpenAIChatCompletionClient
from braintrust.otel import BraintrustSpanProcessor
from dotenv import load_dotenv
from opentelemetry import trace
from opentelemetry.instrumentation.openai import OpenAIInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
 
 
def setup_tracing() -> None:
    provider = TracerProvider()
    provider.add_span_processor(BraintrustSpanProcessor(filter_ai_spans=True))
    trace.set_tracer_provider(provider)
    OpenAIInstrumentor().instrument()
 
 
def percentage_change_tool(start: float, end: float) -> float:
    if start == 0:
        return float("inf")
    return ((end - start) / start) * 100
 
 
async def main(task: Optional[str] = None) -> None:
    load_dotenv()
    setup_tracing()
 
    model_client = OpenAIChatCompletionClient(model="gpt-4o")
 
    tracer = trace.get_tracer("autogen-demo-bot")
    with tracer.start_as_current_span("run_team"):
        planning_agent = AssistantAgent(
            "PlanningAgent",
            description="Plans tasks and delegates.",
            model_client=model_client,
            system_message=(
                "You are a planning agent. You only plan and delegate tasks; do not execute them.\n"
                "When assigning tasks, use this format: 1. <agent> : <task>\n"
                'After all tasks are complete, summarize the findings and end with "TERMINATE".'
            ),
        )
 
        data_analyst_agent = AssistantAgent(
            "DataAnalystAgent",
            description="Performs calculations.",
            model_client=model_client,
            tools=[percentage_change_tool],
            system_message=("You are a data analyst. Use the tools provided to compute numeric results."),
        )
 
        termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(max_messages=25)
 
        if not task:
            task = "You started with 100 apples, now you have 120 apples. what is the percentage change?"
 
        runtime = SingleThreadedAgentRuntime(tracer_provider=trace.get_tracer_provider())
        runtime.start()
 
        selector_prompt = (
            "Select an agent to perform task.\n\n{roles}\n\nCurrent conversation context:\n{history}\n\n"
            "Read the above conversation, then select an agent from {participants} to perform the next task.\n"
            "Make sure the planner agent has assigned tasks before other agents start working.\nOnly select one agent."
        )
 
        team = SelectorGroupChat(
            [planning_agent, data_analyst_agent],
            model_client=model_client,
            termination_condition=termination,
            selector_prompt=selector_prompt,
            allow_repeated_speaker=True,
            runtime=runtime,
        )
 
        await Console(team.run_stream(task=task))
 
        await runtime.stop()
 
    await model_client.close()
 
 
if __name__ == "__main__":
    asyncio.run(main())