autogen/python/packages/autogen-studio/notebooks/tutorial.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## AutoGen Studio Agent Workflow API Example\n",
                "\n",
                "This notebook focuses on demonstrating capabilities of the autogen studio workflow python api.  \n",
                "\n",
                "- Declarative Specification of an Agent Team\n",
                "- Loading the specification and running the resulting agent\n",
                "\n",
                " "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?', type='TextMessage'), ToolCallRequestEvent(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_jcgtAVlBvTFzVpPxKX88Xsa4', arguments='{\"city\":\"New York\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_jcgtAVlBvTFzVpPxKX88Xsa4')], type='ToolCallExecutionEvent'), TextMessage(source='writing_agent', models_usage=None, content='The weather in New York is 73 degrees and Sunny.', type='TextMessage'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14), content='The current weather in New York is 73 degrees and sunny.', type='TextMessage')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=5.103050947189331\n"
                    ]
                }
            ],
            "source": [
                "from autogenstudio.teammanager import TeamManager\n",
                "\n",
                "wm = TeamManager()\n",
                "result = await wm.run(task=\"What is the weather in New York?\", team_config=\"team.json\")\n",
                "print(result)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "source='user' models_usage=None content='What is the weather in New York?' type='TextMessage'\n",
                        "source='writing_agent' models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15) content=[FunctionCall(id='call_EwdwWogp5jDKdB7t9WGCNjZW', arguments='{\"city\":\"New York\"}', name='get_weather')] type='ToolCallRequestEvent'\n",
                        "source='writing_agent' models_usage=None content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_EwdwWogp5jDKdB7t9WGCNjZW')] type='ToolCallExecutionEvent'\n",
                        "source='writing_agent' models_usage=None content='The weather in New York is 73 degrees and Sunny.' type='TextMessage'\n",
                        "source='writing_agent' models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14) content='The weather in New York is currently 73 degrees and sunny.' type='TextMessage'\n",
                        "task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?', type='TextMessage'), ToolCallRequestEvent(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_EwdwWogp5jDKdB7t9WGCNjZW', arguments='{\"city\":\"New York\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_EwdwWogp5jDKdB7t9WGCNjZW')], type='ToolCallExecutionEvent'), TextMessage(source='writing_agent', models_usage=None, content='The weather in New York is 73 degrees and Sunny.', type='TextMessage'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14), content='The weather in New York is currently 73 degrees and sunny.', type='TextMessage')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=1.284574270248413\n"
                    ]
                }
            ],
            "source": [
                "result_stream =  wm.run_stream(task=\"What is the weather in New York?\", team_config=\"team.json\")\n",
                "async for response in result_stream:\n",
                "    print(response)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## AutoGen Studio Database API\n",
                "\n",
                "Api for creating objects and serializing to a database."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "Response(message='Database is ready', status=True, data=None)"
                        ]
                    },
                    "execution_count": 3,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "from autogenstudio.database import DatabaseManager\n",
                "import os\n",
                "# delete database\n",
                "# if os.path.exists(\"test.db\"):\n",
                "#     os.remove(\"test.db\")\n",
                "\n",
                "os.makedirs(\"test\", exist_ok=True)\n",
                "# create a database\n",
                "dbmanager = DatabaseManager(engine_uri=\"sqlite:///test.db\", base_dir=\"test\")\n",
                "dbmanager.initialize_database()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "from sqlmodel import Session, text, select\n",
                "from autogenstudio.datamodel.types import  ModelTypes,  TeamTypes,  AgentTypes,   ToolConfig, ToolTypes, OpenAIModelConfig, RoundRobinTeamConfig, MaxMessageTerminationConfig, AssistantAgentConfig, TerminationTypes\n",
                "\n",
                "from autogenstudio.datamodel.db import Model, Team, Agent, Tool,LinkTypes\n",
                "\n",
                "user_id = \"guestuser@gmail.com\" \n",
                "\n",
                "gpt4_model = Model(user_id=user_id, config= OpenAIModelConfig(model=\"gpt-4o-2024-08-06\", model_type=ModelTypes.OPENAI).model_dump() )\n",
                "\n",
                "weather_tool = Tool(user_id=user_id, config=ToolConfig(name=\"get_weather\", description=\"Get the weather for a city\", content=\"async def get_weather(city: str) -> str:\\n    return f\\\"The weather in {city} is 73 degrees and Sunny.\\\"\",tool_type=ToolTypes.PYTHON_FUNCTION).model_dump() )\n",
                "\n",
                "adding_tool = Tool(user_id=user_id, config=ToolConfig(name=\"add\", description=\"Add two numbers\", content=\"async def add(a: int, b: int) -> int:\\n    return a + b\", tool_type=ToolTypes.PYTHON_FUNCTION).model_dump() )\n",
                "\n",
                "writing_agent = Agent(user_id=user_id,\n",
                "                      config=AssistantAgentConfig(\n",
                "                          name=\"writing_agent\",\n",
                "                          tools=[weather_tool.config],\n",
                "                          agent_type=AgentTypes.ASSISTANT,\n",
                "                          model_client=gpt4_model.config\n",
                "                          ).model_dump()\n",
                "                    )\n",
                "\n",
                "team = Team(user_id=user_id, config=RoundRobinTeamConfig(\n",
                "    name=\"weather_team\",\n",
                "    participants=[writing_agent.config],\n",
                "    termination_condition=MaxMessageTerminationConfig(termination_type=TerminationTypes.MAX_MESSAGES, max_messages=5).model_dump(),\n",
                "    team_type=TeamTypes.ROUND_ROBIN\n",
                "    ).model_dump()\n",
                ")\n",
                "\n",
                "with Session(dbmanager.engine) as session:\n",
                "    session.add(gpt4_model)\n",
                "    session.add(weather_tool)\n",
                "    session.add(adding_tool)\n",
                "    session.add(writing_agent)\n",
                "    session.add(team)\n",
                "    session.commit()\n",
                "\n",
                "    dbmanager.link(LinkTypes.AGENT_MODEL, writing_agent.id, gpt4_model.id)\n",
                "    dbmanager.link(LinkTypes.AGENT_TOOL, writing_agent.id, weather_tool.id)\n",
                "    dbmanager.link(LinkTypes.AGENT_TOOL, writing_agent.id, adding_tool.id)\n",
                "    dbmanager.link(LinkTypes.TEAM_AGENT, team.id, writing_agent.id)\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "2 teams in database\n"
                    ]
                }
            ],
            "source": [
                "all_teams = dbmanager.get(Team)\n",
                "print(len(all_teams.data), \"teams in database\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Configuration Manager\n",
                "\n",
                "Helper class to mostly import teams/agents/models/tools etc into a database."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [],
            "source": [
                "from autogenstudio.database import ConfigurationManager\n",
                "\n",
                "config_manager = ConfigurationManager(dbmanager)\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "message='Team Created Successfully' status=True data={'id': 4, 'updated_at': datetime.datetime(2024, 12, 15, 15, 52, 21, 674916), 'version': '0.0.1', 'created_at': datetime.datetime(2024, 12, 15, 15, 52, 21, 674910), 'user_id': 'user_id', 'config': {'version': '1.0.0', 'component_type': 'team', 'name': 'weather_team', 'participants': [{'version': '1.0.0', 'component_type': 'agent', 'name': 'writing_agent', 'agent_type': 'AssistantAgent', 'description': None, 'model_client': {'version': '1.0.0', 'component_type': 'model', 'model': 'gpt-4o-2024-08-06', 'model_type': 'OpenAIChatCompletionClient', 'api_key': None, 'base_url': None}, 'tools': [{'version': '1.0.0', 'component_type': 'tool', 'name': 'get_weather', 'description': 'Get the weather for a city', 'content': 'async def get_weather(city: str) -> str:\\n    return f\"The weather in {city} is 73 degrees and Sunny.\"', 'tool_type': 'PythonFunction'}], 'system_message': None}], 'team_type': 'RoundRobinGroupChat', 'termination_condition': {'version': '1.0.0', 'component_type': 'termination', 'termination_type': 'MaxMessageTermination', 'max_messages': 5}, 'max_turns': None}}\n"
                    ]
                }
            ],
            "source": [
                "result = await config_manager.import_component(\"team.json\", user_id=\"user_id\", check_exists=True)\n",
                "print(result)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "message='Directory import complete' status=True data=[{'component': 'team', 'status': True, 'message': 'Team Created Successfully', 'id': 5}]\n"
                    ]
                }
            ],
            "source": [
                "result = await config_manager.import_directory(\".\", user_id=\"user_id\", check_exists=False)\n",
                "print(result)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "5 teams in database\n"
                    ]
                }
            ],
            "source": [
                "all_teams = dbmanager.get(Team)\n",
                "print(len(all_teams.data), \"teams in database\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Sample AgentChat Example (Python)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "metadata": {},
            "outputs": [],
            "source": [
                "from autogen_agentchat.agents import AssistantAgent\n",
                "from autogen_agentchat.conditions import TextMentionTermination\n",
                "from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat\n",
                "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
                "\n",
                "planner_agent = AssistantAgent(\n",
                "    \"planner_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can plan trips.\",\n",
                "    system_message=\"You are a helpful assistant that can suggest a travel plan for a user based on their request. Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "local_agent = AssistantAgent(\n",
                "    \"local_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A local assistant that can suggest local activities or places to visit.\",\n",
                "    system_message=\"You are a helpful assistant that can suggest authentic and interesting local activities or places to visit for a user and can utilize any context information provided. Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "language_agent = AssistantAgent(\n",
                "    \"language_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can provide language tips for a given destination.\",\n",
                "    system_message=\"You are a helpful assistant that can review travel plans, providing feedback on important/critical tips about how best to address language or communication challenges for the given destination. If the plan already includes language tips, you can mention that the plan is satisfactory, with rationale.Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "travel_summary_agent = AssistantAgent(\n",
                "    \"travel_summary_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can summarize the travel plan.\",\n",
                "    system_message=\"You are a helpful assistant that can take in all of the suggestions and advice from the other agents and provide a detailed tfinal travel plan. You must ensure th b at the final plan is integrated and complete. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE.Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "termination = TextMentionTermination(\"TERMINATE\")\n",
                "group_chat = RoundRobinGroupChat(\n",
                "    [planner_agent, local_agent, language_agent, travel_summary_agent], termination_condition=termination\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "source='user' models_usage=None content='Plan a 3 day trip to Nepal.' type='TextMessage'\n",
                        "source='planner_agent' models_usage=RequestUsage(prompt_tokens=45, completion_tokens=53) content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.' type='TextMessage'\n",
                        "source='local_agent' models_usage=RequestUsage(prompt_tokens=115, completion_tokens=53) content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.' type='TextMessage'\n",
                        "source='language_agent' models_usage=RequestUsage(prompt_tokens=199, completion_tokens=42) content=\"For your trip to Nepal, it's crucial to learn some phrases in Nepali since English is not widely spoken outside of major cities and tourist areas; even a simple phrasebook or translation app would be beneficial.\" type='TextMessage'\n",
                        "source='travel_summary_agent' models_usage=RequestUsage(prompt_tokens=265, completion_tokens=298) content=\"Day 1: Begin your journey in Kathmandu, where you can visit the historic Durbar Square, a UNESCO World Heritage site that showcases intricate woodcarving and houses the iconic Kasthamandap Temple. From there, proceed to the sacred Pashupatinath Temple, a significant Hindu pilgrimage site on the banks of the holy Bagmati River.\\n\\nDay 2: Embark on an early morning scenic flight over the Everest range. This one-hour flight provides a breathtaking view of the world's highest peak along with other neighboring peaks. Standard flights depart from Tribhuvan International Airport between 6:30 AM to 7:30 AM depending on the weather. Spend the remainder of the day exploring the local markets in Kathmandu, sampling a variety of Nepalese cuisines and shopping for unique souvenirs.\\n\\nDay 3: Finally, take a short flight or drive to Pokhara, the gateway to the Annapurna region. Embark on a guided hike enjoying the stunning backdrop of the Annapurna ranges and the serene Phewa lake.\\n\\nRemember to bring along a phrasebook or translation app, as English is not widely spoken in Nepal, particularly outside of major cities and tourist hotspots. \\n\\nPack comfortable trekking gear, adequate water, medical and emergency supplies. It's also advisable to check on the weather updates, as conditions can change rapidly, particularly in mountainous areas. Enjoy your Nepal expedition!TERMINATE\" type='TextMessage'\n",
                        "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Plan a 3 day trip to Nepal.', type='TextMessage'), TextMessage(source='planner_agent', models_usage=RequestUsage(prompt_tokens=45, completion_tokens=53), content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.', type='TextMessage'), TextMessage(source='local_agent', models_usage=RequestUsage(prompt_tokens=115, completion_tokens=53), content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.', type='TextMessage'), TextMessage(source='language_agent', models_usage=RequestUsage(prompt_tokens=199, completion_tokens=42), content=\"For your trip to Nepal, it's crucial to learn some phrases in Nepali since English is not widely spoken outside of major cities and tourist areas; even a simple phrasebook or translation app would be beneficial.\", type='TextMessage'), TextMessage(source='travel_summary_agent', models_usage=RequestUsage(prompt_tokens=265, completion_tokens=298), content=\"Day 1: Begin your journey in Kathmandu, where you can visit the historic Durbar Square, a UNESCO World Heritage site that showcases intricate woodcarving and houses the iconic Kasthamandap Temple. From there, proceed to the sacred Pashupatinath Temple, a significant Hindu pilgrimage site on the banks of the holy Bagmati River.\\n\\nDay 2: Embark on an early morning scenic flight over the Everest range. This one-hour flight provides a breathtaking view of the world's highest peak along with other neighboring peaks. Standard flights depart from Tribhuvan International Airport between 6:30 AM to 7:30 AM depending on the weather. Spend the remainder of the day exploring the local markets in Kathmandu, sampling a variety of Nepalese cuisines and shopping for unique souvenirs.\\n\\nDay 3: Finally, take a short flight or drive to Pokhara, the gateway to the Annapurna region. Embark on a guided hike enjoying the stunning backdrop of the Annapurna ranges and the serene Phewa lake.\\n\\nRemember to bring along a phrasebook or translation app, as English is not widely spoken in Nepal, particularly outside of major cities and tourist hotspots. \\n\\nPack comfortable trekking gear, adequate water, medical and emergency supplies. It's also advisable to check on the weather updates, as conditions can change rapidly, particularly in mountainous areas. Enjoy your Nepal expedition!TERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")\n"
                    ]
                }
            ],
            "source": [
                "\n",
                "result = group_chat.run_stream(task=\"Plan a 3 day trip to Nepal.\")\n",
                "async for response in result:\n",
                "    print(response)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "agnext",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.9"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}