使用 Python SDK 建構資料代理程式

本頁說明如何使用 Python SDK 向 Conversational Analytics API 發出要求。這個 Python 程式碼範例示範如何完成下列工作：

驗證及設定環境

如要使用 Python SDK for Conversational Analytics API，請按照對話式數據分析 API SDK Colaboratory 筆記本中的操作說明下載並安裝 SDK。請注意，SDK Colab 的下載方法和內容可能會有所變更。

完成筆記本中的設定說明後，您可以使用下列程式碼匯入必要的 SDK 程式庫、在 Colaboratory 環境中驗證 Google 帳戶，以及初始化用戶端以發出 API 要求：

from google.colab import auth
auth.authenticate_user()

from google.cloud import geminidataanalytics

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

指定帳單專案和系統指示

下列 Python 程式碼範例會定義整個指令碼中使用的帳單專案和系統指令：

# Billing project
billing_project = "my_project_name"

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

請依下列方式替換範例值：

  • my_project_name：已啟用必要 API 的帳單專案 ID。
  • Help the user analyze their data.：系統指令，可引導代理程式的行為，並根據您的資料需求進行自訂。舉例來說，您可以使用系統指令定義業務用語、控制回覆長度，或設定資料格式。建議您使用「撰寫有效的系統指令」中建議的 YAML 格式定義系統指令，提供詳細且結構化的指引。

連結至資料來源

下列 Python 程式碼範例說明如何定義 Looker、BigQuery 或 Looker Studio 資料來源的連線詳細資料，以便代理程式查詢資料來回答問題。

連結至 Looker 資料

下列程式碼範例說明如何使用 API 金鑰或存取權杖，定義 Looker 探索的連線詳細資料。使用 Conversational Analytics API 時，一次只能連結一個 Looker 探索。

API 金鑰

您可以透過產生的 Looker API 金鑰與 Looker 執行個體建立連線，詳情請參閱「使用 Conversational Analytics API 驗證及連結資料來源」。

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

請依下列方式替換範例值：

  • my_looker_client_id：您產生的 Looker API 金鑰用戶端 ID。
  • my_looker_client_secret：您產生的 Looker API 金鑰用戶端密鑰。
  • https://my_company.looker.com：Looker 執行個體的完整網址。
  • my_model：包含要連結的「探索」的 LookML 模型名稱。
  • my_explore：您希望資料代理程式查詢的 Looker 探索名稱。

存取權杖

如要與 Looker 執行個體建立連線，請使用存取權權杖，詳情請參閱「使用 Conversational Analytics API 驗證及連線至資料來源」。

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

請依下列方式替換範例值：

  • my_access_token：您產生的 access_token 值，用於向 Looker 驗證。
  • https://my_company.looker.com：Looker 執行個體的完整網址。
  • my_model：包含要連結的「探索」的 LookML 模型名稱。
  • my_explore：您希望資料代理程式查詢的 Looker 探索名稱。

連結至 BigQuery 資料

透過 Conversational Analytics API，您一次最多可以連結及查詢 10 個 BigQuery 資料表。

下列程式碼範例定義與單一 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

請依下列方式替換範例值：

  • my_project_id：包含要連結的 BigQuery 資料集和資料表的專案 ID。 Google Cloud 如要連線至公開資料集，請指定 bigquery-public-data。
  • my_dataset_id：BigQuery 資料集的 ID。例如：san_francisco。
  • my_table_id：BigQuery 資料表的 ID。例如：street_trees。

連結至 Looker Studio 資料

下列程式碼範例定義了與 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]

在上一個範例中，請將 my_datasource_id 替換為資料來源 ID。

設定有狀態或無狀態即時通訊的背景資訊

