Cómo compilar un agente de datos con el SDK de Python

En esta página, se muestra cómo usar el SDK de Python para realizar solicitudes a la API de Conversational Analytics. En el siguiente ejemplo de código de Python, se muestra cómo completar las siguientes tareas:

Autentica y configura tu entorno

Para usar el SDK de Python para la API de Conversational Analytics, sigue las instrucciones del notebook de Colaboratory del SDK de la API de Conversational Analytics para descargar e instalar el SDK. Ten en cuenta que el método de descarga y el contenido del Colab del SDK están sujetos a cambios.

Después de completar las instrucciones de configuración en el notebook, puedes usar el siguiente código para importar las bibliotecas del SDK requeridas, autenticar tu Cuenta de Google en un entorno de Colaboratory y, luego, inicializar un cliente para realizar solicitudes a la API:

from google.colab import auth
auth.authenticate_user()

from google.cloud import geminidataanalytics

data_agent_client = geminidataanalytics.DataAgentServiceClient()
data_chat_client = geminidataanalytics.DataChatServiceClient()

Especifica el proyecto de facturación y las instrucciones del sistema

En el siguiente código de muestra en Python, se definen el proyecto de facturación y las instrucciones del sistema que se usan en toda la secuencia de comandos:

# Billing project
billing_project = "my_project_name"

# System instructions
system_instruction = "Help the user analyze their data."

Reemplaza los valores de muestra de la siguiente manera:

  • my_project_name: Es el ID de tu proyecto de facturación que tiene las APIs requeridas habilitadas.
  • Help the user analyze their data.: Son instrucciones del sistema para guiar el comportamiento del agente y personalizarlo según tus necesidades de datos. Por ejemplo, puedes usar instrucciones del sistema para definir términos comerciales, controlar la longitud de la respuesta o configurar el formato de los datos. Lo ideal es definir las instrucciones del sistema con el formato YAML recomendado en Cómo escribir instrucciones del sistema eficaces para proporcionar orientación detallada y estructurada.

Conectándose a una fuente de datos

En los siguientes ejemplos de código de Python, se muestra cómo definir los detalles de conexión para la fuente de datos de Looker, BigQuery o Looker Studio que consultará tu agente para responder preguntas.

Conéctate a los datos de Looker

En los siguientes ejemplos de código, se muestra cómo definir los detalles de una conexión a un Explorar de Looker con claves de API o un token de acceso. Con la API de Conversational Analytics, solo puedes conectarte a una exploración de Looker a la vez.

Claves de API

Puedes establecer una conexión con una instancia de Looker con claves de API de Looker generadas, como se describe en Autentica y conéctate a una fuente de datos con la API de Conversational Analytics.

looker_client_id = "my_looker_client_id"
looker_client_secret = "my_looker_client_secret"
looker_instance_uri = "https://my_company.looker.com"
lookml_model = "my_model"
explore = "my_explore"

looker_explore_reference = geminidataanalytics.LookerExploreReference()
looker_explore_reference.looker_instance_uri = looker_instance_uri
looker_explore_reference.lookml_model = lookml_model
looker_explore_reference.explore = explore

credentials = geminidataanalytics.Credentials()
credentials.oauth.secret.client_id = looker_client_id
credentials.oauth.secret.client_secret = looker_client_secret

datasource_references = geminidataanalytics.DatasourceReferences()
datasource_references.looker.explore_references = [looker_explore_reference]

# Do not include the following line during agent creation
datasource_references.credentials = credentials

Reemplaza los valores de muestra de la siguiente manera:

  • my_looker_client_id: Es el ID de cliente de la clave de la API de Looker que generaste.
  • my_looker_client_secret: Es el secreto del cliente de la clave de API de Looker que generaste.
  • https://my_company.looker.com: Es la URL completa de tu instancia de Looker.
  • my_model: Es el nombre del modelo de LookML que incluye la exploración a la que deseas conectarte.
  • my_explore: Es el nombre de la exploración de Looker que deseas que consulte el agente de datos.

Token de acceso

Puedes establecer una conexión con una instancia de Looker usando un token de acceso, como se describe en Cómo autenticarse y conectarse a una fuente de datos con la API de Conversational Analytics.

