agents/agentic_rag/notebooks/adk_app_testing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ADK Application Testing\n",
    "\n",
    "This notebook demonstrates how to test an ADK (Agent Development Kit) application.\n",
    "It covers both local and remote testing, both with Agent Engine and Cloud Run.\n",
    "\n",
    "<img src=\"https://github.com/GoogleCloudPlatform/agent-starter-pack/blob/main/docs/images/adk_logo.png?raw=true\" width=\"400\">\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import requests\n",
    "import vertexai.agent_engines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If you are using Agent Engine\n",
    "See more documentation at [Agent Engine Overview](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local Testing\n",
    "\n",
    "You can import directly the AgentEngineApp class within your environment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from app.agent import root_agent\n",
    "from app.agent_engine_app import AgentEngineApp\n",
    "\n",
    "agent_engine = AgentEngineApp(agent=root_agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for event in agent_engine.stream_query(message=\"hi!\", user_id=\"test\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remote Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace with your Agent Engine ID\n",
    "AGENT_ENGINE_ID = \"projects/PROJECT_ID/locations/us-central1/reasoningEngines/ENGINE_ID\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "remote_agent_engine = vertexai.agent_engines.get(AGENT_ENGINE_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for event in remote_agent_engine.stream_query(message=\"hi!\", user_id=\"test\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remote_agent_engine.register_feedback(\n",
    "    feedback={\n",
    "        \"score\": 5,\n",
    "        \"text\": \"Great response!\",\n",
    "        \"invocation_id\": \"test-invocation-123\",\n",
    "        \"user_id\": \"test\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If you are using Cloud Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Local Testing\n",
    "\n",
    "> You can run the application locally via the `make backend` command."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a session\n",
    " Create a new session with user preferences and state information\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = \"test_user_123\"\n",
    "session_id = \"test_session_456\"\n",
    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
    "\n",
    "session_url = f\"http://127.0.0.1:8000/apps/app/users/{user_id}/sessions/{session_id}\"\n",
    "headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
    "print(f\"Session creation status code: {session_response.status_code}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Send a message\n",
    "Send a message to the backend service and receive a streaming response\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_data = {\n",
    "    \"app_name\": \"app\",\n",
    "    \"user_id\": user_id,\n",
    "    \"session_id\": session_id,\n",
    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
    "    \"streaming\": True,\n",
    "}\n",
    "\n",
    "message_url = \"http://127.0.0.1:8000/run_sse\"\n",
    "message_response = requests.post(\n",
    "    message_url, headers=headers, json=message_data, stream=True\n",
    ")\n",
    "\n",
    "print(f\"Message send status code: {message_response.status_code}\")\n",
    "\n",
    "# Print streamed response\n",
    "for line in message_response.iter_lines():\n",
    "    if line:\n",
    "        line_str = line.decode(\"utf-8\")\n",
    "        if line_str.startswith(\"data: \"):\n",
    "            event_json = line_str[6:]\n",
    "            event = json.loads(event_json)\n",
    "            print(f\"Received event: {event}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Remote Testing\n",
    "\n",
    "For more information about authenticating HTTPS requests to Cloud Run services, see:\n",
    "[Cloud Run Authentication Documentation](https://cloud.google.com/run/docs/triggering/https-request)\n",
    "\n",
    "Remote testing involves using a deployed service URL instead of localhost.\n",
    "\n",
    "Authentication is handled using GCP identity tokens instead of local credentials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "ID_TOKEN = get_ipython().getoutput(\"gcloud auth print-identity-token -q\")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVICE_URL = \"YOUR_SERVICE_URL_HERE\"  # Replace with your Cloud Run service URL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You'll need to first create a Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_id = \"test_user_123\"\n",
    "session_id = \"test_session_456\"\n",
    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
    "\n",
    "session_url = f\"{SERVICE_URL}/apps/app/users/{user_id}/sessions/{session_id}\"\n",
    "headers = {\"Content-Type\": \"application/json\", \"Authorization\": f\"Bearer {ID_TOKEN}\"}\n",
    "\n",
    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
    "print(f\"Session creation status code: {session_response.status_code}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you will be able to send a message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "message_data = {\n",
    "    \"app_name\": \"app\",\n",
    "    \"user_id\": user_id,\n",
    "    \"session_id\": session_id,\n",
    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
    "    \"streaming\": True,\n",
    "}\n",
    "\n",
    "message_url = f\"{SERVICE_URL}/run_sse\"\n",
    "message_response = requests.post(\n",
    "    message_url, headers=headers, json=message_data, stream=True\n",
    ")\n",
    "\n",
    "print(f\"Message send status code: {message_response.status_code}\")\n",
    "\n",
    "# Print streamed response\n",
    "for line in message_response.iter_lines():\n",
    "    if line:\n",
    "        line_str = line.decode(\"utf-8\")\n",
    "        if line_str.startswith(\"data: \"):\n",
    "            event_json = line_str[6:]\n",
    "            event = json.loads(event_json)\n",
    "            print(f\"Received event: {event}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}