notebooks/task_vs_message.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "pOgV4IUUZn6B",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "872ae3b8-49fe-46c4-fa55-9ccc889c18b9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m133.2/133.2 kB\u001b[0m \u001b[31m1.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.7/68.7 kB\u001b[0m \u001b[31m3.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.3/44.3 kB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m25.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m239.3/239.3 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m218.1/218.1 kB\u001b[0m \u001b[31m9.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m335.7/335.7 kB\u001b[0m \u001b[31m10.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m160.2/160.2 kB\u001b[0m \u001b[31m8.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m65.6/65.6 kB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m120.0/120.0 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m201.6/201.6 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.8/64.8 kB\u001b[0m \u001b[31m4.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.1/79.1 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m229.5/229.5 kB\u001b[0m \u001b[31m13.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m103.8/103.8 kB\u001b[0m \u001b[31m6.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.2/45.2 kB\u001b[0m \u001b[31m2.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "google-colab 1.0.0 requires requests==2.32.3, but you have requests 2.32.4 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "%pip install google-cloud-aiplatform httpx \"a2a-sdk\" --quiet\n",
        "%pip install --upgrade --quiet  \"google-adk\""
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "PROJECT_ID = '[your-project-id]'  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "LOCATION = '[your-google-cloud-region]'  # @param {type:\"string\", placeholder: \"[your-google-cloud-region]\", isTemplate: true}\n",
        "MODEL = 'gemini-2.5-flash'"
      ],
      "metadata": {
        "id": "M_qUEbg8cRiE"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "\n",
        "try:\n",
        "    auth.authenticate_user()\n",
        "    print('Colab user authenticated.')\n",
        "except Exception as e:\n",
        "    print(\n",
        "        f'Not in a Colab environment or auth failed: {e}. Assuming local gcloud auth.'\n",
        "    )"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QWE4yt7KN-pL",
        "outputId": "62b92a17-4759-42e0-feba-4db8b476eea6"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Colab user authenticated.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import uuid\n",
        "import pprint\n",
        "\n",
        "from collections.abc import AsyncIterable\n",
        "from typing import Any, Optional\n",
        "\n",
        "from a2a.server.agent_execution import AgentExecutor, RequestContext\n",
        "from a2a.server.apps import A2AStarletteApplication\n",
        "from a2a.server.events import EventQueue\n",
        "\n",
        "from pydantic import BaseModel\n",
        "from enum import Enum\n",
        "\n",
        "from a2a.types import (\n",
        "    Part,\n",
        "    Task,\n",
        "    TaskState,\n",
        "    TextPart,\n",
        "    MessageSendParams,\n",
        "    Role,\n",
        "    Message,\n",
        ")\n",
        "\n",
        "from a2a.utils import (\n",
        "    new_agent_parts_message,\n",
        "    new_agent_text_message,\n",
        "    new_task,\n",
        ")\n",
        "\n",
        "from a2a.server.request_handlers.default_request_handler import (\n",
        "    DefaultRequestHandler,\n",
        ")\n",
        "from a2a.server.tasks import InMemoryTaskStore, TaskUpdater\n",
        "\n",
        "from a2a.utils.errors import MethodNotImplementedError\n",
        "\n",
        "# Build agent with adk\n",
        "from google.adk.events import Event\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.adk.tools.tool_context import ToolContext\n",
        "from google.adk.agents.llm_agent import LlmAgent\n",
        "\n",
        "# Evaluate agent\n",
        "from google.cloud import aiplatform\n",
        "from google.genai import types\n",
        "\n",
        "pp = pprint.PrettyPrinter(indent=2, width=120)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gh7mAZrnQ1PY",
        "outputId": "0e71f6e5-2b77-4365-c245-090b70558e09"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.11/dist-packages/pydantic/_internal/_fields.py:198: UserWarning: Field name \"config_type\" in \"SequentialAgent\" shadows an attribute in parent \"BaseAgent\"\n",
            "  warnings.warn(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "if not PROJECT_ID:\n",
        "    raise ValueError('Please set your PROJECT_ID.')\n",
        "os.environ['GOOGLE_CLOUD_PROJECT'] = PROJECT_ID\n",
        "os.environ['GOOGLE_CLOUD_LOCATION'] = LOCATION\n",
        "os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION)"
      ],
      "metadata": {
        "id": "GpNakpiRibH4"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class RouterActionType(str, Enum):\n",
        "    NONE = \"NONE\"\n",
        "    BOOK_FLIGHT = \"BOOK_FLIGHT\"\n",
        "    BOOK_HOTEL = \"BOOK_HOTEL\"\n",
        "\n",
        "class RouterOutput(BaseModel):\n",
        "  message: str\n",
        "  next_step: RouterActionType\n",
        "  next_step_input: Optional[str] = None"
      ],
      "metadata": {
        "id": "RWg_66pBXt0R"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class RouterAgent:\n",
        "    \"\"\"An agent that determines whether to call internal API vs chit-chat\"\"\"\n",
        "\n",
        "    def __init__(self) -> None:\n",
        "        self._agent = self._build_agent()\n",
        "        self._user_id = 'remote_agent'\n",
        "        self._runner = Runner(\n",
        "            app_name=self._agent.name,\n",
        "            agent=self._agent,\n",
        "            session_service=InMemorySessionService(),\n",
        "        )\n",
        "\n",
        "    def _build_agent(self) -> LlmAgent:\n",
        "        \"\"\"Builds the LLM agent for the router agent.\"\"\"\n",
        "        return LlmAgent(\n",
        "            model='gemini-2.5-flash',\n",
        "            name='router_agent',\n",
        "            output_schema=RouterOutput,\n",
        "            instruction=\"\"\"\n",
        "    You are an agent who reponds to user queries on behalf of a booking company. The booking company can book flights, hotels & cars rentals.\n",
        "\n",
        "    Based on user query, you need to suggest next step. Follow below guidelines to choose below next step:\n",
        "    - BOOK_FLIGHT: If the user shows intent to book a flight.\n",
        "    - BOOK_HOTEL: If the user shows intent to book a hotel.\n",
        "    - Otherwise the next step is NONE.\n",
        "\n",
        "    Your reponses should be in JSON in below schema:\n",
        "    {{\n",
        "    \"next_step\": \"NONE | BOOK_FLIGHT | BOOK_HOTEL\",\n",
        "    \"next_step_input\": \"Optional. Not needed in case of next step is NONE. Relevant info from user converstaion required for the selceted next step.\",\n",
        "    \"message\": \"A user visible message, based on the suggested next step. Assume the suggested next step would be auto executed.\"\n",
        "    }}\n",
        "\n",
        "    \"\"\",\n",
        "        )\n",
        "\n",
        "    async def run(self, query, session_id) -> RouterOutput:\n",
        "        session = await self._runner.session_service.get_session(\n",
        "            app_name=self._agent.name,\n",
        "            user_id=self._user_id,\n",
        "            session_id=session_id,\n",
        "        )\n",
        "        content = types.Content(\n",
        "            role='user', parts=[types.Part.from_text(text=query)]\n",
        "        )\n",
        "        if session is None:\n",
        "            session = await self._runner.session_service.create_session(\n",
        "                app_name=self._agent.name,\n",
        "                user_id=self._user_id,\n",
        "                state={},\n",
        "                session_id=session_id,\n",
        "            )\n",
        "        async for event in self._runner.run_async(\n",
        "            user_id=self._user_id, session_id=session.id, new_message=content\n",
        "        ):\n",
        "            if event.is_final_response():\n",
        "                response = ''\n",
        "                if (\n",
        "                    event.content\n",
        "                    and event.content.parts\n",
        "                    and event.content.parts[0].text\n",
        "                ):\n",
        "                    response = '\\n'.join(\n",
        "                        [p.text for p in event.content.parts if p.text]\n",
        "                    )\n",
        "                return RouterOutput.model_validate_json(response)\n",
        "\n",
        "        raise Exception(\"Router failed\")\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "S4uV6pzLWDbq"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "my_agent = RouterAgent()\n",
        "\n",
        "async def run(query: str):\n",
        "    resp = await my_agent.run(query, str(uuid.uuid4()))\n",
        "    pp.pprint(resp)\n",
        "\n",
        "await run(\"what can you do\")\n",
        "await run(\"Book a flight from NY to SF\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gchfa202WXIT",
        "outputId": "8cdcd526-3d1c-4e60-ebde-2afd55180b85"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:google_adk.google.adk.agents.llm_agent:Invalid config for agent router_agent: output_schema cannot co-exist with agent transfer configurations. Setting disallow_transfer_to_parent=True, disallow_transfer_to_peers=True\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "RouterOutput(message='I can help you book flights, hotels, and car rentals. What are you looking to do today?', next_step=<RouterActionType.NONE: 'NONE'>, next_step_input=None)\n",
            "RouterOutput(message='Booking a flight from NY to SF.', next_step=<RouterActionType.BOOK_FLIGHT: 'BOOK_FLIGHT'>, next_step_input='from NY to SF')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class BookingAgentExecutor(AgentExecutor):\n",
        "    \"\"\"Reimbursement AgentExecutor Example.\"\"\"\n",
        "\n",
        "    def __init__(self) -> None:\n",
        "        self.router_agent = RouterAgent()\n",
        "\n",
        "    async def execute(\n",
        "        self,\n",
        "        context: RequestContext,\n",
        "        event_queue: EventQueue,\n",
        "    ) -> None:\n",
        "        query = context.get_user_input()\n",
        "        task = context.current_task\n",
        "\n",
        "        router_output = await self.router_agent.run(query, str(uuid.uuid4()))\n",
        "\n",
        "        if router_output.next_step == RouterActionType.NONE:\n",
        "            await event_queue.enqueue_event(new_agent_text_message(router_output.message, context_id=context.context_id))\n",
        "            return\n",
        "\n",
        "        # Time to create a task.\n",
        "        if not task:\n",
        "            task = new_task(context.message)\n",
        "            await event_queue.enqueue_event(task)\n",
        "\n",
        "        updater = TaskUpdater(event_queue, task.id, task.context_id)\n",
        "        await updater.update_status(\n",
        "            TaskState.working,\n",
        "            new_agent_text_message(router_output.message, context_id=context.context_id)\n",
        "        )\n",
        "\n",
        "        booking_response = ''\n",
        "\n",
        "        if router_output.next_step == RouterActionType.BOOK_FLIGHT:\n",
        "            booking_response = await self.book_flight()\n",
        "        elif router_output.next_step == RouterActionType.BOOK_HOTEL:\n",
        "            booking_response = await self.book_hotel()\n",
        "\n",
        "        await updater.add_artifact(\n",
        "            [Part(root=TextPart(text=booking_response))], name='Booking ID'\n",
        "        )\n",
        "        await updater.complete()\n",
        "\n",
        "    async def book_flight(self) -> str:\n",
        "        return \"PNR: FY1234\"\n",
        "\n",
        "    async def book_hotel(self) -> str:\n",
        "        return \"Hotel Reference No: H789\"\n",
        "\n",
        "    async def cancel(\n",
        "        self, request: RequestContext, event_queue: EventQueue\n",
        "    ) -> Task | None:\n",
        "        raise MethodNotImplementedError(\n",
        "            'ReimbursementAgentExecutor does not support cancel operation.'\n",
        "        )"
      ],
      "metadata": {
        "id": "x_mc6aG1QrrU"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "request_handler = DefaultRequestHandler(\n",
        "    agent_executor=BookingAgentExecutor(),\n",
        "    task_store=InMemoryTaskStore(),\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "M8XfszQ-R9WY",
        "outputId": "8dd34fbf-2877-4857-8857-d5a97e19a62c"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:google_adk.google.adk.agents.llm_agent:Invalid config for agent router_agent: output_schema cannot co-exist with agent transfer configurations. Setting disallow_transfer_to_parent=True, disallow_transfer_to_peers=True\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "async def send_message(query: str = \"hi\"):\n",
        "    task_id = None\n",
        "    context_id = None\n",
        "\n",
        "    user_message = Message(\n",
        "        role=Role.user,\n",
        "        parts=[Part(root=TextPart(text=query))],\n",
        "        message_id=str(uuid.uuid4()),\n",
        "        task_id=task_id,\n",
        "        context_id=context_id,\n",
        "    )\n",
        "    params=MessageSendParams(\n",
        "        message=user_message\n",
        "    )\n",
        "    response_stream = request_handler.on_message_send_stream(params=params)\n",
        "\n",
        "    async for ev in response_stream:\n",
        "        pp.pprint(ev.model_dump(exclude_none=True))"
      ],
      "metadata": {
        "id": "_IOA3tn0SPIR"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "await send_message(\"hey, could you help book my trip\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FqDfuihcVNLl",
        "outputId": "71b11bca-e33b-4974-cc3d-c45bddd40b26"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{ 'contextId': 'f6b13b9d-8b5c-48e5-bc77-c6d42b29a670',\n",
            "  'kind': 'message',\n",
            "  'messageId': '96bb0458-673d-4bb0-8e5f-49f4ba6696d8',\n",
            "  'parts': [ { 'kind': 'text',\n",
            "               'text': 'I can help you with that! Are you looking to book a flight, a hotel, or perhaps a car '\n",
            "                       'rental?'}],\n",
            "  'role': <Role.agent: 'agent'>}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "await send_message(\"book a flight from NY to SF\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Et4bss9Gfsf5",
        "outputId": "bea9637c-7fca-4827-9de0-3f90c988ba96"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{ 'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "  'history': [ { 'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "                 'kind': 'message',\n",
            "                 'messageId': '05200ee2-df85-491c-bd1b-8c430299313e',\n",
            "                 'parts': [{'kind': 'text', 'text': 'book a flight from NY to SF'}],\n",
            "                 'role': <Role.user: 'user'>,\n",
            "                 'taskId': '51cbe03d-904d-4d21-bf07-dce2525bb9cd'}],\n",
            "  'id': '51cbe03d-904d-4d21-bf07-dce2525bb9cd',\n",
            "  'kind': 'task',\n",
            "  'status': {'state': <TaskState.submitted: 'submitted'>}}\n",
            "{ 'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "  'final': False,\n",
            "  'kind': 'status-update',\n",
            "  'status': { 'message': { 'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "                           'kind': 'message',\n",
            "                           'messageId': '6b36bd5c-9d3a-4dcb-acb0-27f1e50cd5e7',\n",
            "                           'parts': [{'kind': 'text', 'text': 'OK. I am booking a flight from NY to SF for you.'}],\n",
            "                           'role': <Role.agent: 'agent'>},\n",
            "              'state': <TaskState.working: 'working'>,\n",
            "              'timestamp': '2025-08-15T22:02:43.064620+00:00'},\n",
            "  'taskId': '51cbe03d-904d-4d21-bf07-dce2525bb9cd'}\n",
            "{ 'artifact': { 'artifactId': '290ee2a8-fd6d-42ec-b3ed-633e1eb82209',\n",
            "                'name': 'Booking ID',\n",
            "                'parts': [{'kind': 'text', 'text': 'PNR: FY1234'}]},\n",
            "  'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "  'kind': 'artifact-update',\n",
            "  'taskId': '51cbe03d-904d-4d21-bf07-dce2525bb9cd'}\n",
            "{ 'contextId': '2480c8df-8569-4e10-8d9f-03c828d949ab',\n",
            "  'final': True,\n",
            "  'kind': 'status-update',\n",
            "  'status': {'state': <TaskState.completed: 'completed'>, 'timestamp': '2025-08-15T22:02:43.064753+00:00'},\n",
            "  'taskId': '51cbe03d-904d-4d21-bf07-dce2525bb9cd'}\n"
          ]
        }
      ]
    }
  ]
}