looker_access_token = "my_access_token"
looker_instance_uri = "https://my_company.looker.com"
lookml_model = "my_model"
explore = "my_explore"

looker_explore_reference = geminidataanalytics.LookerExploreReference()
looker_explore_reference.looker_instance_uri = looker_instance_uri
looker_explore_reference.lookml_model = lookml_model
looker_explore_reference.explore = explore

credentials = geminidataanalytics.Credentials()
credentials.oauth.token.access_token = looker_access_token

datasource_references = geminidataanalytics.DatasourceReferences()
datasource_references.looker.explore_references = [looker_explore_reference]

# Do not include the following line during agent creation
datasource_references.credentials = credentials

Reemplaza los valores de muestra de la siguiente manera:

  • my_access_token: Es el valor de access_token que generas para autenticarte en Looker.
  • https://my_company.looker.com: Es la URL completa de tu instancia de Looker.
  • my_model: Es el nombre del modelo de LookML que incluye la exploración a la que deseas conectarte.
  • my_explore: Es el nombre de la exploración de Looker que deseas que consulte el agente de datos.

Conéctate a los datos de BigQuery

Con la API de Conversational Analytics, puedes conectarte a un máximo de 10 tablas de BigQuery y realizar consultas en ellas al mismo tiempo.

En el siguiente código de ejemplo, se define una conexión a una sola tabla de BigQuery.

bigquery_table_reference = geminidataanalytics.BigQueryTableReference()
bigquery_table_reference.project_id = "my_project_id"
bigquery_table_reference.dataset_id = "my_dataset_id"
bigquery_table_reference.table_id = "my_table_id"

bigquery_table_reference_2 = geminidataanalytics.BigQueryTableReference()
bigquery_table_reference_2.project_id = "my_project_id_2"
bigquery_table_reference_2.dataset_id = "my_dataset_id_2"
bigquery_table_reference_2.table_id = "my_table_id_2"

# Connect to your data source
datasource_references = geminidataanalytics.DatasourceReferences()
datasource_references.bq.table_references = [bigquery_table_reference, bigquery_table_reference_2] # Up to 10 tables

Reemplaza los valores de muestra de la siguiente manera:

  • my_project_id: Es el ID del proyecto Google Cloud que contiene el conjunto de datos y la tabla de BigQuery a los que deseas conectarte. Para conectarte a un conjunto de datos públicos, especifica bigquery-public-data.
  • my_dataset_id: el ID del conjunto de datos de BigQuery. Por ejemplo, san_francisco
  • my_table_id: el ID de la tabla de BigQuery. Por ejemplo, street_trees

Conéctate a los datos de Looker Studio

En el siguiente código de muestra, se define una conexión a una fuente de datos de Looker Studio.

studio_datasource_id = "my_datasource_id"

studio_references = geminidataanalytics.StudioDatasourceReference()
studio_references.datasource_id = studio_datasource_id

## Connect to your data source
datasource_references.studio.studio_references = [studio_references]

En el ejemplo anterior, reemplaza my_datasource_id por el ID de la fuente de datos.

Configura el contexto para el chat con estado o sin estado

La API de Conversational Analytics admite conversaciones de varios turnos, lo que permite a los usuarios hacer preguntas de seguimiento que se basan en el contexto anterior. En el siguiente código de muestra de Python, se muestra cómo configurar el contexto para un chat con estado o sin estado:

  • Chat con estado: Google Cloud Almacena y administra el historial de conversaciones. El chat con estado es inherentemente de varios turnos, ya que la API retiene el contexto de los mensajes anteriores. Solo debes enviar el mensaje actual en cada turno.
  • Chat sin estado: Tu aplicación administra el historial de conversaciones. Debes incluir el historial de conversación completo con cada mensaje nuevo. Para obtener ejemplos detallados sobre cómo administrar conversaciones de varios turnos en modo sin estado, consulta Crea una conversación de varios turnos sin estado.

Chat con estado

En el siguiente muestra de código, se configura el contexto para un chat con estado, en el que Google Cloud almacena y administra el historial de conversaciones. También puedes habilitar de forma opcional el análisis avanzado con Python. Para ello, incluye la línea published_context.options.analysis.python.enabled = True en el siguiente código de muestra.

