vedsadani commited on
Commit
1faa9b0
1 Parent(s): c4081c4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +123 -0
app.py ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import pandas as pd
2
+ from openai import OpenAI
3
+ import os
4
+ from google.colab import userdata
5
+ from google.cloud import bigquery
6
+ import numpy as np
7
+ import gradio as gr
8
+
9
+ from sql_queries import sql_queries
10
+
11
+ os.environ['OPENAI_API_KEY'] = userdata.get('open_ai_fieldops')
12
+ os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/content/aakash-project-422813-1230ad3ba9f1.json'
13
+
14
+ project_id = 'aakash-project-422813'
15
+ dataset_id = 'ved_test'
16
+ table_id = 'synthetic_data'
17
+
18
+ openai_client = OpenAI()
19
+
20
+ def fetch_table_schema(project_id, dataset_id, table_id):
21
+ bqclient = bigquery.Client(project=project_id)
22
+
23
+ table_ref = f"{project_id}.{dataset_id}.{table_id}"
24
+
25
+ table = bqclient.get_table(table_ref)
26
+
27
+ schema_dict = {}
28
+ for schema_field in table.schema:
29
+ schema_dict[schema_field.name] = schema_field.field_type
30
+
31
+ return schema_dict
32
+
33
+ def get_sql_query(description):
34
+ prompt = f'''
35
+ Generate the SQL query for the following task:\n{description}.\n
36
+ The database you need is called {dataset_id} and the table is called {table_id}.
37
+ Use the format {dataset_id}.{table_id} as the table name in the queries.
38
+ Enclose column names in backticks(`) not quotation marks.
39
+ Do not assign aliases to the columns.
40
+ Do not calculate new columns, unless specifically called to.
41
+ Return only the SQL query, nothing else.
42
+ Do not use WITHIN GROUP clause.
43
+ \nThe list of all the columns is as follows: {schema} /n
44
+ '''
45
+ try:
46
+ completion = openai_client.chat.completions.create(
47
+ model='gpt-4o',
48
+ messages = [
49
+ {"role": "system", "content": "You are an expert Data Scientist with in-depth knowledge of SQL, working on Network Telemetry Data."},
50
+ {"role": "user", "content": f'{prompt}'},
51
+ ]
52
+ )
53
+
54
+ except Exception as e:
55
+ print(f'The following error ocurred: {e}\n')
56
+ pass
57
+
58
+ sql_query = completion.choices[0].message.content.strip().split('```sql')[1].split('```')[0]
59
+ return sql_query
60
+
61
+ schema = fetch_table_schema(project_id, dataset_id, table_id)
62
+
63
+ def execute_sql_query(query):
64
+ client = bigquery.Client()
65
+
66
+ try:
67
+ result = client.query(query).to_dataframe()
68
+ message = f'The query : {query}\n was successfully executed and returned the above result.\n'
69
+
70
+ except Exception as e:
71
+ result = 'No output returned'
72
+ message = f'The query : {query}\n could not be executed due to exception {e}\n'
73
+
74
+ return result, message
75
+
76
+ def echo(text):
77
+ query = get_sql_query(text)
78
+ result, message = execute_sql_query(query)
79
+ return result, message
80
+
81
+ def gradio_interface(text):
82
+ result, message = echo(text)
83
+ if isinstance(result, pd.DataFrame):
84
+ return gr.Dataframe(value=result), message
85
+ else:
86
+ return result, message
87
+
88
+ demo = gr.Blocks(
89
+ title="Text-to-SQL",
90
+ theme=gr.themes.Monochrome(),
91
+ )
92
+
93
+ with demo:
94
+
95
+ gr.Markdown(
96
+ '''
97
+ # <p style="text-align: center;">Text to SQL Query Engine</p>
98
+
99
+ <p style="text-align: center;">
100
+ Welcome to our Text2SQL Engine.
101
+ <br>
102
+ Enter your query in natural language and we'll convert it to SQL and return the result to you.
103
+ </p>
104
+ '''
105
+ )
106
+
107
+ with gr.Row():
108
+ with gr.Column():
109
+ text_input = gr.Textbox(label="Enter your query")
110
+ with gr.Column():
111
+ output_text = gr.Textbox(label="Output", interactive=False)
112
+ output_df = gr.Dataframe(interactive=False)
113
+
114
+ def update_output(text):
115
+ result, message = gradio_interface(text)
116
+ if isinstance(result, pd.DataFrame):
117
+ return gr.update(visible=True), result, message
118
+ else:
119
+ return gr.update(visible=False), result, message
120
+
121
+ text_input.submit(update_output, inputs=text_input, outputs=[output_df, output_text])
122
+
123
+ demo.launch(debug=True, auth=("admin", "Text2SQL"))