import flask from google.cloud import bigquery from google.cloud import storage from google.cloud import secretmanager import requests from io import StringIO import os from prettytable import PrettyTable class OpenAIClient: def __init__(self, token): self.token = token def complete(self, prompt): headers = { 'Authorization': f'Bearer {self.token}', } url = 'https://api.openai.com/v1/completions' data = { 'model': 'text-davinci-003', 'prompt': prompt, 'temperature': 0, 'max_tokens': 1000 } response = requests.post(url, headers=headers, json=data) print(response.request.url) print(response.request.body) print(response.request.headers) return response.json() def get_table_desc(table): name = table['name'] columns = table['columns'] desc = table['desc'] s = StringIO() print('Table %s, description=%s, columns = [%s]' % (name, desc, ','.join([x[0] for x in columns])), file=s) for c in columns: name, ty, extra = c print('Column %s, type=%s, %s' % (name, ty, extra), file=s) return s.getvalue() ALL_CLUSTERS_TABLE = { 'name': '`dbaas-eng-bi.user_behavior_analysis.all_clusters`', 'desc': 'general cluster information', 'columns': [ ('cluster_id', 'STRING', None), ('cluster_name', 'STRING', None), ('project_id', 'INTEGER', None), ('project_name', 'STRING', None), ('tenant_id', 'INTEGER', None), ('tenant_name', 'STRING', None), ('tenant_email', 'STRING', None), ('tenant_status', 'STRING', None), ('tenant_kind', 'INTEGER', ['enum=[internal, external]']), ('cluster_created_at', 'TIMESTAMP', None), ('cluster_status', 'STRING', ['enum=[scaling, paused, normal, deletecomplete, cleared]']), ('region', 'STRING', None), ('provider_name', 'STRING', ['enum=[aws, gcp]']), ('is_dev_tier', 'BOOLEAN', None), ('tidb.tier', 'STRING', None), ('tidb.cpu_num', 'INTEGER', None), ('tidb.num_nodes', 'INTEGER', None), ('tidb.storage_size_gib', 'INTEGER', None), ('tikv.tier', 'STRING', None), ('tikv.cpu_num', 'INTEGER', None), ('tikv.num_nodes', 'INTEGER', None), ('tikv.storage_size_gib', 'INTEGER', None), ('tiflash.tier', 'STRING', None), ('tiflash.cpu_num', 'INTEGER', None), ('tiflash.num_nodes', 'INTEGER', None), ('tiflash.storage_size_gib', 'INTEGER', None), ] } EVENTS_TABLE = { 'name': '`dbaas-eng-bi.user_behavior_analysis.prod_view`', 'desc': 'clusters events or activities', 'columns': [ ('cluster_id', 'STRING', None), ('event_timestamp', 'TIMESTAMP', None), ('event_type', 'STRING', [ 'enum=[ClusterCreating, ClusterNormal, ClusterPausing, ClusterPaused, ClusterResuming, ClusterResumed, ClusterScaling, ClusterScaled]' ]), ('duration_secs', 'INTEGER', None), ('infra_provider', 'STRING', ['enum=[aws, gcp, eks, gke, aws-shared]' ]), ('first_cluster_in_project_region', 'BOOLEAN', None), ('project_id', 'INTEGER', None), ('project_name', 'STRING', None), ('tenant_id', 'INTEGER', None), ('tenant_name', 'STRING', None), ('tenant_email', 'STRING', None), ('layout', 'STRING', None), ('is_dev_tier', 'BOOLEAN', None), ('layout_diff', 'STRING', 'layout difference in the case of scaling'), ] } def get_sql(question): api_key = secretmanager.SecretManagerServiceClient().access_secret_version( name= 'projects/728951190684/secrets/yunqing-openai-api-key/versions/latest' ).payload.data.decode('UTF-8') c = OpenAIClient(api_key) r = c.complete(f''' # BigQuery Standard SQL {get_table_desc(ALL_CLUSTERS_TABLE)} - column cluster_status * type in [scaling, normal]: active, non-paused, live * type in [deletecomplete, cleared]: deleted - column is_dev_tier * type is true: dev tier, serverless * type is false: dedicated - column provider_name * type is aws: AWS, aws * type is gcp: GCP, gcp {get_table_desc(EVENTS_TABLE)} - column event_type * type is ClusterNormal: created, cluster created * type is ClusterPaused: paused * type is ClusterResumed: resumed * type is Scaled: scaled # Rules Select statement limit 20 by default Select cluster_id, cluster_name, tenant_id, tenant_name, tenant_email for if possible if query from table `dbaas-eng-bi.user_behavior_analysis.all_clusters` Select event_timestamp, event_type, cluster_id, tenant_name, duration_secs, layout, layout_diff if possible if query from table `dbaas-eng-bi.user_behavior_analysis.prod_view` If not speficied, only select active and paused clusters, excluding deleted clusters Only use table '`dbaas-eng-bi.user_behavior_analysis.prod_view`' when asked for events or activities If question is like 'show layout from cluster' or 'get layout from cluster', use table `dbaas-eng-bi.user_behavior_analysis.all_clusters` # Snippet size of the cluster: tidb.cpu_num * tidb.num_nodes + tikv.cpu_num * tikv.num_nodes + IFNULL(tiflash.cpu_num, 0) * IFNULL(tiflash.num_nodes, 0) for customer XXX, or from XXX: CONTAINS_SUBSTR(tenant_email, 'XXX') events from XXX, or events from customer XXX: CONTAINS_SUBSTR(tenant_email, 'XXX') clusters from XXX, or clusters from customer XXX: CONTAINS_SUBSTR(tenant_email, 'XXX') generate one correct SQL to do query: {question} ''') print(r) sql = r['choices'][0]['text'] return sql app = flask.Flask(__name__) bigquery_client = bigquery.Client() storage_client = storage.Client() @app.route("/") def main(): question = flask.request.args.get("question") if question: question = question.strip() sql = None results = None bq_error = None if question: sql = get_sql(question) query_job = bigquery_client.query(sql) try: results = query_job.result(timeout=30) except Exception as e: bq_error = str(e) # load or append history history = [] blob = storage_client.bucket('dbaas-eng-bi').blob('ai-demo/history.text') if blob.exists(): history = blob.download_as_bytes().decode('utf-8').split('\n##\n') if question: history.append(question) blob.upload_from_string('\n##\n'.join(history)) table = None if results: # Create a PrettyTable instance to hold the results table = PrettyTable() # Add columns to the table for field in results.schema: table.add_column(field.name, []) # Add rows to the table for row in results: table_row = [] for field in results.schema: table_row.append(getattr(row, field.name)) table.add_row(table_row) print(table) return flask.render_template("index.html", sql=sql, results_table=table, question=question, bq_error=bq_error, history=reversed(history[-30:]), code=open(__file__).read()) if __name__ == "__main__": app.run(host="127.0.0.1", port=8080, debug=True)