# Set up context for stateful chat
published_context = geminidataanalytics.Context()
published_context.system_instruction = system_instruction
published_context.datasource_references = datasource_references
# Optional: To enable advanced analysis with Python, include the following line:
published_context.options.analysis.python.enabled = True

Chat sin estado

En el siguiente código de muestra, se configura el contexto para el chat sin estado, en el que debes enviar el historial de conversación completo con cada mensaje. También puedes habilitar de forma opcional el análisis avanzado con Python. Para ello, incluye la línea inline_context.options.analysis.python.enabled = True en el siguiente código de muestra.

# Set up context for stateless chat
# datasource_references.looker.credentials = credentials
inline_context = geminidataanalytics.Context()
inline_context.system_instruction = system_instruction
inline_context.datasource_references = datasource_references
# Optional: To enable advanced analysis with Python, include the following line:
inline_context.options.analysis.python.enabled = True

Crea un agente de datos

El siguiente código de Python de ejemplo realiza una solicitud a la API para crear un agente de datos, que luego puedes usar para tener una conversación sobre tus datos. El agente de datos se configura con la fuente de datos, las instrucciones del sistema y el contexto especificados.

data_agent_id = "data_agent_1"

data_agent = geminidataanalytics.DataAgent()
data_agent.data_analytics_agent.published_context = published_context
data_agent.name = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}" # Optional

request = geminidataanalytics.CreateDataAgentRequest(
    parent=f"projects/{billing_project}/locations/global",
    data_agent_id=data_agent_id, # Optional
    data_agent=data_agent,
)

try:
    data_agent_client.create_data_agent(request=request)
    print("Data Agent created")
except Exception as e:
    print(f"Error creating Data Agent: {e}")

En el ejemplo anterior, reemplaza el valor data_agent_1 por un identificador único para el agente de datos.

Cómo crear una conversación

En el siguiente código de Python de ejemplo, se realiza una solicitud a la API para crear una conversación.

# Initialize request arguments
data_agent_id = "data_agent_1"
conversation_id = "conversation_1"

conversation = geminidataanalytics.Conversation()
conversation.agents = [f'projects/{billing_project}/locations/global/dataAgents/{data_agent_id}']
conversation.name = f"projects/{billing_project}/locations/global/conversations/{conversation_id}"

request = geminidataanalytics.CreateConversationRequest(
    parent=f"projects/{billing_project}/locations/global",
    conversation_id=conversation_id,
    conversation=conversation,
)

# Make the request
response = data_chat_client.create_conversation(request=request)

# Handle the response
print(response)

Reemplaza los valores de muestra de la siguiente manera:

  • data_agent_1: Es el ID del agente de datos, como se define en el bloque de código de ejemplo en Crea un agente de datos.
  • conversation_1: Es un identificador único de la conversación.

Administra agentes de datos y conversaciones

En los siguientes ejemplos de código, se muestra cómo administrar tus agentes de datos y conversaciones con la API de Conversational Analytics. Puedes realizar las siguientes tareas:

Obtén un agente de datos

En la siguiente muestra de código en Python, se muestra cómo realizar una solicitud a la API para recuperar un agente de datos que creaste anteriormente.

# Initialize request arguments
data_agent_id = "data_agent_1"
request = geminidataanalytics.GetDataAgentRequest(
    name=f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}",
)

# Make the request
response = data_agent_client.get_data_agent(request=request)

# Handle the response
print(response)

En el ejemplo anterior, reemplaza el valor data_agent_1 por el identificador único del agente de datos que deseas recuperar.

Enumera los agentes de datos

En el siguiente código, se muestra cómo enumerar todos los agentes de datos de un proyecto determinado llamando al método list_data_agents. Para enumerar todos los agentes, debes tener el permiso geminidataanalytics.dataAgents.list en el proyecto. Para obtener más información sobre qué roles de IAM incluyen este permiso, consulta la lista de roles predefinidos.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
request = geminidataanalytics.ListDataAgentsRequest(
    parent=f"projects/{billing_project}/locations/global",
)

# Make the request
page_result = data_agent_client.list_data_agents(request=request)

