{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\"Open" ], "id": "view-in-github" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Copyright 2026 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." ], "id": "cfNFHkaMKA97" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Intro to Gemini Data Analytics" ], "id": "jD4y5MNU_qi4" }, { "cell_type": "markdown", "metadata": {}, "source": [ "| Author |\n", "| --- |\n", "| [Aditya Verma](https://github.com/vermaAstra) |" ], "id": "RKVaW_pICJ2R" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Background and Overview\n", "The **Conversational Analytics API** lets you chat with your BigQuery or Looker data anywhere, including embedded Looker dashboards, Slack and other chat apps, or even your own web applications. Your team members can get answers where they need them, when they need them, in the applications they use every day. You can find the [Colab example with the Python SDK here](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/gemini_data_analytics/intro_gemini_data_analytics_sdk.ipynb). See [documentation](https://cloud.google.com/gemini/docs/conversational-analytics-api/overview) for more details.\n", "\n", "Please provide feedback to conversational-analytics-api-feedback@google.com\n", "
\n", "### This notebook will help you\n", "1. Authenticate to Google Cloud\n", "2. Add data\n", "3. Perform agent operations (create, list, get, delete)\n", "4. Manage conversations (create, list, get, delete)\n", "5. Ask questions with your agent\n" ], "id": "3Z7ymg4eGyeU" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Get Started" ], "id": "d1divKxHBVyQ" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## API Enablement\n", "\n", "**Please fill in the billing_project form field with your own Google Cloud project. The project must have the following APIs enabled:**\n", "- [cloudaicompanion API](https://console.cloud.google.com/apis/library/cloudaicompanion.googleapis.com)\n", "- [Gemini Data Analytics API](https://console.cloud.google.com/apis/library/geminidataanalytics.googleapis.com)\n", "- [BQ API](https://console.cloud.google.com/marketplace/product/google/bigquery.googleapis.com)\n", "- [Dataform API](https://console.cloud.google.com/apis/library/dataform.googleapis.com)\n", "- [Agent Platform API](https://console.cloud.google.com/apis/library/aiplatform.googleapis.com)\n", "\n", "You may pass in any BigQuery project/dataset/table for which you have read permissions.\n" ], "id": "YGT_vK5nSrCB" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Auth and Imports" ], "id": "DZz5XBuNSRb0" }, { "cell_type": "code", "execution_count": null, "metadata": { "language": "python" }, "outputs": [], "source": [ "import json\n", "import textwrap\n", "\n", "import altair as alt\n", "from google.colab import auth\n", "from IPython.display import HTML, display\n", "import pandas as pd\n", "from pygments import formatters, highlight, lexers\n", "import requests\n", "\n", "auth.authenticate_user()\n", "\n", "access_token = !gcloud auth application-default print-access-token\n", "headers = {\n", " \"Authorization\": f\"Bearer {access_token[0]}\",\n", " \"Content-Type\": \"application/json\",\n", "}" ], "id": "Skzx_KNaG3c-" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Datasource and Billing project\n", "\n", "\n", "* Define your billing project within Google Cloud\n", "* Link your agent to data sources (BigQuery, Looker, Looker Studio)\n" ], "id": "5nEg33LQ4AmH" }, { "cell_type": "code", "execution_count": null, "metadata": { "language": "python" }, "outputs": [], "source": [ "# @title Billing Project and Prompt (System Instruction)\n", "\n", "# fmt: off\n", "billing_project = \"[your-project-id]\" # @param {type:\"string\"}\n", "\n", "# provide critical context for your Conversational Analytics Agent here\n", "system_instruction = \"Think like an Analyst\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "bigquery_data_sources = {}\n", "looker_data_sources = {}\n", "looker_studio_data_sources = {}" ], "id": "a1fb79c902a9" }, { "cell_type": "code", "execution_count": null, "metadata": { "language": "python", "cellView": "form" }, "outputs": [], "source": [ "# @title BigQuery Datasource\n", "\n", "bigquery_data_sources = {\n", " \"bq\": {\n", " \"tableReferences\": [\n", " {\n", " \"projectId\": \"bigquery-public-data\",\n", " \"datasetId\": \"faa\",\n", " \"tableId\": \"us_airports\",\n", " },\n", " {\n", " \"projectId\": \"bigquery-public-data\",\n", " \"datasetId\": \"san_francisco\",\n", " \"tableId\": \"street_trees\",\n", " },\n", " # Add more table references here\n", " ]\n", " }\n", "}\n", "\n", "# Optional example queries (only leveraged for BigQuery datasources currently)\n", "# Example queries can be parameterized in which case parameters must be defined.\n", "# Each parameter requires a name (exactly as it appears in the question and SQL),\n", "# description, and a valid dataType\n", "# (https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_list).\n", "example_queries = [\n", " {\n", " \"naturalLanguageQuestion\": (\n", " \"What is the highest observed positive longitude?\"\n", " ),\n", " \"sqlQuery\": \"SELECT MAX(longitude) FROM airports\",\n", " },\n", " {\n", " \"naturalLanguageQuestion\": (\n", " \"How many airports are in {state} with elevation\"\n", " \" greater than {elevation}?\"\n", " ),\n", " \"sqlQuery\": (\n", " \"SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports` WHERE\"\n", " \" LOWER(state_abbreviation) = @state AND elevation > @elevation\"\n", " ),\n", " \"parameters\": [\n", " {\n", " \"name\": \"state\",\n", " \"dataType\": \"STRING\",\n", " \"description\": \"State abbreviation in lower case\",\n", " },\n", " {\n", " \"name\": \"elevation\",\n", " \"dataType\": \"FLOAT64\",\n", " \"description\": \"Elevation in feet\",\n", " },\n", " ],\n", " },\n", "]\n", "\n", "glossary_terms = [\n", " {\n", " \"display_name\": \"Airport Code\",\n", " \"description\": (\n", " \"A unique identifier for an airport, represented by a 3 character\"\n", " \" long identifier (the IATA code). E.g. 'JFK' for New York.\"\n", " ),\n", " \"labels\": [\"code\", \"IATA\", \"identifier\"],\n", " },\n", "]\n", "\n", "use_glossary_terms = True # @param {type: \"boolean\"}\n", "use_example_queries = True # @param {type:\"boolean\"}" ], "id": "c46d107e7768" }, { "cell_type": "code", "execution_count": null, "metadata": { "language": "python", "cellView": "form" }, "outputs": [], "source": [ "# @title Looker Datasource\n", "\n", "\"\"\"When connecting to the Looker datasource, you can authenticate using either:\n", "\n", "# 1. client_id and client_secret (supported with only PUBLIC instance)\n", "# 2. access_token (supported with both PUBLIC and PRIVATE instance)\n", "\"\"\"\n", "\n", "client_id = \"\" # @param {type:\"string\"}\n", "client_secret = \"\" # @param {type:\"string\"}\n", "looker_credentials = {\n", " \"oauth\": {\n", " \"secret\": {\n", " \"client_id\": client_id,\n", " \"client_secret\": client_secret,\n", " }\n", " }\n", "}\n", "\n", "# Uncomment this and comment out the above one if you are using access_token for authentication\n", "# access_token = \"\" # @param {type:\"string\"}\n", "# looker_credentials = {\n", "# \"oauth\": {\n", "# \"token\": {\n", "# \"access_token\": access_token,\n", "# }\n", "# }\n", "# }\n", "\n", "# Looker datasource 1\n", "lookml_model_1 = \"\" # @param {type:\"string\"}\n", "explore_1 = \"\" # @param {type:\"string\"}\n", "\n", "# Looker datasource 2 (optional)\n", "# lookml_model_2 = \"\" # @param {type:\"string\"}\n", "# explore_2 = \"\" # @param {type:\"string\"}\n", "\n", "\n", "\"\"\" Looker datasources for PUBLIC Instance \"\"\"\n", "\n", "# Required only for Looker PUBLIC instance\n", "looker_instance_uri_1 = \"https://my_company.looker.com\" # @param {type:\"string\"}\n", "# (Optional)\n", "# looker_instance_uri_2 = \"https://my_company.looker.com\" # @param {type:\"string\"}\n", "\n", "looker_data_sources = {\n", " \"looker\": {\n", " \"explore_references\": [\n", " {\n", " \"looker_instance_uri\": looker_instance_uri_1,\n", " \"lookml_model\": lookml_model_1,\n", " \"explore\": explore_1,\n", " },\n", " # {\n", " # \"looker_instance_uri\": looker_instance_uri_2,\n", " # \"lookml_model\": lookml_model_2,\n", " # \"explore\": explore_2,\n", " # },\n", " # Add up to 5 total explore references here\n", " ],\n", " }\n", "}\n", "\n", "\n", "\"\"\" Looker datasource for PRIVATE Instance \"\"\"\n", "\n", "# Required only for Looker PRIVATE instance\n", "# looker_instance_id_1 = \"\" # @param {type:\"string\"}\n", "# looker_service_directory_name_1 = \"\" # @param {type:\"string\"}\n", "\n", "# (Optional)\n", "# looker_instance_id_2 = \"\" # @param {type:\"string\"}\n", "# looker_service_directory_name_2 = \"\" # @param {type:\"string\"}\n", "\n", "# looker_data_sources = {\n", "# \"looker\": {\n", "# \"explore_references\": [\n", "# {\n", "# \"private_looker_instance_info\": {\n", "# \"looker_instance_id\": looker_instance_id_1,\n", "# \"service_directory_name\": looker_service_directory_name_1,\n", "# },\n", "# \"lookml_model\": lookml_model_1,\n", "# \"explore\": explore_1,\n", "# },\n", "# # {\n", "# # \"private_looker_instance_info\": {\n", "# # \"looker_instance_id\": looker_instance_id_2,\n", "# # \"service_directory_name\": looker_service_directory_name_2,\n", "# # },\n", "# # \"lookml_model\": lookml_model_2,\n", "# # \"explore\": explore_2,\n", "# # },\n", "# . # Add up to 5 total explore references here\n", "# ],\n", "# }\n", "\n", "\n", "looker_golden_queries = [\n", " {\n", " \"natural_language_questions\": (\n", " \"What is the highest observed positive longitude?\"\n", " ),\n", " \"looker_query\": {\n", " \"model\": \"airports\",\n", " \"explore\": \"airports\",\n", " \"fields\": \"airports.longitude\",\n", " \"filters\": {\n", " \"field\": \"airports.longitude\",\n", " \"value\": \">0\",\n", " },\n", " \"sorts\": \"airports.longitude desc\",\n", " \"limit\": \"1\",\n", " },\n", " },\n", " {\n", " \"natural_language_questions\": [\n", " \"What are the major airport codes and cities in CA?\",\n", " \"Can you list the cities and airport codes of airports in CA?\",\n", " ],\n", " \"looker_query\": {\n", " \"model\": \"airports\",\n", " \"explore\": \"airports\",\n", " \"fields\": \"airports.city\",\n", " \"fields\": \"airports.code\",\n", " \"filters\": {\n", " \"field\": \"airports.major\",\n", " \"value\": \"Y\",\n", " },\n", " \"filters\": {\n", " \"field\": \"airports.state\",\n", " \"value\": \"CA\",\n", " },\n", " },\n", " },\n", "]\n", "\n", "use_looker_golden_queries = False # @param {type:\"boolean\"}" ], "id": "7ae7bc48aa40" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Looker Studio Datasource\n", "\n", "looker_studio_data_sources = {\n", " \"studio\": {\n", " \"studio_references\": [{\"datasource_id\": \"your_studio_datasource_id\"}]\n", " }\n", "}" ], "id": "61cf3fddad05" }, { "metadata": { "cellView": "form" }, "cell_type": "code", "source": [ "# @title Select datasource\n", "\n", "selected_datasource = \"bigquery_data_sources\" # @param [\"bigquery_data_sources\", \"looker_data_sources\", \"looker_studio_data_sources\"]\n", "\n", "datasource_map = {\n", " \"bigquery_data_sources\": bigquery_data_sources,\n", " \"looker_data_sources\": looker_data_sources,\n", " \"looker_studio_data_sources\": looker_studio_data_sources,\n", "}\n", "\n", "datasource_references = datasource_map[selected_datasource]\n", "\n", "bq_selected = False\n", "looker_selected = False\n", "looker_studio_selected = False\n", "\n", "if selected_datasource == \"looker_data_sources\":\n", " looker_selected = True\n", "elif selected_datasource == \"bigquery_data_sources\":\n", " bq_selected = True\n", "elif selected_datasource == \"looker_studio_data_sources\":\n", " looker_studio_selected = True" ], "outputs": [], "execution_count": null, "id": "5UCbL-xfUNKZ" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Environment" ], "id": "66e0d8583aa7" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# fmt: off\n", "environment = \"prod-global\" # @param [\"prod-global\", \"prod-us\", \"prod-eu\", \"prod-us-east4\", \"staging-global\", \"autopush-global\"]\n", "location = \"\"\n", "api_version = \"v1\"\n", "base_url = \"\"\n", "# fmt: on\n", "\n", "if environment == \"prod-us\":\n", " base_url = \"https://geminidataanalytics.us.rep.googleapis.com\"\n", " location = \"us\"\n", "elif environment == \"prod-eu\":\n", " base_url = \"https://geminidataanalytics.eu.rep.googleapis.com\"\n", " location = \"eu\"\n", "elif environment == \"prod-us-east4\":\n", " base_url = \"https://geminidataanalytics-us-east4.googleapis.com\"\n", " location = \"us-east4\"\n", "elif environment == \"autopush-global\":\n", " base_url = \"https://autopush-geminidataanalytics.sandbox.googleapis.com\"\n", " location = \"global\"\n", "elif environment == \"staging-global\":\n", " base_url = \"https://staging-geminidataanalytics.sandbox.googleapis.com\"\n", " location = \"global\"\n", "else:\n", " base_url = \"https://geminidataanalytics.googleapis.com\"\n", " location = \"global\"" ], "id": "OE6bPKsyw_Jy" }, { "cell_type": "markdown", "source": [ "# Helper Functions\n", "These functions help parse the JSON response and present the data in a cleaner, more readable format." ], "metadata": {}, "id": "GsRvHc1dVOck" }, { "cell_type": "code", "source": [ "# @title Parser for List methods\n", "\n", "def display_resource_list(resource_list):\n", " for i, resource in enumerate(resource_list):\n", " resource_name = resource.get(\"name\").split(\"/\")[-1]\n", " resource_display_name = resource.get(\"displayName\")\n", " formatted_json = json.dumps(resource, indent=2)\n", " div_id = f\"resource_details_{i}\"\n", "\n", " resource_heading = f\"\"\"{resource_name} ( {resource_display_name} )\"\"\" if resource_display_name else resource_name\n", "\n", " display(HTML(f'''\n", "
\n", " {resource_heading}\n", " 
{formatted_json}
\n", "
\n", " '''))" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [], "id": "ByhfACPIVQXd" }, { "cell_type": "code", "source": [ "# @title Streaming Chat Messages\n", "\n", "\n", "def is_json(str):\n", " try:\n", " json_object = json.loads(str)\n", " except ValueError:\n", " return False\n", " return True\n", "\n", "\n", "def handle_text_response(resp):\n", " parts = resp[\"parts\"]\n", " full_text = \"\".join(parts)\n", " if \"\\n\" not in full_text and len(full_text) > 80:\n", " wrapped_text = textwrap.fill(full_text, width=80)\n", " print(wrapped_text)\n", " else:\n", " print(full_text)\n", "\n", "\n", "def get_property(data, field_name, default=\"\"):\n", " return data[field_name] if field_name in data else default\n", "\n", "\n", "def display_schema(data):\n", " fields = data[\"fields\"]\n", " df = pd.DataFrame({\n", " \"Column\": map(lambda field: get_property(field, \"name\"), fields),\n", " \"Type\": map(lambda field: get_property(field, \"type\"), fields),\n", " \"Description\": map(\n", " lambda field: get_property(field, \"description\", \"-\"), fields\n", " ),\n", " \"Mode\": map(lambda field: get_property(field, \"mode\"), fields),\n", " })\n", " display(df)\n", "\n", "\n", "def display_section_title(text):\n", " display(HTML(f\"

{text}

\"))\n", "\n", "\n", "def format_bq_table_ref(table_ref):\n", " return \"{}.{}.{}\".format(\n", " table_ref[\"projectId\"], table_ref[\"datasetId\"], table_ref[\"tableId\"]\n", " )\n", "\n", "\n", "def format_looker_table_ref(table_ref):\n", " return \"lookmlModel: {}, explore: {}\".format(\n", " table_ref[\"lookmlModel\"], table_ref[\"explore\"]\n", " )\n", "\n", "\n", "def display_datasource(datasource):\n", " source_name = \"\"\n", "\n", " if \"studioDatasourceId\" in datasource:\n", " source_name = datasource[\"studioDatasourceId\"]\n", " elif \"lookerExploreReference\" in datasource:\n", " source_name = format_looker_table_ref(datasource[\"lookerExploreReference\"])\n", " else:\n", " source_name = format_bq_table_ref(datasource[\"bigqueryTableReference\"])\n", "\n", " print(source_name)\n", " if \"schema\" in datasource:\n", " display_schema(datasource[\"schema\"])\n", "\n", "\n", "def handle_schema_response(resp):\n", " if \"query\" in resp and \"question\" in resp[\"query\"]:\n", " print(resp[\"query\"][\"question\"])\n", " elif \"result\" in resp:\n", " display_section_title(\"Schema resolved\")\n", " print(\"Data sources:\")\n", " for datasource in resp[\"result\"][\"datasources\"]:\n", " display_datasource(datasource)\n", "\n", "\n", "def handle_data_response(resp):\n", " if \"query\" in resp:\n", " query = resp[\"query\"]\n", " display_section_title(\"Retrieval query\")\n", " if \"name\" in query:\n", " print(\"Query name: {}\".format(query[\"name\"]))\n", " if \"question\" in query:\n", " print(\"Question: {}\".format(query[\"question\"]))\n", " if \"datasources\" in query:\n", " print(\"Data sources:\")\n", " for datasource in query[\"datasources\"]:\n", " display_datasource(datasource)\n", " elif \"generatedSql\" in resp:\n", " display_section_title(\"SQL generated\")\n", " print(resp[\"generatedSql\"])\n", " elif \"result\" in resp:\n", " display_section_title(\"Data retrieved\")\n", "\n", " fields = map(\n", " lambda field: get_property(field, \"name\"),\n", " resp[\"result\"][\"schema\"][\"fields\"],\n", " )\n", " dict = {}\n", "\n", " for field in fields:\n", " dict[field] = [get_property(el, field) for el in resp[\"result\"][\"data\"]]\n", "\n", " display(pd.DataFrame(dict))\n", "\n", "\n", "def handle_chart_response(resp):\n", " if \"query\" in resp and \"instructions\" in resp[\"query\"]:\n", " print(resp[\"query\"][\"instructions\"])\n", " elif \"result\" in resp:\n", " vegaConfig = resp[\"result\"][\"vegaConfig\"]\n", " alt.Chart.from_json(json.dumps(vegaConfig)).display()\n", "\n", "\n", "def handle_error(resp):\n", " display_section_title(\"Error\")\n", " print(\"Code: {}\".format(resp[\"code\"]))\n", " print(\"Message: {}\".format(resp[\"message\"]))\n", "\n", "\n", "def get_stream(url, payload):\n", " s = requests.Session()\n", "\n", " acc = \"\"\n", "\n", " with s.post(url, json=payload, headers=headers, stream=True) as resp:\n", " for line in resp.iter_lines():\n", " if not line:\n", " continue\n", "\n", " decoded_line = str(line, encoding=\"utf-8\")\n", "\n", " if decoded_line == \"[{\":\n", " acc = \"{\"\n", " elif decoded_line == \"}]\":\n", " acc += \"}\"\n", " elif decoded_line == \",\":\n", " continue\n", " else:\n", " acc += decoded_line\n", "\n", " if not is_json(acc):\n", " continue\n", "\n", " data_json = json.loads(acc)\n", "\n", " if \"systemMessage\" not in data_json:\n", " if \"error\" in data_json:\n", " handle_error(data_json[\"error\"])\n", " continue\n", "\n", " if \"text\" in data_json[\"systemMessage\"]:\n", " handle_text_response(data_json[\"systemMessage\"][\"text\"])\n", " elif \"schema\" in data_json[\"systemMessage\"]:\n", " handle_schema_response(data_json[\"systemMessage\"][\"schema\"])\n", " elif \"data\" in data_json[\"systemMessage\"]:\n", " handle_data_response(data_json[\"systemMessage\"][\"data\"])\n", " elif \"chart\" in data_json[\"systemMessage\"]:\n", " handle_chart_response(data_json[\"systemMessage\"][\"chart\"])\n", " else:\n", " colored_json = highlight(\n", " acc, lexers.JsonLexer(), formatters.TerminalFormatter()\n", " )\n", " print(colored_json)\n", " print(\"\\n\")\n", " acc = \"\"" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [], "id": "m47eL3MQVVN1" }, { "cell_type": "code", "source": [ "# @title get_stream function for multi-turn conversation\n", "\n", "# NOTE: This methods is same as get_stream() method present in Streaming Chat Messages section.\n", "# The only difference is that - here we are storing the response in an array \"conversation_messages\" to save the conversation.\n", "\n", "\n", "def get_stream_multi_turn(url, payload, conversation_messages):\n", " s = requests.Session()\n", "\n", " acc = \"\"\n", "\n", " with s.post(url, json=payload, headers=headers, stream=True) as resp:\n", " for line in resp.iter_lines():\n", " if not line:\n", " continue\n", "\n", " decoded_line = str(line, encoding=\"utf-8\")\n", "\n", " if decoded_line == \"[{\":\n", " acc = \"{\"\n", " elif decoded_line == \"}]\":\n", " acc += \"}\"\n", " elif decoded_line == \",\":\n", " continue\n", " else:\n", " acc += decoded_line\n", "\n", " if not is_json(acc):\n", " continue\n", "\n", " data_json = json.loads(acc)\n", " # Store the response to be used in next iteration.\n", " conversation_messages.append(data_json)\n", "\n", " if \"systemMessage\" not in data_json:\n", " if \"error\" in data_json:\n", " handle_error(data_json[\"error\"])\n", " continue\n", "\n", " if \"text\" in data_json[\"systemMessage\"]:\n", " handle_text_response(data_json[\"systemMessage\"][\"text\"])\n", " elif \"schema\" in data_json[\"systemMessage\"]:\n", " handle_schema_response(data_json[\"systemMessage\"][\"schema\"])\n", " elif \"data\" in data_json[\"systemMessage\"]:\n", " handle_data_response(data_json[\"systemMessage\"][\"data\"])\n", " elif \"chart\" in data_json[\"systemMessage\"]:\n", " handle_chart_response(data_json[\"systemMessage\"][\"chart\"])\n", " else:\n", " colored_json = highlight(\n", " acc, lexers.JsonLexer(), formatters.TerminalFormatter()\n", " )\n", " print(colored_json)\n", " print(\"\\n\")\n", " acc = \"\"" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [], "id": "W2CnfYcTVXoD" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Data Agent Services" ], "id": "52oJ2poOfHCp" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Create Data Agent\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents\"\n", "\n", "data_agent_payload = {\n", " \"name\": (\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " ), # Optional\n", " \"description\": \"This is the description of data_agent.\", # Optional\n", " \"data_analytics_agent\": {\n", " \"published_context\": {\n", " \"datasource_references\": datasource_references,\n", " \"system_instruction\": system_instruction,\n", " \"options\": {\n", " \"analysis\": {\n", " # Optional - if wanting to use advanced analysis with python.\n", " # Default is False.\n", " \"python\": {\"enabled\": False}\n", " }\n", " },\n", " }\n", " },\n", "}\n", "\n", "if bq_selected:\n", " if use_example_queries:\n", " data_agent_payload[\"data_analytics_agent\"][\"published_context\"][\n", " \"example_queries\"\n", " ] = example_queries\n", " if use_glossary_terms:\n", " data_agent_payload[\"data_analytics_agent\"][\"published_context\"][\n", " \"glossary_terms\"\n", " ] = glossary_terms\n", "elif looker_selected and use_looker_golden_queries:\n", " data_agent_payload[\"data_analytics_agent\"][\"published_context\"][\n", " \"looker_golden_queries\"\n", " ] = looker_golden_queries\n", "\n", "params = {\"data_agent_id\": data_agent_id} # Optional\n", "\n", "data_agent_response = requests.post(\n", " data_agent_url, params=params, json=data_agent_payload, headers=headers\n", ")\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Data Agent created successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error creating Data Agent: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "RuYxbQSmjSeD" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Get Data Agent\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", "\n", "data_agent_response = requests.get(data_agent_url, headers=headers)\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Fetched Data Agent successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "gKeVQqsGkyqT" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title List Data Agents\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents\"\n", "\n", "params = {\n", " \"orderBy\": \"create_time desc\", # Optional\n", "}\n", "data_agent_response = requests.get(data_agent_url, headers=headers, params=params)\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Data Agents Listed successfully!\")\n", " agents = data_agent_response.json().get(\"dataAgents\", [])\n", " display_resource_list(agents)\n", "else:\n", " print(f\"Error Listing Data Agents: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "VP2oNORjo4bc" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title List Accessible Data Agents\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents:listAccessible\"\n", "\n", "# Choose the desired creator filter\n", "# fmt: off\n", "creator_filter = \"NONE\" # @param [\"NONE\", \"CREATOR_ONLY\", \"NOT_CREATOR_ONLY\"]\n", "# fmt: on\n", "\n", "# Add the creator_filter as a query parameter\n", "params = {\"creator_filter\": creator_filter}\n", "\n", "data_agent_response = requests.get(\n", " data_agent_url, headers=headers, params=params\n", ")\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Data Agents Listed successfully!\")\n", " agents = data_agent_response.json().get(\"dataAgents\", [])\n", " display_resource_list(agents)\n", "else:\n", " print(f\"Error Listing Data Agents: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "be0ebe371ad7" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Update Data Agent\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", "\n", "payload = {\n", " \"description\": \"This is the latest description of data_agent.\",\n", " \"data_analytics_agent\": {\n", " \"published_context\": {\n", " \"datasource_references\": datasource_references,\n", " \"system_instruction\": system_instruction,\n", " }\n", " },\n", "}\n", "\n", "fields = [\"description\", \"data_analytics_agent\"]\n", "params = {\"updateMask\": \",\".join(fields)}\n", "\n", "data_agent_response = requests.patch(\n", " data_agent_url, headers=headers, params=params, json=payload\n", ")\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Data Agent updated successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error Updating Data Agent: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "-S7AIL1OeXnu" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title [Agent Sharing] Set IAM Policy for Data Agent\n", "\n", "\n", "\"\"\"PLEASE NOTE THIS API CALLS OVERRIDES EXISTING PERMISSION FOR THE RESOURCE.\n", "\n", "For preserving existing policy in practise call Get IAM policy to fetch existing\n", "policy and pass it along with additional changes in the call to Set IAM Policy\n", "\"\"\"\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "users = \"abc@google.com, wxyz@google.com\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:setIamPolicy\"\n", "\n", "\n", "# Request body\n", "payload = {\n", " \"policy\": {\n", " \"bindings\": [{\n", " \"role\": \"roles/geminidataanalytics.dataAgentEditor\",\n", " \"members\": [f\"user:{i.strip()}\" for i in users.split(\",\")],\n", " }]\n", " }\n", "}\n", "\n", "data_agent_response = requests.post(\n", " data_agent_url, headers=headers, json=payload\n", ")\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"IAM Policy set successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error setting IAM policy: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "TjZfg2lVF8Jl" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title [Agent Sharing] Get IAM Policy for Data Agent\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}:getIamPolicy\"\n", "\n", "# Request body\n", "payload = {\n", " \"resource\": (\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " )\n", "}\n", "\n", "data_agent_response = requests.post(\n", " data_agent_url, headers=headers, json=payload\n", ")\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"IAM Policy fetched successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error fetching IAM policy: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "kHt_lhjh3LtS" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title [Soft Delete] Delete Data Agent\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", "\n", "data_agent_response = requests.delete(data_agent_url, headers=headers)\n", "\n", "if data_agent_response.status_code == 200:\n", " print(\"Data Agent deleted successfully!\")\n", " print(json.dumps(data_agent_response.json(), indent=2))\n", "else:\n", " print(f\"Error Deleting Data Agent: {data_agent_response.status_code}\")\n", " print(data_agent_response.text)" ], "id": "y-wlEDnPrOvO" }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Data Chat Service" ], "id": "KHQUswQMHkcU" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Create Conversation\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations\"\n", "\n", "conversation_payload = {\n", " \"agents\": [\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " ],\n", " \"name\": (\n", " f\"projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n", " ),\n", "}\n", "params = {\"conversation_id\": conversation_id}\n", "\n", "conversation_response = requests.post(\n", " conversation_url, headers=headers, params=params, json=conversation_payload\n", ")\n", "\n", "if conversation_response.status_code == 200:\n", " print(\"Conversation created successfully!\")\n", " print(json.dumps(conversation_response.json(), indent=2))\n", "else:\n", " print(f\"Error creating Conversation: {conversation_response.status_code}\")\n", " print(conversation_response.text)" ], "id": "c9aCGQAYCMx1" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Get Conversation\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n", "\n", "conversation_response = requests.get(conversation_url, headers=headers)\n", "\n", "# Handle the response\n", "if conversation_response.status_code == 200:\n", " print(\"Conversation fetched successfully!\")\n", " print(json.dumps(conversation_response.json(), indent=2))\n", "else:\n", " print(f\"Error while fetching conversation: {conversation_response.status_code}\")\n", " print(conversation_response.text)" ], "id": "RxxC3Er3HpIR" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title List Conversation\n", "\n", "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations\"\n", "\n", "\n", "conversation_response = requests.get(conversation_url, headers=headers)\n", "\n", "# Handle the response\n", "if conversation_response.status_code == 200:\n", " print(\"Conversation fetched successfully!\")\n", " conversations = conversation_response.json().get(\"conversations\", [])\n", " display_resource_list(conversations)\n", "else:\n", " print(f\"Error while fetching conversations: {conversation_response.status_code}\")\n", " print(conversation_response.text)" ], "id": "3-3otk_aL6Wi" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title [Stateful] Chat using Conversation\n", "\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "\n", "# fmt: off\n", "question = \"Make a bar graph for the top 5 states by the total number of airports\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "chat_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat\"\n", "\n", "# Construct the payload\n", "chat_payload = {\n", " \"parent\": f\"projects/{billing_project}/locations/global\",\n", " \"messages\": [{\"userMessage\": {\"text\": question}}],\n", " \"conversation_reference\": {\n", " \"conversation\": (\n", " f\"projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n", " ),\n", " \"data_agent_context\": {\n", " \"data_agent\": (\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " ),\n", " },\n", " },\n", "}\n", "\n", "if looker_selected:\n", " chat_payload[\"credentials\"] = looker_credentials\n", "\n", "# Call get_stream method to stream the response\n", "get_stream(chat_url, chat_payload)" ], "id": "mQ1G8_UnKQjT" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title [Stateful] Chat using Data Agent\n", "\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: off\n", "question = \"Make a bar graph for the top 5 states by the total number of airports\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "chat_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}:chat\"\n", "\n", "# Construct the payload\n", "chat_payload = {\n", " \"parent\": f\"projects/{billing_project}/locations/global\",\n", " \"messages\": [{\"userMessage\": {\"text\": question}}],\n", " \"data_agent_context\": {\n", " \"data_agent\": (\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " ),\n", " },\n", "}\n", "\n", "if looker_selected:\n", " chat_payload[\"credentials\"] = looker_credentials\n", "\n", "# Call get_stream method to stream the response\n", "get_stream(chat_url, chat_payload)" ], "id": "p0vGFMmzI8BH" }, { "metadata": { "cellView": "form" }, "cell_type": "code", "source": [ "# @title [Stateless] Chat using Inline Context\n", "\n", "chat_url = (\n", " f\"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat\"\n", ")\n", "\n", "# fmt: off\n", "question = \"How many airports are in CA with elevation greater than 1000?\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "# Construct the payload\n", "chat_payload = {\n", " \"parent\": f\"projects/{billing_project}/locations/global\",\n", " \"messages\": [{\"userMessage\": {\"text\": question}}],\n", " \"inline_context\": {\n", " \"datasource_references\": datasource_references,\n", " \"options\": {\n", " \"analysis\": {\n", " # Optional - if wanting to use advanced analysis with python.\n", " # Default is False.\n", " \"python\": {\"enabled\": False}\n", " }\n", " },\n", " },\n", "}\n", "\n", "if bq_selected:\n", " if use_example_queries:\n", " chat_payload[\"inline_context\"][\"example_queries\"] = example_queries\n", " if use_glossary_terms:\n", " chat_payload[\"inline_context\"][\"glossary_terms\"] = glossary_terms\n", "elif looker_selected:\n", " chat_payload[\"credentials\"] = looker_credentials\n", " if use_looker_golden_queries:\n", " chat_payload[\"inline_context\"][\"looker_golden_queries\"] = looker_golden_queries\n", "\n", "# Call get_stream method to stream the response\n", "get_stream(chat_url, chat_payload)" ], "outputs": [], "execution_count": null, "id": "a8ef7ded" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title List Messages\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}/messages\"\n", "\n", "conversation_response = requests.get(conversation_url, headers=headers)\n", "\n", "# Handle the response\n", "if conversation_response.status_code == 200:\n", " print(\"Conversation fetched successfully!\")\n", " print(json.dumps(conversation_response.json(), indent=2))\n", "else:\n", " print(f\"Error while fetching conversation: {conversation_response.status_code}\")\n", " print(conversation_response.text)" ], "id": "Ioj6gG8wMedz" }, { "cell_type": "code", "source": [ "# @title Delete Conversation\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "conversation_url = f\"{base_url}/{api_version}/projects/{billing_project}/locations/{location}/conversations/{conversation_id}\"\n", "\n", "conversation_response = requests.delete(conversation_url, headers=headers)\n", "\n", "# Handle the response\n", "if conversation_response.status_code == 200:\n", " print(\"Conversation deleted successfully!\")\n", " print(json.dumps(conversation_response.json(), indent=2))\n", "else:\n", " print(f\"Error while deleting conversation: {conversation_response.status_code}\")\n", " print(conversation_response.text)" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [], "id": "rnIQ5blAUIuV" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Multi-turn Conversation with Data Agent or Inline Context" ], "id": "wdhjfwM3tF4G" }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form" }, "outputs": [], "source": [ "# @title Multi-turn Conversation helper function\n", "chat_url = (\n", " f\"{base_url}/{api_version}/projects/{billing_project}/locations/global:chat\"\n", ")\n", "\n", "# Re-used across requests to track previous turns\n", "conversation_messages = []\n", "\n", "# fmt: off\n", "context = \"data_agent_context\" # @param [\"data_agent_context\", \"inline_context\"]\n", "# Only for data_agent_context\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "\n", "# Helper function for calling the API\n", "def multi_turn_conversation(msg):\n", " userMessage = {\"userMessage\": {\"text\": msg}}\n", "\n", " # Send multi-turn request by including previous turns, plus new message\n", " conversation_messages.append(userMessage)\n", "\n", " # Construct the payload\n", " chat_payload = (\n", " {\n", " \"parent\": f\"projects/{billing_project}/locations/global\",\n", " \"messages\": conversation_messages,\n", " \"data_agent_context\": {\n", " \"data_agent\": (\n", " f\"projects/{billing_project}/locations/{location}/dataAgents/{data_agent_id}\"\n", " ),\n", " },\n", " }\n", " if context == \"data_agent_context\"\n", " else {\n", " \"parent\": f\"projects/{billing_project}/locations/global\",\n", " \"messages\": conversation_messages,\n", " \"inline_context\": {\n", " \"datasource_references\": datasource_references,\n", " },\n", " }\n", " )\n", "\n", " if looker_selected:\n", " chat_payload[\"credentials\"] = looker_credentials\n", " if context == \"inline_context\" and use_looker_golden_queries:\n", " chat_payload[\"inline_context\"][\n", " \"looker_golden_queries\"\n", " ] = looker_golden_queries\n", " elif bq_selected and context == \"inline_context\":\n", " if use_example_queries:\n", " chat_payload[\"inline_context\"][\"example_queries\"] = example_queries\n", " if use_glossary_terms:\n", " chat_payload[\"inline_context\"][\"glossary_terms\"] = glossary_terms\n", "\n", " # Call get_stream_multi_turn method to stream the response\n", " get_stream_multi_turn(chat_url, chat_payload, conversation_messages)" ], "id": "In8Gtf_Yte_i" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Send first-turn request\n", "multi_turn_conversation(\n", " \"List of the top 5 states by the total number of airports\"\n", ")" ], "id": "M_8Lvo68thp4" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Send follow-up-turn request\n", "multi_turn_conversation(\"Can you please make it a pie chart?\")" ], "id": "RFSy-DNTtjf7" } ], "metadata": { "colab": { "name": "intro_gemini_data_analytics_http.ipynb", "toc_visible": true, "provenance": [], "last_runtime": { "build_target": "//cloud/data_analytics/anarres/colab/cortado:base_colab", "kind": "private" } }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 5 }