When you send a query, the API returns a stream of Message objects. This stream can contain different types of messages, including text, data tables, and charts. Text messages can provide insight into the agent's reasoning, report on its progress, or give the final answer. The purpose of each text message is indicated by its TextType value:
THOUGHT: Shows the agent's internal thought process as it plans how to answer your query.THOUGHTmessages provide step-by-step insight into the agent's reasoning and decision-making process and contain two parts:parts[0]is the thought summary, which briefly summarizes the full thought text, andparts[1]is the full thought text.PROGRESS: Reports the agent's progress on an action, such as data retrieval or a tool that's being invoked. This value is returned only for Looker data sources and contains two parts:parts[0]is the summary, andparts[1]is the full progress text.FINAL_RESPONSE: Provides the final answer to your query.
The following code examples define helper functions to process and display each message in the stream. For guidance on how to render these messages in a user interface when you're using Looker data sources, see Render agent responses for Looker data sources.
Define helper functions
The following sample code contains helper function definitions to parse the response from the API and display the results:
Python SDK
import json as json_lib
import textwrap
import time
import altair as alt
import IPython
import pandas as pd
import proto
import requests
from google.protobuf.json_format import MessageToDict, MessageToJson
from IPython.display import display, HTML
from pygments import highlight, lexers, formatters
from google.protobuf import field_mask_pb2
from google.iam.v1 import policy_pb2
from google.iam.v1 import iam_policy_pb2
def handle_text_response(resp):
parts = resp.parts
full_text = "".join(parts)
if "\n" not in full_text and len(full_text) > 80:
wrapped_text = textwrap.fill(full_text, width=80)
print(wrapped_text)
else:
print(full_text)
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(f"Query name: {query.name}")
if "question" in query:
print(f"Question: {query.question}")
if "datasources" in query:
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')
HTTP
import json as json_lib
import textwrap
import altair as alt
import IPython
from IPython.display import display, HTML
import pandas as pd
from pygments import highlight, lexers, formatters
import requests
def is_json(str):
try:
json_object = json_lib.loads(str)
except ValueError as e:
return False
return True
def handle_text_response(resp):
parts = resp['parts']
full_text = "".join(parts)
if "\n" not in full_text and len(full_text) > 80:
wrapped_text = textwrap.fill(full_text, width=80)
print(wrapped_text)
else:
print(full_text)
def get_property(data, field_name, default = ''):
return data[field_name] if field_name in data else default
def display_schema(data):
fields = data['fields']
df = pd.DataFrame({
"Column": map(lambda field: get_property(field, 'name'), fields),
"Type": map(lambda field: get_property(field, 'type'), fields),
"Description": map(lambda field: get_property(field, 'description', '-'), fields),
"Mode": map(lambda field: get_property(field, 'mode'), fields)
})
display(df)
def display_section_title(text):
display(HTML('<h2>{}</h2>'.format(text)))
def format_bq_table_ref(table_ref):
return '{}.{}.{}'.format(table_ref['projectId'], table_ref['datasetId'], table_ref['tableId'])
def format_looker_table_ref(table_ref):
return 'lookmlModel: {}, explore: {}, lookerInstanceUri: {}'.format(table_ref['lookmlModel'], table_ref['explore'], table_ref['lookerInstanceUri'])
def display_datasource(datasource):
source_name = ''
if 'studioDatasourceId' in datasource:
source_name = datasource['studioDatasourceId']
elif 'lookerExploreReference' in datasource:
source_name = format_looker_table_ref(datasource['lookerExploreReference'])
else:
source_name = format_bq_table_ref(datasource['bigqueryTableReference'])
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']))
if 'question' in query:
print('Question: {}'.format(query['question']))
if 'datasources' in query:
print('Data sources:')
for datasource in query['datasources']:
display_datasource(datasource)
elif 'generatedSql' in resp:
display_section_title('SQL generated')
print(resp['generatedSql'])
elif 'result' in resp:
display_section_title('Data retrieved')
fields = map(lambda field: get_property(field, 'name'), resp['result']['schema']['fields'])
dict = {}
for field in fields:
dict[field] = map(lambda el: get_property(el, field), resp['result']['data'])
display(pd.DataFrame(dict))
def handle_chart_response(resp):
if 'query' in resp:
print(resp['query']['instructions'])
elif 'result' in resp:
vegaConfig = resp['result']['vegaConfig']
alt.Chart.from_json(json_lib.dumps(vegaConfig)).display();
def handle_error(resp):
display_section_title('Error')
print('Code: {}'.format(resp['code']))
print('Message: {}'.format(resp['message']))
def get_stream(url, json):
s = requests.Session()
acc = ''
with s.post(url, json=json, headers=headers, stream=True) as resp:
for line in resp.iter_lines():
if not line:
continue
decoded_line = str(line, encoding='utf-8')
if decoded_line == '[{':
acc = '{'
elif decoded_line == '}]':
acc += '}'
elif decoded_line == ',':
continue
else:
acc += decoded_line
if not is_json(acc):
continue
data_json = json_lib.loads(acc)
if not 'systemMessage' in data_json:
if 'error' in data_json:
handle_error(data_json['error'])
continue
if 'text' in data_json['systemMessage']:
handle_text_response(data_json['systemMessage']['text'])
elif 'schema' in data_json['systemMessage']:
handle_schema_response(data_json['systemMessage']['schema'])
elif 'data' in data_json['systemMessage']:
handle_data_response(data_json['systemMessage']['data'])
elif 'chart' in data_json['systemMessage']:
handle_chart_response(data_json['systemMessage']['chart'])
else:
colored_json = highlight(acc, lexers.JsonLexer(), formatters.TerminalFormatter())
print(colored_json)
print('\n')
acc = ''
def get_stream_multi_turn(url, json, conversation_messages):
s = requests.Session()
acc = ''
with s.post(url, json=json, headers=headers, stream=True) as resp:
for line in resp.iter_lines():
if not line:
continue
decoded_line = str(line, encoding='utf-8')
if decoded_line == '[{':
acc = '{'
elif decoded_line == '}]':
acc += '}'
elif decoded_line == ',':
continue
else:
acc += decoded_line
if not is_json(acc):
continue
data_json = json_lib.loads(acc)
# Store the response that will be used in the next iteration
conversation_messages.append(data_json)
if not 'systemMessage' in data_json:
if 'error' in data_json:
handle_error(data_json['error'])
continue
if 'text' in data_json['systemMessage']:
handle_text_response(data_json['systemMessage']['text'])
elif 'schema' in data_json['systemMessage']:
handle_schema_response(data_json['systemMessage']['schema'])
elif 'data' in data_json['systemMessage']:
handle_data_response(data_json['systemMessage']['data'])
elif 'chart' in data_json['systemMessage']:
handle_chart_response(data_json['systemMessage']['chart'])
else:
colored_json = highlight(acc, lexers.JsonLexer(), formatters.TerminalFormatter())
print(colored_json)
print('\n')
acc = ''