# Handle the response
for response in page_result:
    print(response)

Reemplaza YOUR-BILLING-PROJECT por el ID de tu proyecto de facturación.

Actualiza un agente de datos

En el siguiente código de muestra, se muestra cómo actualizar un agente de datos llamando al método update_data_agent en el recurso del agente de datos. La solicitud requiere un objeto DataAgent que incluya los valores nuevos para los campos que deseas cambiar y un parámetro update_mask que tome un objeto FieldMask para especificar qué campos actualizar.

Para actualizar un agente de datos, debes tener el permiso geminidataanalytics.dataAgents.update de IAM en el agente. Para obtener más información sobre qué roles de IAM incluyen este permiso, consulta la lista de roles predefinidos.

data_agent_id = "data_agent_1"
billing_project = "YOUR-BILLING-PROJECT"
data_agent = geminidataanalytics.DataAgent()
data_agent.data_analytics_agent.published_context = published_context
data_agent.name = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
data_agent.description = "Updated description of the data agent."

update_mask = field_mask_pb2.FieldMask(paths=['description', 'data_analytics_agent.published_context'])

request = geminidataanalytics.UpdateDataAgentRequest(
    data_agent=data_agent,
    update_mask=update_mask,
)

try:
    # Make the request
    data_agent_client.update_data_agent(request=request)
    print("Data Agent Updated")
except Exception as e:
    print(f"Error updating Data Agent: {e}")

Reemplaza los valores de muestra de la siguiente manera:

  • data_agent_1: Es el ID del agente de datos que deseas actualizar.
  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • Updated description of the data agent.: Es una descripción del agente de datos actualizado.

Configura la política de IAM para un agente de datos

Para compartir un agente, puedes usar el método set_iam_policy para asignar roles de IAM a los usuarios en un agente específico. La solicitud incluye vinculaciones que especifican qué roles se deben asignar a qué usuarios.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
data_agent_id = "data_agent_1"
role = "roles/geminidataanalytics.dataAgentEditor"
users = "222larabrown@gmail.com, cloudysanfrancisco@gmail.com"

resource = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"

# Construct the IAM policy
binding = policy_pb2.Binding(
    role=role,
    members= [f"user:{i.strip()}" for i in users.split(",")]
)

policy = policy_pb2.Policy(bindings=[binding])

# Create the request
request = iam_policy_pb2.SetIamPolicyRequest(
    resource=resource,
    policy=policy
)

# Send the request
try:
    response = data_agent_client.set_iam_policy(request=request)
    print("IAM Policy set successfully!")
    print(f"Response: {response}")
except Exception as e:
    print(f"Error setting IAM policy: {e}")

Reemplaza los valores de muestra de la siguiente manera:

  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • data_agent_1: Es el ID del agente de datos para el que deseas establecer la política de IAM.
  • 222larabrown@gmail.com, cloudysanfrancisco@gmail.com: Es una lista separada por comas de las direcciones de correo electrónico de los usuarios a los que deseas otorgar el rol especificado.

Obtén la política de IAM para un agente de datos

En el siguiente código de ejemplo, se muestra cómo usar el método get_iam_policy para recuperar la política de IAM de un agente de datos. La solicitud especifica la ruta del recurso del agente de datos.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
data_agent_id = "data_agent_1"

resource = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
request = iam_policy_pb2.GetIamPolicyRequest(
            resource=resource,
        )
try:
    response = data_agent_client.get_iam_policy(request=request)
    print("IAM Policy fetched successfully!")
    print(f"Response: {response}")
except Exception as e:
    print(f"Error setting IAM policy: {e}")

Reemplaza los valores de muestra de la siguiente manera:

  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • data_agent_1: Es el ID del agente de datos para el que deseas obtener la política de IAM.

Borra un agente de datos

En el siguiente código de ejemplo, se muestra cómo usar el método delete_data_agent para borrar de forma no definitiva un agente de datos. Cuando borras de forma no definitiva un agente, este se borra, pero se puede recuperar en un plazo de 30 días. La solicitud especifica la URL del recurso del agente de datos.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
data_agent_id = "data_agent_1"

request = geminidataanalytics.DeleteDataAgentRequest(
    name=f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}",
)

