{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "7954bec8-6dba-453d-948e-74eada8c87b0", "metadata": { "id": "1Vs2C5ZBkrjmrHTfAyz77hpV" }, "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": "588192ab-3488-4ae2-bf0e-62c837deb91c", "metadata": { "id": "VoJZgH6X77se" }, "source": [ "# **Gemini Enterprise custom agent with prompt management**" ] }, { "cell_type": "markdown", "id": "aecb9f66-0a4c-423b-a56a-cd7464e6206d", "metadata": { "id": "qfDq3m0R8Fe1" }, "source": [ "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \"Google
Open in Colab\n", " \n", " 		\n", " \n", " \"Google
Open in Colab Enterprise\n", " \n", " 		\n", " \n", " \"Vertex
Open in Vertex AI Workbench\n", " \n", " 		\n", " \n", " \"GitHub
View on GitHub\n", " \n", "
\n", "\n", "\n", "\n", "
\n", "
\n", "
\n", "\n", "
\n", "\n", "Share to:\n", "\n", "\n", " \"LinkedIn\n", "\n", "\n", "\n", " \"Bluesky\n", "\n", "\n", "\n", " \"X\n", "\n", "\n", "\n", " \"Reddit\n", "\n", "\n", "\n", " \"Facebook\n", "" ] }, { "cell_type": "markdown", "id": "3f9aed85-8ea3-4f4b-ba16-0849a7f28519", "metadata": { "id": "UD1Cy9O88Zn9" }, "source": [ "| Authors |\n", "| --- |\n", "| Rupjit Chakraborty|\n", "| Sharmila Devi |\n", "| Parag Mhatre |" ] }, { "cell_type": "markdown", "id": "b06fa0ed-4dc1-4bc3-a9b1-e5a91fa9c508", "metadata": { "id": "28b524cdd168" }, "source": [ "## Overview : Custom agent with file context and prompt management in Gemini Enterprise\n", "\n", "This notebook demonstrates an end-to-end example of a custom agent integrated with Gemini Enterprise that can accept a file as context and generate a response using that. In this case, a simple example of a SQL query generator based on a given schema DDL has been demonstrated. We also demonstrate custom agent [prompt management using Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/prompt-classes)\n", "\n", "### Background\n", "\n", "Custom agents integrated with Gemini Enterprise by developers may require instruction modification based on the user choice. Also, the instructions given to agents via prompts, need to be maintainable via versioning. The prompt management sdk by Vertex AI can be used to manage prompts for agents at scale.\n", "\n", "### Business Application Scenario\n", "\n", "We demonstrate a simple scenario where a prompt is acquired by a custom agent deployed on Agent Engine via Vertex AI Prompt Management sdk. The end-user or business users can update the prompt whenever needed and the agent can just be redeployed to read the new prompt, without any changes in the code.\n", "\n", "### Notebook Overview\n", "This notebook guides you through the following key steps:\n", "\n", "**1. Creation of prompt using Vertex AI Prompt Management** \n", "    a. Create the prompt and get the prompt id\n", "\n", "**2. Creation of custom agent that uses the prompt in Vertex AI Prompt Management** \n", "    a. Create the custom agent using the prompt id \n", "    b. Test and Deploy the agent to Agent Engine\n", "\n", "**3. Integration of custom agent to Gemini Enterprise** \n", "    a. Register the agent with Gemini Enterprise and integrate it\n", "\n", "**4. Demonstration of prompt change for custom agent without re-deploying custom agent** \n", "    a. Make a change in prompt being used by the agent and demonstrate the change in output" ] }, { "cell_type": "markdown", "id": "3707f148-1816-4fc5-b60b-a1080685a70e", "metadata": { "id": "37f300fa5e86" }, "source": [ "## Imports" ] }, { "cell_type": "code", "execution_count": null, "id": "dd9e3aa7-6a19-48fb-8d6c-ca5e218dfd4e", "metadata": { "id": "352c647bd47d" }, "outputs": [], "source": [ "import vertexai\n", "from google.adk.agents import Agent\n", "from google.adk.agents.callback_context import CallbackContext\n", "from google.genai import types\n", "from vertexai import agent_engines\n", "from vertexai.preview import prompts, reasoning_engines\n", "from vertexai.preview.prompts import Prompt" ] }, { "cell_type": "markdown", "id": "81440e98-7541-42e6-9fd4-96e6546acc5b", "metadata": { "id": "9a02fc0443de" }, "source": [ "## Constants and Initializations" ] }, { "cell_type": "code", "execution_count": null, "id": "674934a9-d2f0-4644-881b-e6b2b2d9a8ea", "metadata": { "id": "4806e8208f1b" }, "outputs": [], "source": [ "# TODO for Developer: Update project name, location, bucket and schema path (for testing)\n", "PROJECT_ID = \"[your-project-id]\"\n", "PROJECT_NUMBER = \"[your-project-number]\"\n", "LOCATION = \"[your-location]\"\n", "STAGING_BUCKET = \"[your-bucket]\"\n", "SCHEMA = \"[your-schema-txt-containing-DDLs]\"\n", "\n", "vertexai.init(\n", " project=PROJECT_ID,\n", " location=LOCATION,\n", " staging_bucket=STAGING_BUCKET,\n", ")\n", "\n", "try:\n", " with open(SCHEMA) as f_handle:\n", " schema = f_handle.read()\n", "except FileNotFoundError:\n", " print(f\"Error: The schema file was not found at '{SCHEMA}'.\")\n", " print(\"Please update the 'SCHEMA' variable with the correct file path.\")\n", " schema = \"\" # Assign a default value to prevent further errors" ] }, { "cell_type": "markdown", "id": "4aa5342c-3b71-4212-931b-a4ac79d59295", "metadata": { "id": "bf35bbdf7278" }, "source": [ "## Prompt creation using Vertex AI Prompt Management sdk" ] }, { "cell_type": "code", "execution_count": null, "id": "ef0ffa5f-9357-4939-b30c-ff6c1bf75f03", "metadata": { "id": "7937127bcd34" }, "outputs": [], "source": [ "# Add a new prompt to the Prompt Management store via sdk\n", "# This can also be done via the UI\n", "# This task needs to be done only once\n", "\n", "sql_query_gen_prompt = \"\"\"\n", "Look carefully at the schema DDL statements and provide the SQL query for the user's question.\n", "Also provide a detailed explanation of the query generated.\n", "If you cannot find a proper schema return \"Please upload a schema or a valid schema\".\n", "Do not assume anything.\n", "\"\"\"\n", "\n", "sql_query_gen_prompt_obj = Prompt(\n", " prompt_name=\"SQL Generator\",\n", " prompt_data=sql_query_gen_prompt,\n", " model_name=\"gemini-2.0-flash-001\",\n", ")\n", "agent_prompt = prompts.create_version(prompt=sql_query_gen_prompt_obj)\n", "print(\"Prompt id = \", agent_prompt.prompt_id)" ] }, { "cell_type": "markdown", "id": "2aa80a71-7470-4c1d-b995-d7715ac69f31", "metadata": { "id": "37485861d327" }, "source": [ "![SQL generation prompt in Prompt Management\"](https://services.google.com/fh/files/blogs/prompt_mgmt.png)" ] }, { "cell_type": "code", "execution_count": null, "id": "975a9af7-dc54-44f2-bc55-eecbec15ebc2", "metadata": { "id": "a00bbd0b0a53" }, "outputs": [], "source": [ "# List all the prompts and copy the prompt_id of the recently created prompt\n", "# Alternatively we can get is as: pid = agent_prompt.prompt_id\n", "\n", "prompts.list()" ] }, { "cell_type": "markdown", "id": "812fddb2-5f2d-42dc-b9dd-526d8cbfff41", "metadata": { "id": "831950cda4d6" }, "source": [ "## Create Agent" ] }, { "cell_type": "code", "execution_count": null, "id": "321bbedf-2d72-4bff-b325-cc58d6393341", "metadata": { "id": "2aade5b3db6f" }, "outputs": [], "source": [ "\"\"\"SQL query generator agent\"\"\"\n", "\n", "# prompt id obtained after listing prompts\n", "PID = \"[your-prompt-id]\"\n", "\n", "\n", "def update_instructions_add_schema(callback_context: CallbackContext):\n", " \"\"\"Add the contents of the file to the agent context.\n", "\n", " Args:\n", " callback_context: The agent context.\n", " \"\"\"\n", " from vertexai.preview import prompts\n", "\n", " flag = 0\n", " # acquire the prompt added by an end user and add based instructions to it\n", " sql_query_gen_prompt = prompts.get(prompt_id=PID).prompt_data\n", " # add the prompt instruction to the agent context\n", " callback_context._invocation_context.agent.instruction = sql_query_gen_prompt\n", "\n", " # iterate over the callback_context object and check if file data is present\n", " for part in callback_context.user_content.parts:\n", " if hasattr(part, \"inline_data\") and getattr(part.inline_data, \"mime_type\", \"\"):\n", " blob = part.inline_data\n", " raw_bytes = blob.data\n", " flag = 1\n", "\n", " # update the context if there is any file data found\n", " if flag and raw_bytes and isinstance(raw_bytes, (bytes, bytearray)):\n", " uploaded_file_content = raw_bytes.decode(\"utf-8\")\n", " callback_context._invocation_context.agent.instruction += uploaded_file_content\n", "\n", "\n", "# root agent that does the SQL generation\n", "root_agent = Agent(\n", " name=\"SQL_Generation_Agent\",\n", " model=\"gemini-2.0-flash\",\n", " description=(\n", " \"Agent to convert natural language query to SQL based on a given schema DDL.\"\n", " ),\n", " before_agent_callback=update_instructions_add_schema,\n", " generate_content_config=types.GenerateContentConfig(temperature=0.01),\n", ")" ] }, { "cell_type": "markdown", "id": "ac1dafbb-f501-4ca8-8f49-955424571c22", "metadata": { "id": "bddb8ddae816" }, "source": [ "## Test agent locally" ] }, { "cell_type": "code", "execution_count": null, "id": "4481ebf6-f990-45ea-b601-9656a53ede46", "metadata": { "id": "c5d52171d161" }, "outputs": [], "source": [ "# create an AdkApp object for testing and deployment of agent\n", "\n", "app = reasoning_engines.AdkApp(\n", " agent=root_agent,\n", " enable_tracing=True,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "b46b9673-d026-4003-9000-6e929647a93a", "metadata": { "id": "4672a9e0575f" }, "outputs": [], "source": [ "# create a session for interacting with agent\n", "\n", "session = app.create_session(user_id=\"u_123\")\n", "session" ] }, { "cell_type": "code", "execution_count": null, "id": "5ec4b2e9-501e-4d28-a118-f9d1b52a4005", "metadata": { "id": "b2fa3fcc4244" }, "outputs": [], "source": [ "# build the user query object for getting response from agent\n", "\n", "query = \"list all benefits of every employee mapped to the employee name\"\n", "\n", "file_artifact = types.Part.from_bytes(\n", " mime_type=\"text/plain\", data=schema.encode(\"utf-8\")\n", ")\n", "contents = types.Content(\n", " role=\"user\", parts=[types.Part.from_text(text=query), file_artifact]\n", ")\n", "contents.model_dump()" ] }, { "cell_type": "code", "execution_count": null, "id": "071c49a4-0d8f-45bc-ab1e-0935514a9ccd", "metadata": { "id": "3707c84923f3" }, "outputs": [], "source": [ "# get output for the user query\n", "\n", "for event in app.stream_query(\n", " user_id=\"u_123\",\n", " session_id=session.id,\n", " message=contents.model_dump(),\n", "):\n", " print(event[\"content\"][\"parts\"][0][\"text\"])" ] }, { "cell_type": "markdown", "id": "ed2d5839-31cd-4acf-ac43-53b398af4c7c", "metadata": { "id": "f2165b1edbed" }, "source": [ "## Deploy agent to Agent Engine" ] }, { "cell_type": "code", "execution_count": null, "id": "24fe67c2-0239-4ac7-93b7-a85524c17ddc", "metadata": { "id": "0d4c781f4ad9" }, "outputs": [], "source": [ "# add the necessary dependencies and install agent to agent engine\n", "\n", "remote_app = agent_engines.create(\n", " display_name=\"SQL generator\",\n", " agent_engine=app,\n", " requirements=[\n", " \"google-adk (==1.5.0)\",\n", " \"google-genai (==1.24.0)\",\n", " \"pydantic (==2.11.7)\",\n", " \"google-cloud-aiplatform (==1.101.0)\",\n", " ],\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "a9e7131e-a7ba-4b91-a1bb-2096865a7581", "metadata": { "id": "fc1ab19e7f7e" }, "outputs": [], "source": [ "# get the reasoning engine id\n", "\n", "remote_app.resource_name" ] }, { "cell_type": "markdown", "id": "7e9ba57c-8e49-4a63-9b6c-9f67e79338d4", "metadata": { "id": "82981776ff62" }, "source": [ "## Register agent in Gemini Enterprise" ] }, { "cell_type": "code", "execution_count": null, "id": "d5593a82-ffdd-42d8-bd1b-7dd7dcf62c70", "metadata": { "id": "3635a286d182" }, "outputs": [], "source": [ "%%bash\n", "\n", "curl -X POST \\\n", "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n", "-H \"Content-Type: application/json\" \\\n", "-H \"X-Goog-User-Project: [your-project-id]\" \\\n", "\"https://discoveryengine.googleapis.com/v1alpha/projects/[your-project-id]/locations/global/authorizations?authorizationId=sqlgen2\" \\\n", "-d '{\n", "\"name\": \"projects/[your-project-id]/locations/global/authorizations/sqlgen2\",\n", "\"serverSideOauth2\": {\n", "\"clientId\": \"[UPDATE-CLIENT-ID]\",\n", "\"clientSecret\": \"[UPDATE-CLIENT-SECRET]\",\n", "\"authorizationUri\": \"https://accounts.google.com/o/oauth2/v2/auth?client_id=[UPDATE-CLIENT-ID]&scope=https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fcloud-platform&include_granted_scopes=true&response_type=code&access_type=offline&prompt=consent\",\n", "\"tokenUri\": \"https://oauth2.googleapis.com/token\"\n", "}\n", "}'" ] }, { "cell_type": "code", "execution_count": null, "id": "70b2af97-dc7d-4a58-a5a5-ceb3a82c4cb9", "metadata": { "id": "b33096f67166" }, "outputs": [], "source": [ "%%bash\n", "\n", "curl -X POST \\\n", "-H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n", "-H \"Content-Type: application/json\" \\\n", "-H \"X-Goog-User-Project: [your-project-id]\" \\\n", "\"https://discoveryengine.googleapis.com/v1alpha/projects/[your-project-id]/locations/global/collections/default_collection/engines/[your-gemini-enterprise-engine-id]/assistants/default_assistant/agents\" \\\n", "-d '{\n", "\"displayName\": \"DDL SQL Generator\",\n", "\"description\": \"SQL generation agent to help user with SQL queries for a give schema.\",\n", "\"adk_agent_definition\": {\n", "\"tool_settings\": {\n", "\"tool_description\": \"You are an expert SQL developer capable of understanding user queries and converting them to SQL\"\n", "},\n", "\"provisioned_reasoning_engine\": {\n", "\"reasoning_engine\": \"projects/[your-project-number]/locations/us-central1/reasoningEngines/[reasoning-engine-id]\"\n", "},\n", "\"authorizations\": [\"projects/[your-project-number]/locations/global/authorizations/sqlgen2\"]\n", "}\n", "}'" ] }, { "cell_type": "markdown", "id": "79faa8d1-3ee2-4268-8c59-ad1d237b3bcd", "metadata": { "id": "3e61ecea5040" }, "source": [ "## SQL Generator agent in Gemini Enterprise" ] }, { "cell_type": "markdown", "id": "0dce3cf5-189c-43c1-bc75-8cd769df4200", "metadata": { "id": "c18598400989" }, "source": [ "![DDL based SQL query generator agent](https://services.google.com/fh/files/blogs/sql_gen.png)" ] } ], "metadata": { "colab": { "name": "gemini_enterprise_prompt_management.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }