{
 "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?'), ToolCallMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_1LrIS8t1wuK3KzDPP4DuwuiQ', arguments='{\"city\":\"New York\"}', name='get_weather')]), ToolCallResultMessage(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_1LrIS8t1wuK3KzDPP4DuwuiQ')]), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=97, completion_tokens=14), content='The current weather in New York is 73 degrees and sunny.'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=123, completion_tokens=9), content='What is the weather in Los Angeles?')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=2.2458291053771973\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?'\n",
      "source='writing_agent' models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15) content=[FunctionCall(id='call_Otg8Exxp7uAcsZnRyTWRYGkb', arguments='{\"city\":\"New York\"}', name='get_weather')]\n",
      "source='writing_agent' models_usage=None content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_Otg8Exxp7uAcsZnRyTWRYGkb')]\n",
      "source='writing_agent' models_usage=RequestUsage(prompt_tokens=97, completion_tokens=14) content='The weather in New York is currently 73 degrees and sunny.'\n",
      "source='writing_agent' models_usage=RequestUsage(prompt_tokens=123, completion_tokens=16) content='Would you like to know anything else about the weather or any other information?'\n",
      "task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?'), ToolCallMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_Otg8Exxp7uAcsZnRyTWRYGkb', arguments='{\"city\":\"New York\"}', name='get_weather')]), ToolCallResultMessage(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_Otg8Exxp7uAcsZnRyTWRYGkb')]), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=97, completion_tokens=14), content='The weather in New York is currently 73 degrees and sunny.'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=123, completion_tokens=16), content='Would you like to know anything else about the weather or any other information?')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=1.9408440589904785\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": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO  [alembic.runtime.migration] Context impl SQLiteImpl.\n",
      "INFO  [alembic.runtime.migration] Will assume non-transactional DDL.\n",
      "\u001b[32m2024-11-09 11:55:02.090\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mautogenstudio.database.schema_manager\u001b[0m:\u001b[36mupgrade_schema\u001b[0m:\u001b[36m390\u001b[0m - \u001b[1mSchema upgraded successfully to head\u001b[0m\n",
      "\u001b[32m2024-11-09 11:55:02.090\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mautogenstudio.database.db_manager\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m34\u001b[0m - \u001b[1mDatabase schema was upgraded automatically\u001b[0m\n",
      "\u001b[32m2024-11-09 11:55:02.092\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mautogenstudio.database.db_manager\u001b[0m:\u001b[36mcreate_db_and_tables\u001b[0m:\u001b[36m108\u001b[0m - \u001b[1mDatabase tables created successfully\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from autogenstudio.database import DatabaseManager \n",
    "\n",
    "# create a database\n",
    "dbmanager = DatabaseManager(engine_uri=\"sqlite:///test.db\")\n",
    "dbmanager.create_db_and_tables() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from sqlmodel import Session, text, select\n",
    "from autogenstudio.datamodel import Model, ModelConfig, ModelTypes, Team, TeamConfig, TeamTypes, Agent, AgentConfig, AgentTypes, Tool, ToolConfig, LinkTypes,ToolTypes\n",
    "\n",
    "user_id = \"guestuser@gmail.com\"\n",
    "from autogenstudio.datamodel import ModelConfig, Model, TeamConfig, Team, Tool, Agent, AgentConfig, TerminationConfig, TerminationTypes, ModelTypes, TeamTypes, AgentTypes, ToolConfig, LinkTypes, TerminationTypes\n",
    "\n",
    "gpt4_model = Model(user_id=user_id, config= ModelConfig(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=AgentConfig(\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=TeamConfig(\n",
    "    name=\"weather_team\",\n",
    "    participants=[writing_agent.config],\n",
    "    termination_condition=TerminationConfig(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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "43 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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogenstudio.database import ConfigurationManager  \n",
    "config_manager = ConfigurationManager(dbmanager)\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "message=\"team with team_type='TeamTypes.ROUND_ROBIN' and name='weather_team' already exists\" status=True data={'id': 42}\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": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "message='Directory import complete' status=True data=[{'component': 'team', 'status': True, 'message': 'Team Created Successfully', 'id': 44}]\n"
     ]
    }
   ],
   "source": [
    "result = await config_manager.import_directory(\".\", user_id=\"user_id\", check_exists=False)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "44 teams in database\n"
     ]
    }
   ],
   "source": [
    "all_teams = dbmanager.get(Team)\n",
    "print(len(all_teams.data), \"teams in database\")"
   ]
  }
 ],
 "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
}