try:
    # Make the request
    data_agent_client.delete_data_agent(request=request)
    print("Data Agent Deleted")
except Exception as e:
    print(f"Error deleting Data Agent: {e}")

Reemplaza los valores de muestra de la siguiente manera:

  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • data_agent_1: Es el ID del agente de datos que deseas borrar.

Obtén una conversación

En el siguiente código de ejemplo, se muestra cómo usar el método get_conversation para recuperar información sobre una conversación existente. La solicitud especifica la ruta de acceso al recurso de conversación.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
conversation_id = "conversation_1"

request = geminidataanalytics.GetConversationRequest(
    name = f"projects/{billing_project}/locations/global/conversations/{conversation_id}"
)

# Make the request
response = data_chat_client.get_conversation(request=request)

# Handle the response
print(response)

Reemplaza los valores de muestra de la siguiente manera:

  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • conversation_1: ID de la conversación que deseas recuperar.

Mostrar lista de conversaciones

En el siguiente ejemplo de código, se muestra cómo enumerar conversaciones para un proyecto determinado llamando al método list_conversations. La solicitud especifica la URL del recurso principal, que es el proyecto y la ubicación (por ejemplo, projects/my-project/locations/global).

De forma predeterminada, este método devuelve las conversaciones que creaste. Los administradores (usuarios con el rol de IAM de cloudaicompanion.topicAdmin) pueden ver todas las conversaciones del proyecto.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"
request = geminidataanalytics.ListConversationsRequest(
    parent=f"projects/{billing_project}/locations/global",
)

# Make the request
response = data_chat_client.list_conversations(request=request)

# Handle the response
print(response)

Reemplaza YOUR-BILLING-PROJECT por el ID del proyecto de facturación en el que habilitaste las APIs requeridas.

Cómo enumerar los mensajes de una conversación

En el siguiente código de ejemplo, se muestra cómo usar el método list_messages para recuperar todos los mensajes de una conversación. La solicitud especifica la ruta de acceso al recurso de conversación.

Para enumerar los mensajes, debes tener el permiso cloudaicompanion.topics.get en la conversación.

billing_project = "YOUR-BILLING-PROJECT"
location = "global"

conversation_id = "conversation_1"

request = geminidataanalytics.ListMessagesRequest(
    parent=f"projects/{billing_project}/locations/global/conversations/{conversation_id}",
)

# Make the request
response = data_chat_client.list_messages(request=request)

# Handle the response
print(response)

Reemplaza los valores de muestra de la siguiente manera:

  • YOUR-BILLING-PROJECT: Es el ID de tu proyecto de facturación.
  • conversation_1: Es el ID de la conversación para la que deseas enumerar los mensajes.

Usa la API para hacer preguntas

Después de crear un agente de datos y una conversación, el siguiente código de Python de ejemplo envía una consulta al agente. El código usa el contexto que configuraste para el chat con o sin estado. La API devuelve un flujo de mensajes que representan los pasos que sigue el agente para responder la búsqueda.

Chat con estado

Envía una solicitud de chat con estado con una referencia Conversation

Puedes enviar una solicitud de chat con estado al agente de datos haciendo referencia a un recurso Conversation que creaste anteriormente.

# Create a request that contains a single user message (your question)
question = "Which species of tree is most prevalent?"
messages = [geminidataanalytics.Message()]
messages[0].user_message.text = question

data_agent_id = "data_agent_1"
conversation_id = "conversation_1"

# Create a conversation_reference
conversation_reference = geminidataanalytics.ConversationReference()
conversation_reference.conversation = f"projects/{billing_project}/locations/global/conversations/{conversation_id}"
conversation_reference.data_agent_context.data_agent = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
# conversation_reference.data_agent_context.credentials = credentials

# Form the request
request = geminidataanalytics.ChatRequest(
    parent = f"projects/{billing_project}/locations/global",
    messages = messages,
    conversation_reference = conversation_reference
)

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
    show_message(response)

Reemplaza los valores de muestra de la siguiente manera:

  • Which species of tree is most prevalent?: Es una pregunta en lenguaje natural que se envía al agente de datos.
  • data_agent_1: Es el identificador único del agente de datos, como se define en Crea un agente de datos.
  • conversation_1: Es el identificador único de la conversación, como se define en Crea una conversación.

