Criar um agente de dados usando o SDK do Python

Nesta página, mostramos como usar o SDK do Python para fazer solicitações à API Conversational Analytics. O exemplo de código Python mostra como fazer as seguintes tarefas:

Autenticar-se e configurar o ambiente

Para usar o SDK do Python para a API Conversational Analytics, siga as instruções no bloco Colaboratory do SDK da API Conversational Analytics para baixar e instalar o SDK. O método de download e o conteúdo do Colab do SDK estão sujeitos a mudanças.

Depois de seguir as instruções de configuração no bloco, use o seguinte código para importar as bibliotecas do SDK necessárias, autenticar sua Conta do Google em um ambiente do Colaboratory e inicializar um cliente para fazer solicitações de API:

from google.colab import auth
auth.authenticate_user()

from google.cloud import geminidataanalytics

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

Especificar o projeto de faturamento e as instruções do sistema

Este exemplo de código em Python define o projeto de faturamento e as instruções do sistema usadas em todo o script:

# Billing project
billing_project = "my_project_name"

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

Substitua os valores de exemplo da seguinte forma:

  • my_project_name: o ID do seu projeto de faturamento que tem as APIs necessárias ativadas.
  • Help the user analyze their data.: instruções do sistema para orientar o comportamento do agente e personalizá-lo de acordo com suas necessidades de dados. Por exemplo, é possível usar instruções do sistema para definir termos comerciais, controlar o tamanho da resposta ou definir a formatação de dados. O ideal é definir instruções do sistema usando o formato YAML recomendado em Escrever instruções eficazes para o sistema para dar orientações detalhadas e estruturadas.

Conectar a uma fonte de dados

Os exemplos de código Python a seguir mostram como definir os detalhes da conexão para a fonte de dados do Looker, do BigQuery ou do Looker Studio que o agente vai consultar para responder a perguntas.

Conectar aos dados do Looker

Os exemplos de código a seguir mostram como definir os detalhes de uma conexão com uma Análise do Looker usando chaves de API ou um token de acesso.

Chaves de API

É possível estabelecer uma conexão com uma instância do Looker usando chaves de API geradas, conforme descrito em Autenticar e se conectar a uma fonte de dados com a API 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]

Substitua os valores de exemplo da seguinte forma:

  • my_looker_client_id: o ID do cliente da chave de API do Looker gerada.
  • my_looker_client_secret: a chave secreta do cliente da chave de API do Looker gerada.
  • https://my_company.looker.com: o URL completo da sua instância do Looker.
  • my_model: o nome do modelo do LookML que inclui a Análise de destino.
  • my_explore: o nome da Análise do Looker que o agente de dados deve consultar.

Token de acesso

É possível estabelecer uma conexão com uma instância do Looker usando um token de acesso, conforme descrito em Autenticar e se conectar a uma fonte de dados com a API 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]

Substitua os valores de exemplo da seguinte forma:

  • my_access_token: o valor access_token gerado para autenticar no Looker.
  • https://my_company.looker.com: o URL completo da sua instância do Looker.
  • my_model: o nome do modelo do LookML que inclui a Análise de destino.
  • my_explore: o nome da Análise do Looker que o agente de dados deve consultar.

Conectar-se aos dados do BigQuery

Com a API Conversational Analytics, é possível acessar e consultar até 10 tabelas do BigQuery por vez.

O exemplo de código a seguir define uma conexão com uma tabela do BigQuery.

bq_project_id = "my_project_id"
bq_dataset_id = "my_dataset_id"
bq_table_id = "my_table_id"

bigquery_table_reference = geminidataanalytics.BigQueryTableReference()
bigquery_table_reference.project_id = bq_project_id
bigquery_table_reference.dataset_id = bq_dataset_id
bigquery_table_reference.table_id = bq_table_id

# Connect to your data source
datasource_references = geminidataanalytics.DatasourceReferences()
datasource_references.bq.table_references = [bigquery_table_reference]

Substitua os valores de exemplo da seguinte forma:

  • my_project_id: o ID do projeto do Google Cloud que contém o conjunto de dados e a tabela do BigQuery a que você quer se conectar. Para se conectar a um conjunto de dados público, especifique bigquery-public-data.
  • my_dataset_id: ID do conjunto de dados do BigQuery. Por exemplo, san_francisco.
  • my_table_id: ID da tabela do BigQuery. Por exemplo, street_trees.

Conectar aos dados do Looker Studio

O exemplo de código a seguir define uma conexão com uma fonte de dados do 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]

No exemplo anterior, substitua my_datasource_id pelo ID da origem de dados.

