{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github" }, "source": [ "\"Open" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "LUi4gYAeKS9W" }, "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." ] }, { "cell_type": "markdown", "metadata": { "id": "3decEzdcBiOl" }, "source": [ "# Intro to Gemini Data Analytics" ] }, { "cell_type": "markdown", "metadata": { "id": "idN2YJQiFa5P" }, "source": [ "| Author |\n", "| --- |\n", "| [Aditya Verma](https://github.com/vermaAstra) |" ] }, { "cell_type": "markdown", "metadata": { "id": "rCxXUPdGTOSC" }, "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 for HTTP here](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/agents/gemini_data_analytics/intro_gemini_data_analytics_http.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. Setting up the SDK\n", "2. Authenticate to Google Cloud\n", "3. Add data\n", "4. Perform agent operations (create, list, get, delete)\n", "5. Manage conversations (create, list, get, delete)\n", "6. Ask questions with your agent" ] }, { "cell_type": "markdown", "metadata": { "id": "PmGa6wSWBnvp" }, "source": [ "# Get Started" ] }, { "cell_type": "markdown", "metadata": { "id": "FZyNm_3Cq5hA" }, "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" ] }, { "cell_type": "markdown", "metadata": { "id": "ecf71a32661a" }, "source": [ "## Install the client library" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "4cda8ebb5891" }, "outputs": [], "source": [ "%pip install google-cloud-geminidataanalytics" ] }, { "cell_type": "markdown", "metadata": { "id": "e0c53abfd17c" }, "source": [ "# Setup" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "3SIB1yG9HvOq" }, "outputs": [], "source": [ "# @title Authenticate\n", "from google.colab import auth\n", "\n", "auth.authenticate_user()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "62cHAeMKPeh1" }, "outputs": [], "source": [ "# @title Imports\n", "import json\n", "import textwrap\n", "import time\n", "\n", "import altair as alt\n", "from google.iam.v1 import iam_policy_pb2, policy_pb2\n", "from google.protobuf import field_mask_pb2\n", "from google.protobuf.json_format import MessageToDict\n", "from IPython.display import HTML, display\n", "import pandas as pd\n", "import proto" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6rrLdjvm7QrX" }, "outputs": [], "source": [ "# @title Import the Gemini Data Analytics client library\n", "\n", "\n", "from google.cloud import geminidataanalytics_v1 as geminidataanalytics\n", "\n", "data_agent_client = geminidataanalytics.DataAgentServiceClient()\n", "data_chat_client = geminidataanalytics.DataChatServiceClient()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "beec2cd80e0a" }, "outputs": [], "source": [ "# @title Billing Project and sample Prompt (System Instruction)\n", "\n", "# fmt: off\n", "billing_project = \"[your-project-id]\" # @param {type:\"string\"}\n", "\n", "location = \"global\"\n", "\n", "# provide critical context for your Conversational Analytics Agent here\n", "system_instruction = \"Think like an Analyst\" # @param {type:\"string\"}\n", "# fmt: on" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "c09d0b409c82", "cellView": "form" }, "outputs": [], "source": [ "# @title BigQuery Datasource\n", "\n", "# List datasource(s) you'd like to ask questions of.\n", "# This example includes multiple (two) sources.\n", "# Make sure you have permission to query any tables listed here.\n", "\n", "# fmt: off\n", "bq_project_id_1 = \"bigquery-public-data\" # @param {type:\"string\"}\n", "bq_dataset_id_1 = \"faa\" # @param {type:\"string\"}\n", "bq_table_id_1 = \"us_airports\" # @param {type:\"string\"}\n", "\n", "bq_project_id_2 = \"bigquery-public-data\" # @param {type:\"string\"}\n", "bq_dataset_id_2 = \"san_francisco\" # @param {type:\"string\"}\n", "bq_table_id_2 = \"street_trees\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "\n", "# Setup\n", "bigquery_table_reference_1 = geminidataanalytics.BigQueryTableReference(\n", " project_id=bq_project_id_1,\n", " dataset_id=bq_dataset_id_1,\n", " table_id=bq_table_id_1,\n", ")\n", "bigquery_table_reference_2 = geminidataanalytics.BigQueryTableReference(\n", " project_id=bq_project_id_2,\n", " dataset_id=bq_dataset_id_2,\n", " table_id=bq_table_id_2,\n", ")\n", "\n", "# optional example queries (only leveraged for BigQuery datasources currently)\n", "example_queries = [\n", " geminidataanalytics.ExampleQuery(\n", " natural_language_question=\"How many airports are there?\",\n", " sql_query=\"SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports`\",\n", " ),\n", " geminidataanalytics.ExampleQuery(\n", " natural_language_question=(\n", " \"How many airports are in {state} with elevation greater than\"\n", " \" {elevation}?\"\n", " ),\n", " sql_query=(\n", " \"SELECT COUNT(*) FROM `bigquery-public-data.faa.us_airports` WHERE\"\n", " \" LOWER(state_abbreviation) = @state AND elevation > @elevation\"\n", " ),\n", " parameters=[\n", " geminidataanalytics.QueryParameter(\n", " name=\"state\",\n", " description=\"The state abbreviation in lower case.\",\n", " data_type=\"STRING\",\n", " ),\n", " geminidataanalytics.QueryParameter(\n", " name=\"elevation\",\n", " description=\"The elevation as a float value.\",\n", " data_type=\"FLOAT64\",\n", " ),\n", " ],\n", " ),\n", "]\n", "\n", "use_example_queries = True # @param {type:\"boolean\"}\n", "\n", "glossary_terms = [\n", " geminidataanalytics.GlossaryTerm(\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", "\n", "bq = geminidataanalytics.BigQueryTableReferences(\n", " table_references=[\n", " bigquery_table_reference_1,\n", " bigquery_table_reference_2,\n", " ]\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "c8d22eb7bf9e", "cellView": "form" }, "outputs": [], "source": [ "# @title Looker Datasource\n", "\n", "# List up to 5 Looker datasources you'd like to ask questions of.\n", "# This example includes two datasources.\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", "# User Credentials (apply to all Looker datasources)\n", "looker_client_id = \"\" # @param {type:\"string\"}\n", "looker_client_secret = \"\" # @param {type:\"string\"}\n", "# OR\n", "# looker_access_token = \"\" # @param {type:\"string\"}\n", "\n", "\n", "\"\"\" When connecting to the Looker datasource, you can authenticate using either:\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", "credentials = geminidataanalytics.Credentials(\n", " oauth=geminidataanalytics.OAuthCredentials(\n", " secret=geminidataanalytics.OAuthCredentials.SecretBased(\n", " client_id=looker_client_id, client_secret=looker_client_secret\n", " )\n", " )\n", ")\n", "# Uncomment this and comment out the above one if you are using access_token for authentication\n", "# credentials = geminidataanalytics.Credentials(\n", "# oauth=geminidataanalytics.OAuthCredentials(\n", "# token=geminidataanalytics.OAuthCredentials.TokenBased(\n", "# access_token = looker_access_token\n", "# )\n", "# )\n", "# )\n", "\n", "\n", "\"\"\" Looker datasource setup 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_explore_reference_1 = geminidataanalytics.LookerExploreReference(\n", " looker_instance_uri=looker_instance_uri_1,\n", " lookml_model=lookml_model_1,\n", " explore=explore_1,\n", ")\n", "# looker_explore_reference_2 = geminidataanalytics.LookerExploreReference(\n", "# looker_instance_uri=looker_instance_uri_2, lookml_model=lookml_model_2, explore=explore_2\n", "# )\n", "\n", "\n", "\"\"\" Looker datasource setup 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", "# (Optional)\n", "# looker_instance_id_2 = \"\" # @param {type:\"string\"}\n", "# looker_service_directory_name_2 = \"\" # @param {type:\"string\"}\n", "\n", "# looker_explore_reference_1 = geminidataanalytics.LookerExploreReference(\n", "# private_looker_instance_info=geminidataanalytics.PrivateLookerInstanceInfo(\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", "# (Optional)\n", "# looker_explore_reference_2 = geminidataanalytics.LookerExploreReference(\n", "# private_looker_instance_info=geminidataanalytics.PrivateLookerInstanceInfo(\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", "\n", "\n", "looker = geminidataanalytics.LookerExploreReferences(\n", " explore_references=[\n", " looker_explore_reference_1,\n", " # looker_explore_reference_2\n", " ]\n", ")\n", "\n", "looker_golden_queries = [\n", " geminidataanalytics.LookerGoldenQuery(\n", " natural_language_questions=[\n", " \"What is the highest observed positive longitude?\"\n", " ],\n", " looker_query=geminidataanalytics.LookerQuery(\n", " model=\"airports\",\n", " explore=\"airports\",\n", " fields=[\"airports.longitude\"],\n", " filters=[\n", " geminidataanalytics.LookerQuery.Filter(\n", " field=\"airports.longitude\", value=\">0\"\n", " )\n", " ],\n", " sorts=[\"airports.longitude desc\"],\n", " limit=\"1\",\n", " ),\n", " ),\n", " geminidataanalytics.LookerGoldenQuery(\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=geminidataanalytics.LookerQuery(\n", " model=\"airports\",\n", " explore=\"airports\",\n", " fields=[\"airports.city\", \"airports.code\"],\n", " filters=[\n", " geminidataanalytics.LookerQuery.Filter(\n", " field=\"airports.major\", value=\"Y\"\n", " ),\n", " geminidataanalytics.LookerQuery.Filter(\n", " field=\"airports.state\", value=\"CA\"\n", " ),\n", " ],\n", " ),\n", " ),\n", "]\n", "\n", "use_looker_golden_queries = False # @param {type:\"boolean\"}" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "e83dc687b9e0", "cellView": "form" }, "outputs": [], "source": [ "# @title Looker Studio Datasource\n", "\n", "# fmt: off\n", "studio_datasource_id = \"\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "# Setup\n", "studio_references = geminidataanalytics.StudioDatasourceReference(\n", " datasource_id=studio_datasource_id\n", ")\n", "\n", "studio = geminidataanalytics.StudioDatasourceReferences(\n", " studio_references=[studio_references]\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "d2de560116e2", "cellView": "form" }, "outputs": [], "source": [ "# @title Select Data Source\n", "\n", "selected_datasource = \"bigquery_data_sources\" # @param [\"bigquery_data_sources\", \"looker_data_sources\", \"looker_studio_data_sources\"]\n", "\n", "bq_selected = False\n", "looker_selected = False\n", "looker_studio_selected = False\n", "datasource_references = geminidataanalytics.DatasourceReferences()\n", "\n", "if selected_datasource == \"looker_data_sources\":\n", " looker_selected = True\n", " datasource_references.looker = looker\n", "elif selected_datasource == \"bigquery_data_sources\":\n", " bq_selected = True\n", " datasource_references.bq = bq\n", "elif selected_datasource == \"looker_studio_data_sources\":\n", " looker_studio_selected = True\n", " datasource_references.studio = studio\n", "\n", "\n", "# Context set-up\n", "published_context = geminidataanalytics.Context(\n", " system_instruction=system_instruction,\n", " datasource_references=datasource_references,\n", " options=geminidataanalytics.ConversationOptions(\n", " analysis=geminidataanalytics.AnalysisOptions(\n", " python=geminidataanalytics.AnalysisOptions.Python(\n", " enabled=False\n", " ) # if wanting to use advanced analysis with python. Default is False.\n", " )\n", " ),\n", ")\n", "\n", "# Context set-up for chat using inline context\n", "inline_context = geminidataanalytics.Context(\n", " system_instruction=system_instruction,\n", " datasource_references=datasource_references,\n", " options=geminidataanalytics.ConversationOptions(\n", " analysis=geminidataanalytics.AnalysisOptions(\n", " python=geminidataanalytics.AnalysisOptions.Python(\n", " enabled=False\n", " ) # if wanting to use advanced analysis with python. Default is False.\n", " )\n", " ),\n", ")" ] }, { "cell_type": "markdown", "source": [ "# Helper Functions" ], "metadata": { "id": "YsOsYWaPWj7I" } }, { "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", " # Convert the proto message to a dictionary for field access and JSON serialization\n", " # 'resource._pb' accesses the underlying protobuf message required by MessageToDict\n", " resource_dict = MessageToDict(resource._pb)\n", "\n", " resource_name = resource_dict.get(\"name\").split(\"/\")[-1]\n", " resource_display_name = resource_dict.get(\"displayName\")\n", " formatted_json = json.dumps(resource_dict, 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", "id": "W5nyWlf7WyFF" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "3sPlf0l09t2y", "cellView": "form" }, "outputs": [], "source": [ "# @title Define utilities for viewing responses\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 display_schema(data):\n", " fields = data.fields\n", " df = pd.DataFrame({\n", " \"Column\": map(lambda field: field.name, fields),\n", " \"Type\": map(lambda field: field.type, fields),\n", " \"Description\": map(\n", " lambda field: getattr(field, \"description\", \"-\"), fields\n", " ),\n", " \"Mode\": map(lambda field: 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_looker_table_ref(table_ref):\n", " return f\"lookmlModel: {table_ref.lookml_model}, explore: {table_ref.explore}\"\n", "\n", "\n", "def format_bq_table_ref(table_ref):\n", " return f\"{table_ref.project_id}.{table_ref.dataset_id}.{table_ref.table_id}\"\n", "\n", "\n", "def display_datasource(datasource):\n", " source_name = \"\"\n", " if \"studio_datasource_id\" in datasource:\n", " source_name = datasource.studio_datasource_id\n", " elif \"looker_explore_reference\" in datasource:\n", " source_name = format_looker_table_ref(datasource.looker_explore_reference)\n", " else:\n", " source_name = format_bq_table_ref(datasource.bigquery_table_reference)\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(f\"Query name: {query.name}\")\n", " if \"question\" in query:\n", " print(f\"Question: {query.question}\")\n", " if \"datasources\" in query:\n", " print(\"Data sources:\")\n", " for datasource in query.datasources:\n", " display_datasource(datasource)\n", " elif \"generated_sql\" in resp:\n", " display_section_title(\"SQL generated\")\n", " print(resp.generated_sql)\n", " elif \"result\" in resp:\n", " display_section_title(\"Data retrieved\")\n", "\n", " if \"schema\" in resp.result:\n", " fields = [field.name for field in resp.result.schema.fields]\n", " d = {}\n", " for el in resp.result.data:\n", " for field in fields:\n", " if field in d:\n", " d[field].append(el[field])\n", " else:\n", " d[field] = [el[field]]\n", "\n", " display(pd.DataFrame(d))\n", "\n", "\n", "def handle_chart_response(resp):\n", " def _value_to_dict(v):\n", " if isinstance(v, proto.marshal.collections.maps.MapComposite):\n", " return _map_to_dict(v)\n", " if isinstance(v, proto.marshal.collections.RepeatedComposite):\n", " return [_value_to_dict(el) for el in v]\n", " if isinstance(v, (int, float, str, bool)):\n", " return v\n", " return MessageToDict(v)\n", "\n", " def _map_to_dict(d):\n", " out = {}\n", " for k in d:\n", " if isinstance(d[k], proto.marshal.collections.maps.MapComposite):\n", " out[k] = _map_to_dict(d[k])\n", " else:\n", " out[k] = _value_to_dict(d[k])\n", " return out\n", "\n", " if \"query\" in resp and \"instructions\" in resp.query:\n", " print(resp.query.instructions)\n", " elif \"result\" in resp:\n", " vegaConfig = resp.result.vega_config\n", " vegaConfig_dict = _map_to_dict(vegaConfig)\n", " alt.Chart.from_json(json.dumps(vegaConfig_dict)).display()\n", "\n", "\n", "def show_message(msg):\n", " m = msg.system_message\n", " if \"text\" in m:\n", " handle_text_response(m.text)\n", " elif \"schema\" in m:\n", " handle_schema_response(m.schema)\n", " elif \"data\" in m:\n", " handle_data_response(m.data)\n", " elif \"chart\" in m:\n", " handle_chart_response(m.chart)\n", " print(\"\\n\")" ] }, { "cell_type": "markdown", "metadata": { "id": "xMjbfg7YUnTc" }, "source": [ "### Write an effective prompt (system_instruction)\n", "- Providing more context around your business and your data improves the quality of answers. We recommend (though do not require) you follow our example `system_instruction` below and fill out the stringified YAML **template** provided with field instructions, validated (\"golden\") queries, relationships, business terms, etc.\n", "- We are working to improve how our product ingests and uses this context to improve question-answering accuracy" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Bla2oAfRHcOx" }, "outputs": [], "source": [ "###############################################################################\n", "####### Define prompt / system instruction for the question. #######\n", "####### It's optional and can guide the agent's response strategy. #######\n", "###############################################################################\n", "\n", "###############################################################################\n", "##### Example template for system instruction for a sales analyst agent: #####\n", "###############################################################################\n", "\n", "# \"\"\"\n", "# - system_instruction: >-\n", "# You are an expert sales analyst and understand how to answer questions about\n", "# the sales data for a fictitious e-commerce store\n", "# - tables:\n", "# - table:\n", "# - name: bigquery-public-data.thelook_ecommerce.orders\n", "# - description: orders for The Look fictitious e-commerce store.\n", "# - synonyms: sales\n", "# - tags: 'sale, order, sales_order'\n", "# - fields:\n", "# - field:\n", "# - name: order_id\n", "# - description: unique identifier for each order\n", "# - field:\n", "# - name: user_id\n", "# - description: unique identifier for each user\n", "# - field:\n", "# - name: status\n", "# - description: status of the order\n", "# - sample_values:\n", "# - complete\n", "# - shipped\n", "# - returned\n", "# - field:\n", "# - name: gender\n", "# - description: gender of the user\n", "# - sample_values:\n", "# - male\n", "# - female\n", "# - field:\n", "# - name: created_at\n", "# - description: date and time when the order was created in timestamp format\n", "# - field:\n", "# - name: returned_at\n", "# - description: >-\n", "# date and time when the order was returned in timestamp\n", "# format\n", "# - field:\n", "# - name: num_of_item\n", "# - description: number of items in the order\n", "# - aggregations: 'sum, avg'\n", "# - field:\n", "# - name: earnings\n", "# - description: total sales from the order\n", "# - aggregations: 'sum, avg'\n", "# - field:\n", "# - name: cost\n", "# - description: total cost to get the items for the order\n", "# - aggregations: 'sum, avg'\n", "# - measures:\n", "# - measure:\n", "# - name: profit\n", "# - description: raw profit\n", "# - exp: cost - earnings\n", "# - synonyms: gains\n", "# - golden_queries:\n", "# - golden_query:\n", "# - natural_language_query: How many orders are there?\n", "# - sql_query: SELECT COUNT(*) FROM sqlgen-testing.thelook_ecommerce.orders\n", "# - golden_query:\n", "# - natural_language_query: How many orders were shipped?\n", "# - sql_query: >-\n", "# SELECT COUNT(*) FROM sqlgen-testing.thelook_ecommerce.orders\n", "# WHERE status = 'shipped'\n", "# - golden_action_plans:\n", "# - golden_action_plan:\n", "# - natural_language_query: Show me the number of orders broken down by gender.\n", "# - action_plan:\n", "# - step: >-\n", "# Run a SQL query on the table\n", "# sqlgen-testing.thelook_ecommerce.orders to get a\n", "# breakdown of order count by gender.\n", "# - step: >-\n", "# Create a vertical bar plot using the retrieved data,\n", "# with one bar per gender.\n", "# - table:\n", "# - name: sqlgen-testing.thelook_ecommerce.users\n", "# - description: user of The Look fictitious e-commerce store.\n", "# - synonyms: customers\n", "# - tags: 'user, customer, buyer'\n", "# - fields:\n", "# - field:\n", "# - name: id\n", "# - description: unique identifier for each user\n", "# - field:\n", "# - name: first_name\n", "# - description: first name of the user\n", "# - tag: person\n", "# - sample_values: 'graham, adam, brian'\n", "# - field:\n", "# - name: last_name\n", "# - description: first name of the user\n", "# - tag: person\n", "# - sample_values: 'warmer, stilles, smith'\n", "# - field:\n", "# - name: gender\n", "# - description: gender of the user\n", "# - sample_values:\n", "# - male\n", "# - female\n", "# - field:\n", "# - name: email\n", "# - description: email of the user\n", "# - tag: contact\n", "# - sample_values: 'brian@gmail.com, graham@gmail.com'\n", "# - golden_queries:\n", "# - golden_query:\n", "# - natural_language_query: How many unique customers are there?\n", "# - sql_query: >-\n", "# SELECT COUNT(DISTINCT id) FROM\n", "# sqlgen-testing.thelook_ecommerce.users\n", "# - golden_query:\n", "# - natural_language_query: How many female users have gmail email id?\n", "# - sql_query: >-\n", "# SELECT COUNT(DISTINCT id) FROM\n", "# sqlgen-testing.thelook_ecommerce.users WHERE users.gender =\n", "# 'F' AND users.email LIKE '%@gmail.com';\n", "# - relationships:\n", "# - relationship:\n", "# - name: earnings_to_user\n", "# - description: >-\n", "# Sales table is related to the users table and can be joined for\n", "# aggregated view.\n", "# - relationship_type: many-to-one\n", "# - join_type: left\n", "# - left_table: str sqlgen-testing.thelook_ecommerce.orders\n", "# - right_table: sqlgen-testing.thelook_ecommerce.users\n", "# - relationship_columns: '// Join columns - left_column:''user_id'' - right_column:''id'''\n", "# - glossaries:\n", "# - glossary:\n", "# - term: male\n", "# - description: male gender\n", "# - glossary:\n", "# - term: female\n", "# - description: female gender\n", "# - glossary:\n", "# - term: complete\n", "# - description: complete status\n", "# - synonyms: 'finish, done, fulfilled'\n", "# - glossary:\n", "# - term: shipped\n", "# - description: shipped status\n", "# - glossary:\n", "# - term: returned\n", "# - description: returned status\n", "# - glossary:\n", "# - term: OMPF\n", "# - description: Order Management and Product Fulfillment\n", "# - additional_descriptions:\n", "# - text: All the sales data is for Looker organization.\n", "# - text: 'Orders can be of three categories, food , clothes, electronics.'\n", "# \"\"\"\n", "\n", "###############################################################################\n", "###### Use the schema below to provide your own system instruction. ######\n", "###############################################################################\n", "\n", "# \"\"\"\n", "# - system_instruction: str # Expected behavior of the agent. Eg. You are a sales analyst.\n", "# - tables:\n", "# - table: # Table relevant for this agent.\n", "# - name: str # Name of the table.\n", "# - description: str # Description of the table.\n", "# - synonyms: list[str] # List of synonyms used to refer to the table.\n", "# - tags: list[str] # List of tags associated with the table.\n", "# - fields: # Fields in the table.\n", "# - field:\n", "# - name: str # Name of the column.\n", "# - description: str # Description of the column.\n", "# - synonyms: list[str] # List of synonyms used to refer to the column.\n", "# - tags: list[str] # List of tags associated with the column.\n", "# - sample_values: list[str] # List of sample values in the column.\n", "# - aggregations: list[str] # Any commonly used or default aggregations associated with the column.\n", "# - measures: # Measures for the table.\n", "# - measure:\n", "# - name: str # Name of the measure.\n", "# - description: str # Description of the measure.\n", "# - exp: str # Expression to construct the measure.\n", "# - synonyms: list[str] # List of synonyms used to refer to the measure.\n", "# - golden_queries: # Golden or popular queries for the table.\n", "# - golden_query:\n", "# - natural_language_query: str # Natural language query.\n", "# - sql_query: str # SQL query.\n", "# - golden_action_plans: # Golden action plans as the suggested steps to take (in order) to answer the query.\n", "# - golden_action_plan:\n", "# - natural_language_query: str # Natural language query.\n", "# - action_plan:\n", "# - step: str # Step to take.\n", "# - relationships: # Join relationships between tables.\n", "# - relationship:\n", "# - name: str # Name of the relationship.\n", "# - description: str # Description of the relationship.\n", "# - relationship_type: str # Relationship type: one-to-one, one-to-many, many-to-one, many-to-many.\n", "# - join_type: str # Join type: inner, outer, left, right, full\n", "# - left_table: str # Left table name.\n", "# - right_table: str # Right table name.\n", "# - relationship_columns: # Join columns.\n", "# - left_column: str # Join column from left table.\n", "# - right_column: str # Join column from right table.\n", "# - glossaries: # Business glossary, jargon, etc.\n", "# - glossary:\n", "# - term: str # Name of the term. Term can be a word, phrase, abbreviation, etc.\n", "# - description: str # Description or definition of the term.\n", "# - synonyms: list[str] # List of synonyms for the term.\n", "# - additional_descriptions:\n", "# - text: str # Any additional description that was not covered above.\n", "# \"\"\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "S6-3l1ggHK2_" }, "outputs": [], "source": [ "##################### OPTION 1 ###############################\n", "####### #######\n", "####### Fill the above YAML template for #######\n", "####### system_instruction as shown in example above. #######\n", "####### #######\n", "##############################################################\n", "\n", "system_instruction = \"\"\"\"\"\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "9Ac6KtdMHhea" }, "outputs": [], "source": [ "##################### OPTION 2 ###############################\n", "####### #######\n", "####### Leave system_instruction empty #######\n", "####### or provide description in form #######\n", "####### of long string. #######\n", "####### #######\n", "##############################################################\n", "\n", "# Define context for question.\n", "# system_instruction is optional, can be used to steer the agent's answering\n", "# strategy.\n", "# fmt: off\n", "system_instruction = \"Think like an Analyst\" # @param {type:\"string\"}\n", "# fmt: on" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ezPjVAXRHxFc" }, "outputs": [], "source": [ "# @title Helper functions for processing the system_instruction yaml string\n", "\n", "from collections import Counter\n", "import re\n", "\n", "import yaml\n", "from yaml.scanner import ScannerError\n", "\n", "\n", "def cleanup_path(input_string: str) -> str:\n", " \"\"\"Removes bracketed numbers and preceding dots from an input string.\n", "\n", " Args:\n", " input_string: The input string.\n", "\n", " Returns:\n", " The string with bracketed numbers and preceding dots removed.\n", " \"\"\"\n", " cleaned_string = re.sub(r\"\\[\\d+]\", \"\", input_string)\n", " if cleaned_string.startswith(\".\"):\n", " cleaned_string = cleaned_string[1:]\n", " return cleaned_string\n", "\n", "\n", "def is_valid_yaml(yaml_string: str) -> bool:\n", " try:\n", " yaml.safe_load(yaml_string)\n", " return True\n", " except ScannerError:\n", " return False\n", " except Exception as e:\n", " print(f\"An unexpected error occurred: {e}\")\n", " return False\n", "\n", "\n", "def print_counter_desc(data: list[str]) -> None:\n", " # Count the occurrences of each item in the list.\n", " item_counts = Counter(data)\n", "\n", " # Sort the items in descending order based on their counts.\n", " sorted_items = sorted(\n", " item_counts.items(), key=lambda item: item[1], reverse=True\n", " )\n", "\n", " # Print the items and their counts.\n", " for item, count in sorted_items:\n", " print(f\"{item}: {count}\")\n", "\n", "\n", "def extract_insights_from_system_instruction(yaml_string: str) -> None:\n", " \"\"\"Extracts insights into used fields from a YAML string.\n", "\n", " Args:\n", " yaml_string: The YAML string to analyze.\n", " \"\"\"\n", " if not is_valid_yaml(yaml_string):\n", " print(\"Received Invalid or not a YAML format. Cannot extract insights.\")\n", " return\n", "\n", " try:\n", " data = yaml.safe_load(yaml_string)\n", " used_fields = []\n", "\n", " def traverse(node, path=\"\"):\n", " if isinstance(node, dict):\n", " for key, value in node.items():\n", " new_path = f\"{path}.{key}\" if path else key\n", " traverse(value, new_path)\n", " elif isinstance(node, list):\n", " for i, item in enumerate(node):\n", " new_path = f\"{path}[{i}]\"\n", " traverse(item, new_path)\n", " else:\n", " used_fields.append(cleanup_path(path))\n", "\n", " traverse(data)\n", " print(\"Insights:\")\n", " print(\"Number of unique fields used: \", len(set(used_fields)))\n", " print(\"Used fields in YAML:\")\n", " print()\n", " print_counter_desc(used_fields)\n", "\n", " except ScannerError as e:\n", " print(f\"Invalid YAML: {e}\")\n", " return\n", " except Exception as e:\n", " print(f\"An unexpected error occurred: {e}\")\n", " return\n", "\n", "\n", "# Printing insights assuming 'system_instruction' is defined above.\n", "extract_insights_from_system_instruction(system_instruction)" ] }, { "cell_type": "markdown", "metadata": { "id": "TGAOb15dlRKj" }, "source": [ "# Example API call\n", "**All API calls** take `billing_project` (a Google Cloud billing project), a `question` (first in a list of Messages), and `prompt` / `system_instruction` (optional)\n", "\n", "\n", "**BigQuery calls** create a BigQueryTableReference via `project_id`, `dataset_id`, and `table_id`. Note that multiple table references can be passed in a list.\n", "\n", "\n", "**Looker calls** create a LookerExploreReference with `looker_instance_uri` for public instance or `private_looker_instance_info` for private instance, `lookml_model`, and `explore`. Auth credentials are created with `client_id` and `client_secret` or `access_token`.\n", "\n", "*Note: options for advanced system instructions are in the cells below*" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "qXplnBRAMOXN" }, "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 = geminidataanalytics.DataAgent(\n", " data_analytics_agent=geminidataanalytics.DataAnalyticsAgent(\n", " published_context=published_context\n", " ),\n", ")\n", "\n", "request = geminidataanalytics.CreateDataAgentRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " data_agent_id=data_agent_id, # Optional\n", " data_agent=data_agent,\n", ")\n", "\n", "if bq_selected:\n", " if use_example_queries:\n", " request.data_agent.data_analytics_agent.published_context.example_queries = (\n", " example_queries\n", " )\n", " if use_glossary_terms:\n", " request.data_agent.data_analytics_agent.published_context.glossary_terms = (\n", " glossary_terms\n", " )\n", "elif looker_selected and use_looker_golden_queries:\n", " request.data_agent.data_analytics_agent.published_context.looker_golden_queries = (\n", " looker_golden_queries\n", " )\n", "\n", "try:\n", " operation = data_agent_client.create_data_agent(request=request)\n", "\n", " while not operation.done():\n", " print(\"Still in progress...\")\n", " time.sleep(2)\n", "\n", " if operation.done():\n", " if operation.exception():\n", " print(\"Operation failed:\", operation.exception())\n", " else:\n", " print(\"Operation succeeded:\", operation.result())\n", "except Exception as e:\n", " print(f\"Error creating Data Agent: {e}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "27yX_ItjMqWS", "cellView": "form" }, "outputs": [], "source": [ "# @title Create Conversation\n", "\n", "# Initialize request argument(s)\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 = geminidataanalytics.Conversation(\n", " agents=[\n", " data_chat_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " )\n", " ],\n", ")\n", "\n", "request = geminidataanalytics.CreateConversationRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " conversation_id=conversation_id,\n", " conversation=conversation,\n", ")\n", "\n", "# Make the request\n", "response = data_chat_client.create_conversation(request=request)\n", "\n", "# Handle the response\n", "print(response)" ] }, { "cell_type": "markdown", "metadata": { "id": "t6453RzShYWB" }, "source": [ "# Chat request\n", "1. Chat request with Conversation reference\n", "2. Chat request Data Agent\n", "3. Chat request with Inline Context" ] }, { "cell_type": "markdown", "metadata": { "id": "Q2mALm-giwgj" }, "source": [ "## [Stateful] Chat request with conversation reference" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "vkk9yYAHWnKU" }, "outputs": [], "source": [ "# Create a request containing a single user message -- your question.\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\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", "# Create the user message\n", "messages = [\n", " geminidataanalytics.Message(\n", " user_message=geminidataanalytics.UserMessage(text=question)\n", " )\n", "]\n", "\n", "# Create a conversation_reference\n", "conversation_reference = geminidataanalytics.ConversationReference(\n", " conversation=data_chat_client.conversation_path(\n", " billing_project, location, conversation_id\n", " ),\n", " data_agent_context=geminidataanalytics.DataAgentContext(\n", " data_agent=data_chat_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " ),\n", " ),\n", ")\n", "\n", "# Form the request\n", "request = geminidataanalytics.ChatRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " messages=messages,\n", " conversation_reference=conversation_reference,\n", ")\n", "\n", "if looker_selected:\n", " request.credentials = credentials\n", "\n", "# Make the request\n", "stream = data_chat_client.chat(request=request)\n", "\n", "# Handle the response\n", "for response in stream:\n", " show_message(response)" ] }, { "cell_type": "markdown", "metadata": { "id": "LdN2o3qUi99z" }, "source": [ "## [Stateful] Chat request with data_agent" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "_t3_8o_qgxzp" }, "outputs": [], "source": [ "# Create a request containing a single user message -- your question.\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\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", "# Create the message\n", "messages = [\n", " geminidataanalytics.Message(\n", " user_message=geminidataanalytics.UserMessage(text=question)\n", " )\n", "]\n", "\n", "# Create the data agent context\n", "data_agent_context = geminidataanalytics.DataAgentContext(\n", " data_agent=data_chat_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " ),\n", ")\n", "\n", "# Form the request\n", "request = geminidataanalytics.ChatRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " messages=messages,\n", " data_agent_context=data_agent_context,\n", ")\n", "\n", "if looker_selected:\n", " request.credentials = credentials\n", "\n", "# Make the request\n", "stream = data_chat_client.chat(request=request)\n", "\n", "# Handle the response\n", "for response in stream:\n", " show_message(response)" ] }, { "cell_type": "markdown", "metadata": { "id": "NKmRVLyDlt0n" }, "source": [ "## [Stateless] Chat request with Inline Context" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "jhXiPBZY8QM7" }, "outputs": [], "source": [ "# Create a request containing a single user message -- your question.\n", "# fmt: off\n", "question = \"How many airports are in CA with elevation greater than 1000?\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "messages = [\n", " geminidataanalytics.Message(\n", " user_message=geminidataanalytics.UserMessage(text=question)\n", " )\n", "]\n", "\n", "request = geminidataanalytics.ChatRequest(\n", " inline_context=inline_context,\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " messages=messages,\n", ")\n", "if bq_selected:\n", " if use_example_queries:\n", " request.inline_context.example_queries = example_queries\n", " if use_glossary_terms:\n", " request.inline_context.glossary_terms = glossary_terms\n", "elif looker_selected:\n", " request.credentials = credentials\n", " if use_looker_golden_queries:\n", " request.inline_context.looker_golden_queries = looker_golden_queries\n", "\n", "# Make the request\n", "stream = data_chat_client.chat(request=request)\n", "\n", "# Handle the response\n", "for response in stream:\n", " show_message(response)" ] }, { "cell_type": "markdown", "metadata": { "id": "SAMel8bWaoVs" }, "source": [ "# Multi-turn Conversation with Data Agent or Inline Context" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "SyBA1KHtyOgD" }, "outputs": [], "source": [ "# @title Multi-turn Conversation\n", "\n", "# Re-used across requests to track previous turns\n", "conversation_messages = []\n", "\n", "# fmt: off\n", "context = \"inline_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", "data_agent_context = geminidataanalytics.DataAgentContext(\n", " data_agent=data_chat_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " ),\n", ")\n", "\n", "# Helper function for calling the API\n", "def multi_turn_conversation(msg):\n", " message = geminidataanalytics.Message(\n", " user_message=geminidataanalytics.UserMessage(text=msg)\n", " )\n", "\n", " # Send multi-turn request by including previous turns plus new message\n", " conversation_messages.append(message)\n", "\n", " request = geminidataanalytics.ChatRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " messages=conversation_messages,\n", " **(\n", " {\"data_agent_context\": data_agent_context}\n", " if context == \"data_agent_context\"\n", " else {\"inline_context\": inline_context}\n", " ),\n", " )\n", "\n", " if looker_selected:\n", " request.credentials = credentials\n", " if context == \"inline_context\" and use_looker_golden_queries:\n", " inline_context.looker_golden_queries = looker_golden_queries\n", " elif bq_selected and context == \"inline_context\":\n", " if use_example_queries:\n", " inline_context.example_queries = example_queries\n", " if use_glossary_terms:\n", " inline_context.glossary_terms = glossary_terms\n", "\n", " # Make the request\n", " stream = data_chat_client.chat(request=request)\n", "\n", " # Handle the response\n", " for response in stream:\n", " show_message(response)\n", " if response.system_message or response.user_message:\n", " conversation_messages.append(response)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "pmdkDdNFzZFg" }, "outputs": [], "source": [ "# Send first-turn request\n", "multi_turn_conversation(\n", " \"List of the top 5 states by the total number of airports\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "WnCikHbCOy3o" }, "outputs": [], "source": [ "# Send follow-up-turn request\n", "multi_turn_conversation(\"Can you please make it a pie chart?\")" ] }, { "cell_type": "markdown", "metadata": { "id": "mmz03WdUJ4Ya" }, "source": [ "# Other DataAgent Service Methods" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "Y1QOX0AFKGDh" }, "outputs": [], "source": [ "# @title Get Data Agent\n", "\n", "# Initialize request argument(s)\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "request = geminidataanalytics.GetDataAgentRequest(\n", " name=data_agent_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " )\n", ")\n", "\n", "# Make the request\n", "response = data_agent_client.get_data_agent(request=request)\n", "\n", "# Handle the response\n", "print(response)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "dBVG9CEPKwuF" }, "outputs": [], "source": [ "# @title List Data Agents\n", "\n", "# Initialize request argument(s)\n", "request = geminidataanalytics.ListDataAgentsRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " order_by=\"create_time desc\", # Optional\n", ")\n", "\n", "# Make the request\n", "data_agents = data_agent_client.list_data_agents(request=request)\n", "\n", "# Handle the response\n", "display_resource_list(data_agents)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "cc98454d36ea", "cellView": "form" }, "outputs": [], "source": [ "# @title List Accessible Data Agents\n", "\n", "# fmt: off\n", "creator_filter = \"NONE\" # @param [\"NONE\", \"CREATOR_ONLY\", \"NOT_CREATOR_ONLY\"]\n", "# fmt: on\n", "\n", "# Initialize request argument(s)\n", "request = geminidataanalytics.ListAccessibleDataAgentsRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", " creator_filter=creator_filter,\n", ")\n", "\n", "# Make the request\n", "data_agents = data_agent_client.list_accessible_data_agents(request=request)\n", "\n", "# Handle the response\n", "display_resource_list(data_agents)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "pSAaAcUdLIvO" }, "outputs": [], "source": [ "# @title Update Data Agent\n", "\n", "# Initialize request argument(s)\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "data_agent = geminidataanalytics.DataAgent(\n", " data_analytics_agent=geminidataanalytics.DataAnalyticsAgent(\n", " published_context=published_context\n", " ),\n", " name=data_agent_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " ),\n", " description=\"This is my new updated description.\",\n", ")\n", "\n", "\n", "update_mask = field_mask_pb2.FieldMask(\n", " paths=[\"description\", \"data_analytics_agent.published_context\"]\n", ")\n", "\n", "request = geminidataanalytics.UpdateDataAgentRequest(\n", " data_agent=data_agent,\n", " update_mask=update_mask,\n", ")\n", "\n", "try:\n", " operation = data_agent_client.update_data_agent(request=request)\n", "\n", " while not operation.done():\n", " print(\"Still in progress...\")\n", " time.sleep(2)\n", "\n", " if operation.done():\n", " if operation.exception():\n", " print(\"Operation failed:\", operation.exception())\n", " else:\n", " print(\"Operation succeeded:\", operation.result())\n", "except Exception as e:\n", " print(f\"Error updating Data Agent: {e}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "e6598c46fe18", "cellView": "form" }, "outputs": [], "source": [ "# @title [Agent Sharing] Set IAM Policy for Data Agent\n", "\n", "\"\"\"PLEASE NOTE THIS API CALLS OVERRIDE EXISTING PERMISSION FOR THE RESOURCE.\n", "\n", "For preserving existing policy in practice 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", "# Inputs\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "role = \"roles/geminidataanalytics.dataAgentEditor\" # @param {type:\"string\"}\n", "users = \"abc@google.com, wxyz@google.com\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "resource = (\n", " f\"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}\"\n", ")\n", "\n", "# Construct the IAM policy\n", "binding = policy_pb2.Binding(\n", " role=role, members=[f\"user:{i.strip()}\" for i in users.split(\",\")]\n", ")\n", "\n", "policy = policy_pb2.Policy(bindings=[binding])\n", "\n", "# Create the request\n", "request = iam_policy_pb2.SetIamPolicyRequest(resource=resource, policy=policy)\n", "\n", "# Send the request\n", "try:\n", " response = data_agent_client.set_iam_policy(request=request)\n", " print(\"IAM Policy set successfully!\")\n", " print(f\"Response: {response}\")\n", "except Exception as e:\n", " print(f\"Error setting IAM policy: {e}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "6ea4d2206c55", "cellView": "form" }, "outputs": [], "source": [ "# @title [Agent Sharing] GET IAM Policy for Data Agent\n", "\n", "# Initialize request argument(s)\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "resource = (\n", " f\"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}\"\n", ")\n", "request = iam_policy_pb2.GetIamPolicyRequest(\n", " resource=resource,\n", ")\n", "try:\n", " response = data_agent_client.get_iam_policy(request=request)\n", " print(\"IAM Policy fetched successfully!\")\n", " print(f\"Response: {response}\")\n", "except Exception as e:\n", " print(f\"Error setting IAM policy: {e}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "GDiXL18MNaAn" }, "outputs": [], "source": [ "# @title [Soft Delete] Delete Data Agent\n", "\n", "# Initialize request argument(s)\n", "\n", "# fmt: off\n", "data_agent_id = \"data_agent_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "request = geminidataanalytics.DeleteDataAgentRequest(\n", " name=data_agent_client.data_agent_path(\n", " billing_project, location, data_agent_id\n", " )\n", ")\n", "\n", "try:\n", " # Make the request\n", " data_agent_client.delete_data_agent(request=request)\n", " print(\"Data Agent Deleted\")\n", "except Exception as e:\n", " print(f\"Error deleting Data Agent: {e}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "t2gxdjLKCe9B" }, "source": [ "# Other DataChat Service Methods" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "clN0SioB6ejc", "cellView": "form" }, "outputs": [], "source": [ "# @title Get Conversation\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "request = geminidataanalytics.GetConversationRequest(\n", " name=data_chat_client.conversation_path(\n", " billing_project, location, conversation_id\n", " ),\n", ")\n", "\n", "# Make the request\n", "response = data_chat_client.get_conversation(request=request)\n", "\n", "# Handle the response\n", "print(response)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "ibniY5_vBDnd" }, "outputs": [], "source": [ "# @title List Conversation\n", "\n", "request = geminidataanalytics.ListConversationsRequest(\n", " parent=f\"projects/{billing_project}/locations/{location}\",\n", ")\n", "\n", "# Make the request\n", "conversations = data_chat_client.list_conversations(request=request)\n", "\n", "# Handle the response\n", "display_resource_list(conversations)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "xqeWlpaRB_66" }, "outputs": [], "source": [ "# @title List Messages\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "request = geminidataanalytics.ListMessagesRequest(\n", " parent=data_chat_client.conversation_path(\n", " billing_project, location, conversation_id\n", " ),\n", ")\n", "\n", "# Make the request\n", "response = data_chat_client.list_messages(request=request)\n", "\n", "# Handle the response\n", "print(response)" ] }, { "cell_type": "code", "source": [ "# @title Delete Conversation\n", "\n", "# fmt: off\n", "conversation_id = \"conversation_1\" # @param {type:\"string\"}\n", "# fmt: on\n", "\n", "request = geminidataanalytics.DeleteConversationRequest(\n", " name=data_chat_client.conversation_path(\n", " billing_project, location, conversation_id\n", " ),\n", ")\n", "\n", "# Make the request\n", "response = data_chat_client.delete_conversation(request=request)\n", "\n", "# Handle the response\n", "print(response)" ], "metadata": { "id": "VutwJEtlbTP1", "cellView": "form" }, "execution_count": null, "outputs": [] } ], "metadata": { "colab": { "collapsed_sections": [ "xMjbfg7YUnTc", "Q2mALm-giwgj", "LdN2o3qUi99z", "NKmRVLyDlt0n" ], "name": "intro_gemini_data_analytics_sdk.ipynb", "toc_visible": true, "provenance": [ { "file_id": "/piper/depot/google3/google/cloud/geminidataanalytics/colab/intro_gemini_data_analytics_sdk.ipynb", "timestamp": 1760373388601 } ], "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": 0 }