Chat sin estado

En los siguientes ejemplos de código, se muestra cómo enviar una consulta al agente de datos cuando configuraste el contexto para el chat sin estado. Puedes enviar consultas sin estado haciendo referencia a un recurso DataAgent definido previamente o usando contexto intercalado en la solicitud.

Envía una solicitud de chat sin estado con una referencia DataAgent.

Puedes enviar una consulta al agente de datos haciendo referencia a un recurso DataAgent que creaste anteriormente.

# Create a request that contains a single user message (your question)
question = "Which species of tree is most prevalent?"
messages = [geminidataanalytics.Message()]
messages[0].user_message.text = question

data_agent_id = "data_agent_1"

data_agent_context = geminidataanalytics.DataAgentContext()
data_agent_context.data_agent = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
# data_agent_context.credentials = credentials

# Form the request
request = geminidataanalytics.ChatRequest(
    parent=f"projects/{billing_project}/locations/global",
    messages=messages,
    data_agent_context = data_agent_context
)

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
    show_message(response)

Reemplaza los valores de muestra de la siguiente manera:

  • Which species of tree is most prevalent?: Es una pregunta en lenguaje natural que se envía al agente de datos.
  • data_agent_1: Es el identificador único del agente de datos, como se define en Crea un agente de datos.

Envía una solicitud de chat sin estado con contexto intercalado

En el siguiente código de ejemplo, se muestra cómo usar el parámetro inline_context para proporcionar contexto directamente en tu solicitud de chat sin estado.

# Create a request that contains a single user message (your question)
question = "Which species of tree is most prevalent?"
messages = [geminidataanalytics.Message()]
messages[0].user_message.text = question

request = geminidataanalytics.ChatRequest(
    inline_context=inline_context,
    parent=f"projects/{billing_project}/locations/global",
    messages=messages,
)

# Make the request
stream = data_chat_client.chat(request=request)

# Handle the response
for response in stream:
    show_message(response)

En el ejemplo anterior, reemplaza Which species of tree is most prevalent? por una pregunta en lenguaje natural para enviarle al agente de datos.

Crea una conversación sin estado de varios turnos

Para hacer preguntas relacionadas en una conversación sin estado, tu aplicación debe administrar el contexto de la conversación enviando todo el historial de mensajes con cada solicitud nueva. En el siguiente ejemplo, se muestra cómo crear una conversación de varios turnos haciendo referencia a un agente de datos o usando el contexto intercalado para proporcionar la fuente de datos directamente.

# List that is used to track previous turns and is reused across requests
conversation_messages = []

data_agent_id = "data_agent_1"

# Use data agent context
data_agent_context = geminidataanalytics.DataAgentContext()
data_agent_context.data_agent = f"projects/{billing_project}/locations/global/dataAgents/{data_agent_id}"
# data_agent_context.credentials = credentials

# Helper function for calling the API
def multi_turn_Conversation(msg):

    message = geminidataanalytics.Message()
    message.user_message.text = msg

    # Send a multi-turn request by including previous turns and the new message
    conversation_messages.append(message)

    request = geminidataanalytics.ChatRequest(
        parent=f"projects/{billing_project}/locations/global",
        messages=conversation_messages,
        # Use data agent context
        data_agent_context=data_agent_context,
        # Use inline context
        # inline_context=inline_context,
    )

    # Make the request
    stream = data_chat_client.chat(request=request)

    # Handle the response
    for response in stream:
      show_message(response)
      conversation_messages.append(response)

# Send the first turn request
multi_turn_Conversation("Which species of tree is most prevalent?")

# Send follow-up turn request
multi_turn_Conversation("Can you show me the results as a bar chart?")

En el ejemplo anterior, reemplaza los valores de muestra de la siguiente manera:

  • data_agent_1: Es el identificador único del agente de datos, como se define en el bloque de código de muestra en Crea un agente de datos.
  • Which species of tree is most prevalent?: Es una pregunta en lenguaje natural que se envía al agente de datos.
  • Can you show me the results as a bar chart?: Es una pregunta de seguimiento que se basa en la pregunta anterior o la refina.

