{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Agents\n",
    "\n",
    "You may have agents with behaviors that do not fall into a preset. \n",
    "In such cases, you can build custom agents.\n",
    "\n",
    "All agents in AgentChat inherit from {py:class}`~autogen_agentchat.agents.BaseChatAgent` \n",
    "class and implement the following abstract methods and attributes:\n",
    "\n",
    "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages`: The abstract method that defines the behavior of the agent in response to messages. This method is called when the agent is asked to provide a response in {py:meth}`~autogen_agentchat.agents.BaseChatAgent.run`. It returns a {py:class}`~autogen_agentchat.base.Response` object.\n",
    "- {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_reset`: The abstract method that resets the agent to its initial state. This method is called when the agent is asked to reset itself.\n",
    "- {py:attr}`~autogen_agentchat.agents.BaseChatAgent.produced_message_types`: The list of possible {py:class}`~autogen_agentchat.messages.ChatMessage` message types the agent can produce in its response.\n",
    "\n",
    "Optionally, you can implement the the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream` method to stream messages as they are generated by the agent. If this method is not implemented, the agent\n",
    "uses the default implementation of {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages_stream`\n",
    "that calls the {py:meth}`~autogen_agentchat.agents.BaseChatAgent.on_messages` method and\n",
    "yields all messages in the response."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CounterDownAgent\n",
    "\n",
    "In this example, we create a simple agent that counts down from a given number to zero,\n",
    "and produces a stream of messages with the current count."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import AsyncGenerator, List, Sequence\n",
    "\n",
    "from autogen_agentchat.agents import BaseChatAgent\n",
    "from autogen_agentchat.base import Response\n",
    "from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage\n",
    "from autogen_core import CancellationToken\n",
    "\n",
    "\n",
    "class CountDownAgent(BaseChatAgent):\n",
    "    def __init__(self, name: str, count: int = 3):\n",
    "        super().__init__(name, \"A simple agent that counts down.\")\n",
    "        self._count = count\n",
    "\n",
    "    @property\n",
    "    def produced_message_types(self) -> List[type[ChatMessage]]:\n",
    "        return [TextMessage]\n",
    "\n",
    "    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n",
    "        # Calls the on_messages_stream.\n",
    "        response: Response | None = None\n",
    "        async for message in self.on_messages_stream(messages, cancellation_token):\n",
    "            if isinstance(message, Response):\n",
    "                response = message\n",
    "        assert response is not None\n",
    "        return response\n",
    "\n",
    "    async def on_messages_stream(\n",
    "        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken\n",
    "    ) -> AsyncGenerator[AgentMessage | Response, None]:\n",
    "        inner_messages: List[AgentMessage] = []\n",
    "        for i in range(self._count, 0, -1):\n",
    "            msg = TextMessage(content=f\"{i}...\", source=self.name)\n",
    "            inner_messages.append(msg)\n",
    "            yield msg\n",
    "        # The response is returned at the end of the stream.\n",
    "        # It contains the final message and all the inner messages.\n",
    "        yield Response(chat_message=TextMessage(content=\"Done!\", source=self.name), inner_messages=inner_messages)\n",
    "\n",
    "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
    "        pass\n",
    "\n",
    "\n",
    "async def run_countdown_agent() -> None:\n",
    "    # Create a countdown agent.\n",
    "    countdown_agent = CountDownAgent(\"countdown\")\n",
    "\n",
    "    # Run the agent with a given task and stream the response.\n",
    "    async for message in countdown_agent.on_messages_stream([], CancellationToken()):\n",
    "        if isinstance(message, Response):\n",
    "            print(message.chat_message.content)\n",
    "        else:\n",
    "            print(message.content)\n",
    "\n",
    "\n",
    "# Use asyncio.run(run_countdown_agent()) when running in a script.\n",
    "await run_countdown_agent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UserProxyAgent \n",
    "\n",
    "A common use case for building a custom agent is to create an agent that acts as a proxy for the user.\n",
    "\n",
    "In the example below we show how to implement a `UserProxyAgent` - an agent that asks the user to enter\n",
    "some text through console and then returns that message as a response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "from typing import List, Sequence\n",
    "\n",
    "from autogen_agentchat.agents import BaseChatAgent\n",
    "from autogen_agentchat.base import Response\n",
    "from autogen_agentchat.messages import ChatMessage\n",
    "from autogen_core import CancellationToken\n",
    "\n",
    "\n",
    "class UserProxyAgent(BaseChatAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name, \"A human user.\")\n",
    "\n",
    "    @property\n",
    "    def produced_message_types(self) -> List[type[ChatMessage]]:\n",
    "        return [TextMessage]\n",
    "\n",
    "    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:\n",
    "        user_input = await asyncio.get_event_loop().run_in_executor(None, input, \"Enter your response: \")\n",
    "        return Response(chat_message=TextMessage(content=user_input, source=self.name))\n",
    "\n",
    "    async def on_reset(self, cancellation_token: CancellationToken) -> None:\n",
    "        pass\n",
    "\n",
    "\n",
    "async def run_user_proxy_agent() -> None:\n",
    "    user_proxy_agent = UserProxyAgent(name=\"user_proxy_agent\")\n",
    "    response = await user_proxy_agent.on_messages([], CancellationToken())\n",
    "    print(response.chat_message.content)\n",
    "\n",
    "\n",
    "# Use asyncio.run(run_user_proxy_agent()) when running in a script.\n",
    "await run_user_proxy_agent()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}