Create app.py

app.py

@@ -0,0 +1,123 @@
+import pandas as pd
+from openai import OpenAI
+import os
+from google.colab import userdata
+from google.cloud import bigquery
+import numpy as np
+import gradio as gr
+
+from sql_queries import sql_queries
+
+os.environ['OPENAI_API_KEY'] = userdata.get('open_ai_fieldops')
+os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/content/aakash-project-422813-1230ad3ba9f1.json'
+
+project_id = 'aakash-project-422813'
+dataset_id = 'ved_test'
+table_id = 'synthetic_data'
+
+openai_client = OpenAI()
+
+def fetch_table_schema(project_id, dataset_id, table_id):
+    bqclient = bigquery.Client(project=project_id)
+
+    table_ref = f"{project_id}.{dataset_id}.{table_id}"
+
+    table = bqclient.get_table(table_ref)
+
+    schema_dict = {}
+    for schema_field in table.schema:
+        schema_dict[schema_field.name] = schema_field.field_type
+
+    return schema_dict
+
+def get_sql_query(description):
+    prompt = f'''
+    Generate the SQL query for the following task:\n{description}.\n
+    The database you need is called {dataset_id} and the table is called {table_id}.
+    Use the format {dataset_id}.{table_id} as the table name in the queries.
+    Enclose column names in backticks(`) not quotation marks.
+    Do not assign aliases to the columns.
+    Do not calculate new columns, unless specifically called to.
+    Return only the SQL query, nothing else.
+    Do not use WITHIN GROUP clause.
+    \nThe list of all the columns is as follows: {schema} /n
+    '''
+    try:
+      completion = openai_client.chat.completions.create(
+          model='gpt-4o',
+          messages = [
+                  {"role": "system", "content": "You are an expert Data Scientist with in-depth knowledge of SQL, working on Network Telemetry Data."},
+                  {"role": "user", "content": f'{prompt}'},
+                ]
+      )
+
+    except Exception as e:
+     print(f'The following error ocurred: {e}\n')
+     pass
+
+    sql_query = completion.choices[0].message.content.strip().split('```sql')[1].split('```')[0]
+    return sql_query
+
+schema = fetch_table_schema(project_id, dataset_id, table_id)
+
+def execute_sql_query(query):
+    client = bigquery.Client()
+
+    try:
+     result = client.query(query).to_dataframe()
+     message = f'The query : {query}\n was successfully executed and returned the above result.\n'
+
+    except Exception as e:
+     result = 'No output returned'
+     message = f'The query : {query}\n could not be executed due to exception {e}\n'
+
+    return result, message
+
+def echo(text):
+  query = get_sql_query(text)
+  result, message = execute_sql_query(query)
+  return result, message
+
+def gradio_interface(text):
+    result, message = echo(text)
+    if isinstance(result, pd.DataFrame):
+        return gr.Dataframe(value=result), message
+    else:
+        return result, message
+
+demo = gr.Blocks(
+        title="Text-to-SQL",
+        theme=gr.themes.Monochrome(),
+)
+
+with demo:
+
+  gr.Markdown(
+    '''
+    # <p style="text-align: center;">Text to SQL Query Engine</p>
+
+    <p style="text-align: center;">
+    Welcome to our Text2SQL Engine.
+    <br>
+    Enter your query in natural language and we'll convert it to SQL and return the result to you.
+    </p>
+    '''
+    )
+
+  with gr.Row():
+    with gr.Column():
+      text_input = gr.Textbox(label="Enter your query")
+    with gr.Column():
+      output_text = gr.Textbox(label="Output", interactive=False)
+      output_df = gr.Dataframe(interactive=False)
+
+    def update_output(text):
+        result, message = gradio_interface(text)
+        if isinstance(result, pd.DataFrame):
+            return gr.update(visible=True), result, message
+        else:
+            return gr.update(visible=False), result, message
+
+    text_input.submit(update_output, inputs=text_input, outputs=[output_df, output_text])
+
+demo.launch(debug=True, auth=("admin", "Text2SQL"))