Define funciones auxiliares

El siguiente código de muestra contiene definiciones de funciones auxiliares que se usan en las muestras de código anteriores. Estas funciones ayudan a analizar la respuesta de la API y mostrar los resultados.

from pygments import highlight, lexers, formatters
import pandas as pd
import requests
import json as json_lib
import altair as alt
import IPython
from IPython.display import display, HTML

import proto
from google.protobuf.json_format import MessageToDict, MessageToJson

def handle_text_response(resp):
  parts = getattr(resp, 'parts')
  print(''.join(parts))

def display_schema(data):
  fields = getattr(data, 'fields')
  df = pd.DataFrame({
    "Column": map(lambda field: getattr(field, 'name'), fields),
    "Type": map(lambda field: getattr(field, 'type'), fields),
    "Description": map(lambda field: getattr(field, 'description', '-'), fields),
    "Mode": map(lambda field: getattr(field, 'mode'), fields)
  })
  display(df)

def display_section_title(text):
  display(HTML('<h2>{}</h2>'.format(text)))

def format_looker_table_ref(table_ref):
 return 'lookmlModel: {}, explore: {}, lookerInstanceUri: {}'.format(table_ref.lookml_model, table_ref.explore, table_ref.looker_instance_uri)

def format_bq_table_ref(table_ref):
  return '{}.{}.{}'.format(table_ref.project_id, table_ref.dataset_id, table_ref.table_id)

def display_datasource(datasource):
  source_name = ''
  if 'studio_datasource_id' in datasource:
   source_name = getattr(datasource, 'studio_datasource_id')
  elif 'looker_explore_reference' in datasource:
   source_name = format_looker_table_ref(getattr(datasource, 'looker_explore_reference'))
  else:
    source_name = format_bq_table_ref(getattr(datasource, 'bigquery_table_reference'))

  print(source_name)
  display_schema(datasource.schema)

def handle_schema_response(resp):
  if 'query' in resp:
    print(resp.query.question)
  elif 'result' in resp:
    display_section_title('Schema resolved')
    print('Data sources:')
    for datasource in resp.result.datasources:
      display_datasource(datasource)

def handle_data_response(resp):
  if 'query' in resp:
    query = resp.query
    display_section_title('Retrieval query')
    print('Query name: {}'.format(query.name))
    print('Question: {}'.format(query.question))
    print('Data sources:')
    for datasource in query.datasources:
      display_datasource(datasource)
  elif 'generated_sql' in resp:
    display_section_title('SQL generated')
    print(resp.generated_sql)
  elif 'result' in resp:
    display_section_title('Data retrieved')

    fields = [field.name for field in resp.result.schema.fields]
    d = {}
    for el in resp.result.data:
      for field in fields:
        if field in d:
          d[field].append(el[field])
        else:
          d[field] = [el[field]]

    display(pd.DataFrame(d))

def handle_chart_response(resp):
  def _value_to_dict(v):
    if isinstance(v, proto.marshal.collections.maps.MapComposite):
      return _map_to_dict(v)
    elif isinstance(v, proto.marshal.collections.RepeatedComposite):
      return [_value_to_dict(el) for el in v]
    elif isinstance(v, (int, float, str, bool)):
      return v
    else:
      return MessageToDict(v)

  def _map_to_dict(d):
    out = {}
    for k in d:
      if isinstance(d[k], proto.marshal.collections.maps.MapComposite):
        out[k] = _map_to_dict(d[k])
      else:
        out[k] = _value_to_dict(d[k])
    return out

  if 'query' in resp:
    print(resp.query.instructions)
  elif 'result' in resp:
    vegaConfig = resp.result.vega_config
    vegaConfig_dict = _map_to_dict(vegaConfig)
    alt.Chart.from_json(json_lib.dumps(vegaConfig_dict)).display();

def show_message(msg):
  m = msg.system_message
  if 'text' in m:
    handle_text_response(getattr(m, 'text'))
  elif 'schema' in m:
    handle_schema_response(getattr(m, 'schema'))
  elif 'data' in m:
    handle_data_response(getattr(m, 'data'))
  elif 'chart' in m:
    handle_chart_response(getattr(m, 'chart'))
  print('\n')