{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e339eff-e13f-48c6-99bc-1027a9e12341",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Copyright 2025 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "title-overview",
   "metadata": {},
   "source": [
    "# ADK Session with Cloud SQL\n",
    "\n",
    "This notebook demonstrates how to build and deploy an Agent Development Kit (ADK) agent that leverages different session services for conversation memory. Specifically, it covers:\n",
    "\n",
    "1.  Using a local SQLite database for quick development and testing.\n",
    "2.  Migrating to a production-ready Cloud SQL for PostgreSQL database for robust, scalable session management.\n",
    "\n",
    "It walks through setting up the agent, configuring both session types, and observing how conversation history is stored."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "setup-section",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "install-libraries-explanation",
   "metadata": {},
   "source": [
    "### Install Google Cloud Libraries\n",
    "\n",
    "This cell installs all the required Python libraries for this notebook:\n",
    "\n",
    "*   `google-cloud-aiplatform`: The client library for Vertex AI and Agent Engine.\n",
    "*   `google-adk`: The Agent Development Kit for building the agent.\n",
    "*   `pandas`: Used for displaying database query results in a clean table format.\n",
    "*   `pg8000`: A Python driver needed to connect to the PostgreSQL database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750ce816-d9c3-49c0-8d08-97251680e4f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q google-cloud-aiplatform[reasoning_engines] google-adk pandas pg8000\n",
    "\n",
    "print(\"\u2705 All required libraries are installed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "import-libraries-explanation",
   "metadata": {},
   "source": [
    "### Import Libraries\n",
    "\n",
    "This cell imports the necessary Python libraries and modules required to run the notebook, including components from the Google ADK and Vertex AI SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31e710bf-46d5-4c2a-8cda-02a572b5ac6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import vertexai\n",
    "import sqlite3\n",
    "import pandas as pd\n",
    "from google.adk.agents import Agent\n",
    "from google.adk.sessions import BaseSessionService, DatabaseSessionService\n",
    "from vertexai.preview.reasoning_engines import AdkApp\n",
    "from vertexai import agent_engines"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gcp-config-explanation",
   "metadata": {},
   "source": [
    "### Configure Google Cloud Environment\n",
    "\n",
    "This cell sets the core configuration for your Google Cloud environment. You need to replace the placeholder values below with your actual Project ID, Location, and the name of a Cloud Storage bucket for staging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e38a85cf-97d8-45ae-83d0-53df1ff88f3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = \"your-project-id\" # ACTION REQUIRED: Replace with your Project ID\n",
    "os.environ[\"GOOGLE_CLOUD_LOCATION\"] = \"us-central1\" # ACTION REQUIRED: Replace with your preferred region\n",
    "os.environ[\"STAGING_BUCKET\"] = \"gs://your-staging-bucket\" # ACTION REQUIRED: Replace with a GCS bucket for staging\n",
    "\n",
    "print(\"\u2705 Google Cloud environment variables configured.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "load-env-explanation",
   "metadata": {},
   "source": [
    "This cell loads the Google Cloud project ID, location, and staging bucket from the environment variables set in the previous cell into Python variables for use throughout the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "638bc946-df49-42bb-8372-b4b11bde842f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
    "LOCATION = os.environ[\"GOOGLE_CLOUD_LOCATION\"]\n",
    "STAGING_BUCKET = os.environ[\"STAGING_BUCKET\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vertex-ai-init-explanation",
   "metadata": {},
   "source": [
    "This cell initializes the Vertex AI SDK. It authenticates your environment and sets the default project, location, and staging bucket for all subsequent Vertex AI calls made within this session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49df3124-4c5c-48b7-8600-da2638e3443b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vertexai.init(\n",
    "    project=PROJECT_ID,\n",
    "    location=LOCATION,\n",
    "    staging_bucket=STAGING_BUCKET,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "define-agent-section",
   "metadata": {},
   "source": [
    "## Define the Agent\n",
    "\n",
    "This section defines the core Agent, which is the brain of our application. It combines a specific Gemini model with a detailed instruction prompt that defines its personality, goals, and constraints."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "define-root-agent-explanation",
   "metadata": {},
   "source": [
    "### Define the Root Agent\n",
    "\n",
    "This cell defines the core `Agent` instance that will be used throughout the notebook. It sets the agent's name, the generative model (`gemini-2.5-flash`), a brief description, and a detailed instruction prompt that dictates the agent's behavior and personality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c902ace1-bab3-469b-bf1e-32e81cb712ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simple_prompt = \"\"\"You are a helpful example bot that tries to answer anything the client wants.\n",
    "Help the client by answering his questions. Be respectful and kind, but if the client asks something bad feel free to deny it\"\"\"\n",
    "\n",
    "root_agent = Agent(\n",
    "    name=\"rag_agent\",\n",
    "    model=\"gemini-2.5-flash\",\n",
    "    description=(\n",
    "        \"Help the user by answering questions.\"\n",
    "    ),\n",
    "    instruction=simple_prompt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6616ab13-42bc-4987-bcbd-ab33745f41f1",
   "metadata": {},
   "source": [
    "## Option 1: Using SQLite (Local Development & Testing)\n",
    "\n",
    "To get started, we'll use the simplest method for storing conversation history: a local **SQLite** database.\n",
    "\n",
    "### Why Use SQLite?\n",
    "*   **Zero Setup**: It doesn't require any separate database server, authentication, or network configuration. The ADK handles everything for you.\n",
    "*   **File-Based**: Your entire database is a single file (`sessions.db`) that gets created in your project directory, making it easy to inspect or delete.\n",
    "*   **Perfect for Development**: It's the ideal choice for quickly building and testing your agent on your local machine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sqlite-builder-explanation",
   "metadata": {},
   "source": [
    "### Define the Session Service Builder for SQLite\n",
    "\n",
    "To enable the agent to remember conversations, it requires a session service. This function defines a builder that configures the agent to store conversation history in a local SQLite database file (`sessions.db`). This is ideal for quick local development and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6dcf0d-1b52-40d8-9c62-94221bcae022",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_database_session_service_sqlite() -> BaseSessionService:\n",
    "    \"\"\"\n",
    "    Creates and configures the DatabaseSessionService for SQLite.\n",
    "    This function acts as a \"builder\" that the AdkApp can call whenever\n",
    "    it needs to create a session service instance.\n",
    "    \"\"\"\n",
    "    db_connection_string = \"sqlite:///sessions.db\"\n",
    "\n",
    "    print(f\"Session service configured to use URL: {db_connection_string}\")\n",
    "    return DatabaseSessionService(db_url=db_connection_string)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-app-sqlite-explanation",
   "metadata": {},
   "source": [
    "### Create the AdkApp Instance with SQLite Session\n",
    "\n",
    "The AdkApp acts as the main entry point for interacting with our agent. We initialize it by combining our defined agent logic with the SQLite session management service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e2d155-0308-49fd-abd1-b9efc2d27eeb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "app = AdkApp(\n",
    "    agent=root_agent,\n",
    "    session_service_builder=create_database_session_service_sqlite,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-session-explanation",
   "metadata": {},
   "source": [
    "### Create a New Conversation Session\n",
    "\n",
    "Before we can send messages, a conversation session must be created. This step generates a unique session ID for a specified user, which will be used to track the conversation history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc64097f-95c2-4553-8d68-018178222cfb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "user_id = \"local_user_demo\"\n",
    "\n",
    "session = app.create_session(user_id=user_id)\n",
    "\n",
    "print(f\"Created a new session for user '{user_id}' with session ID: {session.id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "send-message-explanation",
   "metadata": {},
   "source": [
    "### Send a Message and Stream the Response\n",
    "\n",
    "This cell sends a user message to the agent within the created session and streams the agent's response. The conversation turn will be saved in the local SQLite database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91b9f927-567d-49c0-8952-d1fa74bc2a4f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"\\n--- User's Message ---\")\n",
    "message = \"Hello whats up?\"\n",
    "print(f\"You: {message}\\n\")\n",
    "\n",
    "print(\"--- Agent's Response ---\")\n",
    "for event in app.stream_query(user_id = user_id, \n",
    "                              session_id=session.id, \n",
    "                              message=message):\n",
    "    print(event['content'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sqlite-db-look-explanation",
   "metadata": {},
   "source": [
    "### Inspect the SQLite Database\n",
    "\n",
    "After the conversation, the DatabaseSessionService has persisted the history into the 'sessions.db' file. This section explores the tables created by the ADK to store session data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fcb4fd7-a955-4349-a083-59e7349433cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# We have the following tables: app_states, events, sessions, user_states\n",
    "# Lets take a look inside each one of them"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db-connection-setup-explanation",
   "metadata": {},
   "source": [
    "#### Establish Database Connection\n",
    "\n",
    "This cell establishes a connection to the local 'sessions.db' SQLite file using the `sqlite3` library, allowing us to query and inspect its contents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "470c4973-4a38-46ca-ad2b-809f1bcada1f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "DB_FILE = \"sessions.db\"\n",
    "\n",
    "pd.set_option('display.max_colwidth', None)\n",
    "\n",
    "conn = sqlite3.connect(DB_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "query-sessions-explanation",
   "metadata": {},
   "source": [
    "#### Query the 'sessions' Table\n",
    "\n",
    "The 'sessions' table stores the primary record for each conversation, including the full chat history between the user and the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22b9a925-d11f-4060-8825-9eb6ff00b798",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"--- Contents of the 'sessions' table ---\")\n",
    "sessions_df = pd.read_sql_query(\"SELECT * FROM sessions\", conn)\n",
    "display(sessions_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "query-events-explanation",
   "metadata": {},
   "source": [
    "#### Query the 'events' Table\n",
    "\n",
    "The 'events' table is used by the ADK to log specific actions or occurrences during agent request processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22dd6fd8-a0f0-40a4-b627-de5ee87507d4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"--- Contents of the 'events' table ---\")\n",
    "events_df = pd.read_sql_query(\"SELECT * FROM events\", conn)\n",
    "display(events_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "query-app-states-explanation",
   "metadata": {},
   "source": [
    "#### Query the 'app_states' Table\n",
    "\n",
    "The 'app_states' table stores the agent's internal application state or tool memory that needs to persist across turns, separate from the chat history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e13443-eb80-40e0-84f2-895d204510ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"--- Contents of the 'app_states' table ---\")\n",
    "app_states_df = pd.read_sql_query(\"SELECT * FROM app_states\", conn)\n",
    "display(app_states_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "query-user-states-explanation",
   "metadata": {},
   "source": [
    "#### Query the 'user_states' Table\n",
    "\n",
    "The 'user_states' table is designed to hold long-term memory about a specific user, such as preferences, which can be recalled across multiple conversations or sessions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad7c9993-8208-4c8e-8a02-637211bbaea2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"--- Contents of the 'user_states' table ---\")\n",
    "user_states_df = pd.read_sql_query(\"SELECT * FROM user_states\", conn)\n",
    "display(user_states_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4b1c3d2-e4f5-4678-9876-f5d4e3c2b1a0",
   "metadata": {},
   "source": [
    "## Option 2: Using a Cloud SQL Database (Production Setup)\n",
    "\n",
    "While SQLite is great for getting started, a managed cloud database like **Cloud SQL** is the right choice for a more robust or production-ready application.\n",
    "\n",
    "### Why Use Cloud SQL?\n",
    "*   **Scalability**: It can handle many more simultaneous connections than a local SQLite file.\n",
    "*   **Persistence & Reliability**: As a managed service, Google handles backups, failovers, and maintenance, making your data much safer.\n",
    "*   **Accessibility**: Your deployed applications (e.g., on Agent Engine) can connect to it from anywhere within your Google Cloud project.\n",
    "\n",
    "The following steps will guide you through creating a PostgreSQL instance on Cloud SQL using `gcloud` commands."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5c4d3e2-f6a7-4890-a1b2-c3d4e5f6a7b8",
   "metadata": {},
   "source": [
    "### Step 1: Enable APIs\n",
    "\n",
    "Before creating an instance, you need to ensure the necessary APIs are enabled for your project. The following command enables the Cloud SQL Admin API (for managing instances) and the Vertex AI API."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "enable-apis-explanation",
   "metadata": {},
   "source": [
    "This command enables the necessary Google Cloud APIs: the Cloud SQL Admin API for managing database instances and the Vertex AI API for agent deployment and interaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d5e4f3-a8b9-4c12-b3c4-d5e6f7a8b9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud services enable sqladmin.googleapis.com aiplatform.googleapis.com --project=$PROJECT_ID\n",
    "\n",
    "print(\"\u2705 APIs enabled.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7e6f5a4-b9c0-4d34-a5b6-c7d8e9f0a1b2",
   "metadata": {},
   "source": [
    "### Step 2: Create the Cloud SQL Instance\n",
    "\n",
    "Next, we'll create the Cloud SQL instance. This command will provision a new PostgreSQL database.\n",
    "\n",
    "*   **Instance ID**: A unique name for your instance, like `adk-agent-db`.\n",
    "*   **Password**: You will be prompted to set a strong password for the default `postgres` user. **Important: Save this password somewhere secure!** You will need it for your `.env` file.\n",
    "*   **Region**: To get the best performance, choose the **same region** where your other Vertex AI resources are located (e.g., `us-central1`).\n",
    "*   **Database Version**: We'll use PostgreSQL 13, but you can choose another supported version.\n",
    "*   **Tier**: The `db-g1-small` is a cost-effective option for development and testing.\n",
    "\n",
    "Replace `adk-agent-db` with your desired instance ID in the command below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-cloudsql-instance-explanation",
   "metadata": {},
   "source": [
    "This cell provisions a new PostgreSQL database instance on Cloud SQL. You will be prompted to set a strong password for the 'postgres' user. Ensure the instance ID is unique and the region matches your other Vertex AI resources. This command can take several minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f7a6b5-c0d1-4e56-b7c8-d9e0f1a2b3c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import shlex\n",
    "import os\n",
    "\n",
    "INSTANCE_ID = \"my-instance-name\" # ACTION REQUIRED: Replace with your desired instance ID\n",
    "REGION = LOCATION \n",
    "\n",
    "print(f\"Creating Cloud SQL instance '{INSTANCE_ID}'. This will take several minutes...\")\n",
    "!gcloud sql instances create {INSTANCE_ID} \\\n",
    "    --database-version=POSTGRES_13 \\\n",
    "    --tier=db-g1-small \\\n",
    "    --region={REGION} \\\n",
    "    --project={PROJECT_ID}\n",
    "print(\"Instance created.\")\n",
    "\n",
    "db_password = getpass.getpass(\"Enter a secure password for the 'postgres' user: \")\n",
    "\n",
    "os.environ[\"DB_PASSWORD\"] = db_password\n",
    "print(\"DB_PASSWORD environment variable has been set in this session.\")\n",
    "\n",
    "quoted_password = shlex.quote(db_password)\n",
    "\n",
    "!gcloud sql users set-password postgres \\\n",
    "    --instance={INSTANCE_ID} \\\n",
    "    --project={PROJECT_ID} \\\n",
    "    --password={quoted_password}\n",
    "\n",
    "print(f\"\u2705 Cloud SQL instance '{INSTANCE_ID}' created and password set for 'postgres' user.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9a8b7c6-d1e2-4f78-a9b0-c1d2e3f4a5b6",
   "metadata": {},
   "source": [
    "### Step 3: Create a Database\n",
    "\n",
    "Once your instance is running, you need to create a specific database inside it for your agent's sessions. We'll name it `agent_sessions`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-database-explanation",
   "metadata": {},
   "source": [
    "This command creates a dedicated database within your Cloud SQL instance to store the agent's session data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab9c8d7-e2f3-4089-badc-1d2e3f4a5b6c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "DATABASE_NAME = \"my-database-name\" # ACTION REQUIRED: Replace with your desired database name\n",
    "\n",
    "!gcloud sql databases create {DATABASE_NAME} \\\n",
    "    --instance={INSTANCE_ID} \\\n",
    "    --project={PROJECT_ID}\n",
    "\n",
    "print(f\"\u2705 Database '{DATABASE_NAME}' created in instance '{INSTANCE_ID}'.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bcad8e-f3a4-419a-aedc-2d3e4f5a6b7d",
   "metadata": {},
   "source": [
    "### Step 4: Get the Instance Connection Name\n",
    "\n",
    "The instance connection name is a unique identifier for your database instance that is required for connecting securely, especially with the Cloud SQL Auth Proxy. The following command retrieves this name."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "get-connection-name-explanation",
   "metadata": {},
   "source": [
    "This cell retrieves the unique instance connection name, which is essential for securely connecting to your Cloud SQL database via the Cloud SQL Auth Proxy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cdbe9f-a4b5-42ab-bfed-3e4f5a6b7c8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "CONNECTION_NAME_VAR = !gcloud sql instances describe {INSTANCE_ID} --project={PROJECT_ID} --format='value(connectionName)'\n",
    "CONNECTION_NAME = CONNECTION_NAME_VAR[0]\n",
    "\n",
    "print(f\"Instance Connection Name: {CONNECTION_NAME}\")\n",
    "print(\"ACTION: Copy this value and paste it into the 'Configure Cloud SQL Database Connection Details' cell below.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3defa0b-b5c6-43bc-80fd-4f5a6b7c8d9e",
   "metadata": {},
   "source": [
    "### Step 5: Download and Run the Cloud SQL Auth Proxy\n",
    "\n",
    "To connect your local machine to your Cloud SQL instance securely, you must use the **Cloud SQL Auth Proxy**. This tool creates a secure local connection without needing to configure IP allowlisting.\n",
    "\n",
    "The next cell automates the process of downloading the correct proxy for your operating system and starting it as a background process. The proxy must be running whenever you want to connect to your database from this notebook.\n",
    "\n",
    "**Note**: The proxy will run in the background. To stop it, you will need to interrupt or restart the notebook's kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "download-run-proxy-explanation",
   "metadata": {},
   "source": [
    "This cell downloads the Cloud SQL Auth Proxy for your operating system and starts it as a background process. The proxy establishes a secure, authenticated connection to your Cloud SQL instance without needing IP whitelisting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "new-proxy-cell-12345",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import platform\n",
    "import subprocess\n",
    "import time\n",
    "\n",
    "if 'CONNECTION_NAME' not in locals():\n",
    "    raise NameError(\"CONNECTION_NAME is not defined. Please run the previous cell to get the instance connection name.\")\n",
    "\n",
    "system = platform.system()\n",
    "machine = platform.machine()\n",
    "proxy_url = \"\"\n",
    "executable_name = \"cloud-sql-proxy\"\n",
    "\n",
    "if system == \"Linux\" and machine == \"x86_64\":\n",
    "    proxy_url = \"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.8.2/cloud-sql-proxy.linux.amd64\"\n",
    "elif system == \"Darwin\":\n",
    "    if machine == \"arm64\":\n",
    "        proxy_url = \"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.8.2/cloud-sql-proxy.darwin.arm64\"\n",
    "    elif machine == \"x86_64\":\n",
    "        proxy_url = \"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.8.2/cloud-sql-proxy.darwin.amd64\"\n",
    "elif system == \"Windows\":\n",
    "     proxy_url = \"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.8.2/cloud-sql-proxy.x64.exe\"\n",
    "     executable_name += \".exe\"\n",
    "\n",
    "if not proxy_url:\n",
    "    print(f\"Unsupported OS/architecture: {system}/{machine}. Please install the Cloud SQL Auth Proxy manually.\")\n",
    "else:\n",
    "    print(f\"Downloading Cloud SQL Auth Proxy for {system}/{machine}...\")\n",
    "    !curl -o {executable_name} {proxy_url}\n",
    "    \n",
    "    if system != \"Windows\":\n",
    "        print(\"Making it executable...\")\n",
    "        !chmod +x {executable_name}\n",
    "\n",
    "    print(f\"Starting proxy for instance: {CONNECTION_NAME}\")\n",
    "    proxy_process = subprocess.Popen([f\"./{executable_name}\", CONNECTION_NAME])\n",
    "    \n",
    "    time.sleep(5)\n",
    "    \n",
    "    print(f\"\u2705 Cloud SQL Auth Proxy started in the background with PID: {proxy_process.pid}\")\n",
    "    print(\"You can now proceed to the next cells to connect to the database.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "set-db-details-explanation",
   "metadata": {},
   "source": [
    "### Configure Cloud SQL Database Connection Details\n",
    "\n",
    "This cell sets up the necessary environment variables and Python variables for connecting to your Cloud SQL database. Ensure you replace the placeholder values for `DB_NAME`, `CONNECTION_NAME`, `DB_HOST`, and `DB_PORT` with your actual Cloud SQL details. The `DB_PASSWORD` should have been set interactively earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15f78f49-07a2-474b-bf86-d89631cae02e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ[\"DB_USER\"] = \"postgres\"\n",
    "os.environ[\"DB_NAME\"] = \"my-database-name\" # ACTION REQUIRED: Use the database name from Step 3\n",
    "os.environ[\"CONNECTION_NAME\"] = \"your-project-id:your-region:my-instance-name\" # ACTION REQUIRED: Paste Connection Name from Step 4\n",
    "os.environ[\"DB_HOST\"] = \"127.0.0.1\" # Use localhost since Cloud SQL Auth Proxy is running locally\n",
    "os.environ[\"DB_PORT\"] = \"5432\" # Default PostgreSQL port\n",
    "\n",
    "DB_USER = os.environ[\"DB_USER\"]\n",
    "DB_PASSWORD = os.environ[\"DB_PASSWORD\"]\n",
    "DB_NAME = os.environ[\"DB_NAME\"]\n",
    "CONNECTION_NAME = os.environ[\"CONNECTION_NAME\"]\n",
    "DB_HOST = os.environ[\"DB_HOST\"]\n",
    "DB_PORT = os.environ[\"DB_PORT\"]\n",
    "\n",
    "if not DB_PASSWORD:\n",
    "    print(\"\u26a0\ufe0f DB_PASSWORD is not set. Please re-run the 'Create the Cloud SQL Instance' cell and set the password.\")\n",
    "else:\n",
    "    print(\"\u2705 Cloud SQL environment variables configured.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cloudsql-builder-explanation",
   "metadata": {},
   "source": [
    "### Define the Session Service Builder for Cloud SQL\n",
    "\n",
    "This function defines the session service builder that connects to the Cloud SQL for PostgreSQL instance. It constructs a PostgreSQL connection string using the credentials and host details configured previously, enabling the agent to use Cloud SQL for conversation memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c734514-841c-4a8a-8e1e-c4c56183eade",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_database_session_service_sql() -> BaseSessionService:\n",
    "    db_url = f\"postgresql+pg8000://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}\"\n",
    "\n",
    "    session_db_kwargs = {\n",
    "        \"pool_recycle\": 3600,\n",
    "        \"echo\": False\n",
    "    }\n",
    "\n",
    "    print(f\"Creating DatabaseSessionService with Cloud SQL URL.\")\n",
    "    return DatabaseSessionService(db_url=db_url, **session_db_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-app-cloudsql-explanation",
   "metadata": {},
   "source": [
    "### Create the AdkApp Instance with Cloud SQL Session\n",
    "\n",
    "This cell creates a new `AdkApp` instance configured to use our Cloud SQL database for session management. By swapping the session service builder, all subsequent conversations with this `sql_app` will be stored in PostgreSQL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d92aaa30-f1be-4db8-a6f9-316e614f913d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_app = AdkApp(\n",
    "    agent=root_agent,\n",
    "    session_service_builder=create_database_session_service_sql,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "create-session-cloudsql-explanation",
   "metadata": {},
   "source": [
    "### Create a New Conversation Session (Cloud SQL)\n",
    "\n",
    "Similar to the SQLite setup, this step creates a new session specifically for interactions that will be backed by the Cloud SQL database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "765327bf-5e75-4afe-81d7-250432071daa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "user_id = \"sql_user_demo\"\n",
    "\n",
    "session = sql_app.create_session(user_id=user_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interact-cloudsql-explanation",
   "metadata": {},
   "source": [
    "### Interact with the Agent (Using Cloud SQL)\n",
    "\n",
    "This cell sends a message to the agent using the Cloud SQL-backed application instance. The conversation turn will be saved to your PostgreSQL database in Cloud SQL, demonstrating its use as a persistent memory store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "587c153e-118f-4d00-a662-1bc539ff081a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"\\n--- User's Message ---\")\n",
    "message = \"Hello whats up?\"\n",
    "print(f\"You: {message}\\n\")\n",
    "\n",
    "\n",
    "print(\"--- Agent's Response ---\")\n",
    "for event in sql_app.stream_query(user_id=user_id, session_id=session.id, message=message):\n",
    "    print(event['content'])"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "agentic_ssh",
   "name": "workbench-notebooks.m132",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m132"
  },
  "kernelspec": {
   "display_name": "agents (Local)",
   "language": "python",
   "name": "agentic_ssh"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}