對話式數據分析 API 支援多輪對話，可讓使用者根據先前的脈絡提出後續問題。下列 Python 程式碼範例示範如何為有狀態或無狀態對話設定環境：

  • 有狀態的對話： Google Cloud 儲存及管理對話記錄。有狀態的即時通訊本質上是多輪對話，因為 API 會保留先前訊息的背景資訊。您只需要在每個回合中傳送當前訊息。
  • 無狀態對話：應用程式會管理對話記錄。每則新訊息都必須包含完整的對話記錄。如需在無狀態模式下管理多輪對話的詳細範例，請參閱「建立無狀態多輪對話」。

有狀態的對話

下列程式碼範例會設定有狀態的對話內容，其中 Google Cloud 會儲存及管理對話記錄。您也可以在下列程式碼範例中加入 published_context.options.analysis.python.enabled = True 行，選擇啟用 Python 進階分析。

# 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

無狀態對話

下列程式碼範例會設定無狀態的即時通訊情境，您必須在每則訊息中傳送完整的對話記錄。您也可以在下列程式碼範例中加入 inline_context.options.analysis.python.enabled = True 行，選擇啟用 Python 進階分析。

# 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

建立資料代理程式

下列 Python 程式碼範例會發出 API 要求來建立資料代理程式，您之後就能使用該代理程式，與資料進行對話。資料代理程式會根據指定的資料來源、系統指令和脈絡進行設定。

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

在先前的範例中，請將 data_agent_1 值替換為資料代理程式的專屬 ID。

建立對話

下列 Python 程式碼範例會發出 API 要求，建立對話。

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

請依下列方式替換範例值：

  • data_agent_1：資料代理程式的 ID，如「建立資料代理程式」中的程式碼範例區塊所定義。
  • conversation_1：對話的專屬 ID。

管理資料代理程式和對話

下列程式碼範例說明如何使用 Conversational Analytics API 管理資料代理程式和對話。您可以執行下列工作：

取得資料虛擬服務專員

下列 Python 程式碼範例示範如何發出 API 要求，擷取先前建立的資料代理程式。

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

在先前的範例中，請將 data_agent_1 值替換為要擷取的資料代理程式專屬 ID。

列出資料代理

下列程式碼示範如何呼叫 list_data_agents 方法，列出指定專案的所有資料代理程式。如要列出所有代理程式，您必須具備專案的 geminidataanalytics.dataAgents.list 權限。如要進一步瞭解哪些 IAM 角色包含這項權限，請參閱預先定義的角色清單。

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)

將 YOUR-BILLING-PROJECT 替換為帳單專案的 ID。

更新資料代理程式

以下範例程式碼示範如何呼叫資料代理程式資源的 update_data_agent 方法，更新資料代理程式。這項要求需要 DataAgent 物件，其中包含要變更的欄位新值，以及 update_mask 參數，該參數會採用 FieldMask 物件來指定要更新的欄位。

如要更新資料代理程式，您必須對該代理程式具備 geminidataanalytics.dataAgents.update IAM 權限。如要進一步瞭解哪些 IAM 角色包含這項權限，請參閱預先定義的角色清單。

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

請依下列方式替換範例值：

  • data_agent_1：要更新的資料代理程式 ID。
  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • Updated description of the data agent.：更新後的資料代理程式說明。

設定資料代理程式的 IAM 政策

如要共用代理程式，可以使用 set_iam_policy 方法，將 IAM 角色指派給特定代理程式的使用者。要求包含繫結，可指定要將哪些角色指派給哪些使用者。

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

請依下列方式替換範例值：

  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • data_agent_1：您要設定 IAM 政策的資料代理程式 ID。
  • 222larabrown@gmail.com, cloudysanfrancisco@gmail.com：以半形逗號分隔的使用者電子郵件地址清單，您要將指定角色授予這些使用者。

取得資料代理程式的身分與存取權管理政策

下列程式碼範例示範如何使用 get_iam_policy 方法，擷取資料代理程式的 IAM 政策。要求會指定資料代理程式資源路徑。

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

請依下列方式替換範例值：

  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • data_agent_1：您要取得 IAM 政策的資料代理程式 ID。

刪除資料代理程式

