SQL Generation by PaLM API

main.py

import io
import sys

from google.cloud import bigquery
from vertexai.preview.language_models import TextGenerationModel


client = bigquery.Client()

def render_prompt(full_table_id: str, user_prompt: str) -> str:
    table = client.get_table(full_table_id)

    restriction = "None"
    if (partitioning := table.time_partitioning) and partitioning.require_partition_filter:
        restriction = f"You MUST filter `{partitioning.field}` field in a `where` clause."

    with io.StringIO("") as buf:
        client.schema_to_json(table.schema, buf)
        schema = buf.getvalue()

    return f"""You are an experienced data analyst. Write a BigQuery SQL to answer the user's prompt based on the following context.

---- Context ----
Format: Plain SQL only, no Markdown
Table: `{table.full_table_id}`
Restriction: {restriction}
Schema as JSON:
{schema}
----

User's prompt: {user_prompt}"""

full_table_id = sys.argv[1]
user_prompt = sys.argv[2]

prompt = render_prompt(full_table_id, user_prompt)
print(f"==== Prompt ====\n{prompt}")

model = TextGenerationModel.from_pretrained("text-bison@001")
result = model.predict(prompt)
print(f"\n==== Response ====\n{result.text}")

rows = client.query(result.text).result()
print(f"\n==== Result ====\n")
for row in rows:
    print("\t".join(map(str, row.values())))