Configurar o contexto para conversa com ou sem estado

A API Conversational Analytics é compatível com conversas multiturno, em que os usuários fazem perguntas com base no contexto anterior. O exemplo de código Python a seguir demonstra como configurar o contexto para uma conversa com estado ou sem estado:

  • Conversa com estado:o Google Cloud armazena e gerencia o histórico de conversas. A conversa com estado é inerentemente multiturno, já que a API retém o contexto das mensagens anteriores. Você só precisa enviar a mensagem atual a cada turno.
  • Conversa sem estado: o aplicativo gerencia o histórico de conversas. É preciso incluir todo o histórico da conversa em cada nova mensagem. Confira exemplos detalhados de como gerenciar conversas multiturno no modo sem estado em Criar uma conversa multiturno sem estado.

Conversa com estado

O exemplo de código a seguir configura o contexto para conversa com estado, em que o Google Cloud armazena e gerencia o histórico. Também é possível ativar a análise avançada com Python incluindo a linha published_context.options.analysis.python.enabled = True no exemplo de código abaixo.

# 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

Conversa sem estado

O exemplo de código a seguir configura o contexto para conversa sem estado, em que é necessário enviar todo o histórico com cada mensagem. Também é possível ativar a análise avançada com Python incluindo a linha inline_context.options.analysis.python.enabled = True no exemplo de código abaixo.

# 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

Criar um agente de dados

O exemplo de código Python a seguir faz uma solicitação de API para criar um agente de dados, que pode ser usado para conversar sobre os dados. O agente de dados é configurado com a fonte de dados, as instruções do sistema e o 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}")

No exemplo anterior, substitua o valor data_agent_1 por um identificador exclusivo do agente de dados.

Criar uma conversa

O exemplo de código Python a seguir faz uma solicitação de API para criar uma conversa.

# 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)

Substitua os valores de exemplo da seguinte forma:

  • data_agent_1: o ID do agente de dados, conforme definido no exemplo de bloco de código em Criar um agente de dados.
  • conversation_1: um identificador exclusivo da conversa.

Gerenciar agentes de dados e conversas

Os exemplos de código a seguir mostram como gerenciar agentes de dados e conversas usando a API Conversational Analytics. É possível fazer as seguintes tarefas:

Recuperar um agente de dados

O exemplo de código Python a seguir demonstra como fazer uma solicitação de API para recuperar um agente de dados criado 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)

No exemplo anterior, substitua o valor data_agent_1 pelo identificador exclusivo do agente de dados que você quer recuperar.

Listar agentes de dados

O código a seguir demonstra como listar todos os agentes de dados de um determinado projeto com o método list_data_agents. Para listar todos os agentes, você precisa ter a permissão geminidataanalytics.dataAgents.list no projeto. Confira quais papéis do IAM incluem essa permissão na lista de papéis 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)

Substitua YOUR-BILLING-PROJECT pelo ID do seu projeto de faturamento.

Atualizar um agente de dados

O exemplo de código a seguir demonstra como atualizar um agente de dados chamando o método update_data_agent no recurso do agente. A solicitação exige um objeto DataAgent que inclua os novos valores dos campos que você quer mudar e um parâmetro update_mask que receba um objeto FieldMask para especificar quais campos atualizar.

Para atualizar um agente de dados, você precisa ter a permissão geminidataanalytics.dataAgents.update do IAM no agente. Confira quais papéis do IAM incluem essa permissão na lista de papéis 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}")

Substitua os valores de exemplo da seguinte forma:

  • data_agent_1: o ID do agente de dados que você quer atualizar.
  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • Updated description of the data agent.: uma descrição do agente de dados atualizado.

Definir a política do IAM de um agente de dados

Para compartilhar um agente, use o método set_iam_policy e atribua papéis do IAM a usuários em um agente específico. A solicitação inclui vinculações que especificam quais papéis devem ser atribuídos a quais usuários.

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}")

Substitua os valores de exemplo da seguinte forma:

  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • data_agent_1: o ID do agente de dados para o qual você quer definir a política do IAM.
  • 222larabrown@gmail.com, cloudysanfrancisco@gmail.com: uma lista separada por vírgulas de e-mails de usuários a quem você quer conceder o papel especificado.

Acessar a política do IAM de um agente de dados

O exemplo de código a seguir demonstra como usar o método get_iam_policy para buscar a política do IAM de um agente de dados. A solicitação especifica o caminho do recurso do agente de dados.

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}")

Substitua os valores de exemplo da seguinte forma:

  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • data_agent_1: o ID do agente de dados cuja política do IAM você quer buscar.

Excluir um agente de dados