下列程式碼範例說明如何使用 delete_data_agent 方法，以軟刪除方式刪除資料代理程式。虛刪除代理程式時，該代理程式會遭到刪除，但仍可在 30 天內擷取。要求會指定資料代理程式資源網址。

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

請依下列方式替換範例值：

  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • data_agent_1：要刪除的資料代理程式 ID。

取得對話

下列程式碼範例示範如何使用 get_conversation 方法擷取現有對話的相關資訊。要求會指定對話資源路徑。

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)

請依下列方式替換範例值：

  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • conversation_1：要擷取的對話 ID。

列出對話

下列程式碼範例示範如何呼叫 list_conversations 方法，列出特定專案的對話。要求會指定父項資源網址，也就是專案和位置 (例如 projects/my-project/locations/global)。

根據預設，這個方法會傳回您建立的對話。管理員 (具備 cloudaicompanion.topicAdmin IAM 角色的使用者) 可以查看專案中的所有對話。

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)

將 YOUR-BILLING-PROJECT 替換為已啟用必要 API 的帳單專案 ID。

列出對話中的訊息

下列程式碼範例示範如何使用 list_messages 方法，擷取對話中的所有訊息。要求會指定對話資源路徑。

如要列出訊息，您必須具備對話的 cloudaicompanion.topics.get 權限。

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)

請依下列方式替換範例值：

  • YOUR-BILLING-PROJECT：計費專案的 ID。
  • conversation_1：您要列出訊息的對話 ID。

使用 API 提問

建立資料代理程式和對話後，下列 Python 程式碼範例會將查詢傳送至代理程式。程式碼會使用您為有狀態或無狀態即時通訊設定的內容。API 會傳回一連串訊息，代表代理程式為回答查詢所採取的步驟。

有狀態的對話

傳送含有 Conversation 參照的具狀態即時通訊要求

您可以參照先前建立的 Conversation 資源，將具狀態的即時通訊要求傳送給資料代理程式。

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

請依下列方式替換範例值：

  • Which species of tree is most prevalent?：要傳送給資料代理程式的自然語言問題。
  • data_agent_1：資料代理程式的專屬 ID，如「建立資料代理程式」一文所述。
  • conversation_1：對話的專屬 ID，如「建立對話」中所定義。

無狀態對話

下列程式碼範例示範在設定無狀態對話的內容後，如何將查詢傳送至資料代理程式。您可以參照先前定義的 DataAgent 資源，或在要求中使用內嵌環境，傳送無狀態查詢。

傳送含有 DataAgent 參照的無狀態即時通訊要求

您可以參照先前建立的 DataAgent 資源，將查詢傳送至資料代理程式。

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

請依下列方式替換範例值：

  • Which species of tree is most prevalent?：要傳送給資料代理程式的自然語言問題。
  • data_agent_1：資料代理程式的專屬 ID，如「建立資料代理程式」一文所述。

傳送內含內嵌背景資訊的無狀態即時通訊要求

以下程式碼範例說明如何使用 inline_context 參數，直接在無狀態的即時通訊要求中提供背景資訊。

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

在上一個範例中，請將 Which species of tree is most prevalent? 換成要傳送給資料代理程式的自然語言問題。

建立無狀態多輪對話

如要在無狀態對話中詢問後續問題，應用程式必須管理對話內容，方法是在每次提出新要求時，傳送完整的訊息記錄。以下範例說明如何參照資料代理程式，或使用內嵌背景資訊直接提供資料來源，建立多輪對話。

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

在先前的範例中，請依下列方式替換範例值：

  • data_agent_1：資料代理程式的專屬 ID，如「建立資料代理程式」中的程式碼範例區塊所定義。
  • Which species of tree is most prevalent?：要傳送給資料代理程式的自然語言問題。
  • Can you show me the results as a bar chart?：根據或修正先前問題的後續問題。

定義輔助函式

下列程式碼範例包含先前程式碼範例中使用的輔助函式定義。這些函式有助於剖析 API 的回應並顯示結果。

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