O exemplo de código a seguir demonstra como usar o método delete_data_agent para excluir um agente de dados de forma reversível. Quando você faz uma exclusão reversível de um agente, ele é excluído, mas ainda pode ser recuperado em até 30 dias. A solicitação especifica o URL do recurso do agente de dados.

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}")

Substitua os valores de exemplo da seguinte forma:

  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • data_agent_1: o ID do agente de dados que você quer excluir.

Acessar uma conversa

O exemplo de código a seguir demonstra como usar o método get_conversation para buscar informações sobre uma conversa. A solicitação especifica o caminho do recurso de conversa.

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)

Substitua os valores de exemplo da seguinte forma:

  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • conversation_1: o ID da conversa que você quer buscar.

Listar conversas

O exemplo de código a seguir demonstra como listar as conversas de um determinado projeto com o método list_conversations. A solicitação especifica o URL do recurso pai, que é o projeto e o local (por exemplo, projects/my-project/locations/global).

Por padrão, esse método retorna as conversas que você criou. Os administradores (usuários com o papel do IAM cloudaicompanion.topicAdmin) podem acessar todas as conversas no projeto.

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)

Substitua YOUR-BILLING-PROJECT pelo ID do projeto de faturamento em que você ativou as APIs necessárias.

Listar mensagens em uma conversa

O exemplo de código a seguir demonstra como usar o método list_messages para buscar todas as mensagens em uma conversa. A solicitação especifica o caminho do recurso de conversa.

Para listar mensagens, é necessário ter a permissão cloudaicompanion.topics.get na conversa.

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)

Substitua os valores de exemplo da seguinte forma:

  • YOUR-BILLING-PROJECT: o ID do seu projeto de faturamento.
  • conversation_1: o ID da conversa cujas mensagens você quer listar.

Usar a API para fazer perguntas

Depois que você cria um agente de dados e uma conversa, o exemplo de código Python a seguir envia uma consulta ao agente. O código usa o contexto que você configurou para conversas com ou sem estado. A API retorna um fluxo de mensagens que representam as etapas que o agente realiza para responder à consulta.

Conversa com estado

Enviar uma solicitação de conversa com estado com uma referência Conversation

Para enviar uma solicitação de conversa com estado ao agente de dados, faça referência a um recurso Conversation já criado.

# 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)

Substitua os valores de exemplo da seguinte forma:

  • Which species of tree is most prevalent?: uma pergunta em linguagem natural para enviar ao agente de dados.
  • data_agent_1: o identificador exclusivo do agente de dados, definido em Criar um agente de dados.
  • conversation_1: o identificador exclusivo da conversa, definido em Criar uma conversa.

Conversa sem estado

Os exemplos de código a seguir mostram como enviar uma consulta ao agente de dados quando você configura o contexto para conversa sem estado. Para enviar consultas sem estado, faça referência a um recurso DataAgent já definido ou use o contexto inline na solicitação.

Enviar uma solicitação de conversa sem estado com uma referência DataAgent

Para enviar uma consulta ao agente de dados, faça referência a um recurso DataAgent já criado.

# 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)

Substitua os valores de exemplo da seguinte forma:

  • Which species of tree is most prevalent?: uma pergunta em linguagem natural para enviar ao agente de dados.
  • data_agent_1: o identificador exclusivo do agente de dados, definido em Criar um agente de dados.

Enviar uma solicitação de conversa sem estado com contexto inline

O exemplo de código a seguir demonstra como usar o parâmetro inline_context para incluir o contexto na solicitação de conversa sem 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)

No exemplo anterior, substitua Which species of tree is most prevalent? por uma pergunta em linguagem natural para enviar ao agente de dados.

Criar uma conversa multiturno sem estado

Para o usuário fazer perguntas complementares em uma conversa sem estado, o aplicativo precisa enviar todo o histórico de mensagens em cada nova solicitação. O exemplo a seguir mostra como criar uma conversa multiturno referenciando um agente de dados ou usando o contexto inline para informar a fonte de dados diretamente.

# 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?")

No exemplo anterior, substitua os valores de amostra da seguinte forma:

  • data_agent_1: o identificador exclusivo do agente de dados, definido no exemplo de bloco de código em Criar um agente de dados.
  • Which species of tree is most prevalent?: uma pergunta em linguagem natural para enviar ao agente de dados.
  • Can you show me the results as a bar chart?: uma pergunta complementar, baseada na pergunta anterior.

Definir funções auxiliares

O exemplo de código a seguir contém as definições das funções auxiliares usadas nos exemplos anteriores. As funções ajudam a analisar a resposta da API e mostrar os 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')