Prompt generator, updated demo

John Doe · John Doe · commit 085cc5519229 · 2023-02-04T12:31:07.000+01:00
API Token now can be set in localStorage
diff --git a/README.md b/README.md
@@ -3,20 +3,21 @@
 This is simple python application to generate SQL Schema + prompt to ask GPT-3 to generate SQL queries.
 It also has a simple UI to show results in a table.
 
+**It can generate a prompt for you, if you don't know what to ask GPT-3.**
+
 ![](https://github.com/Hormold/gpt-sql-box/blob/master/docs/demo.gif?raw=true)
 
 
 ## How it works:
-1. Getting SQL schemas from PostgreSQL
-2. Compile prompt from SQL Schema
+1. Getting SQL schemas from PostgreSQL and compile prompt from SQL Schema
 3. Wait for user input
 4. Generate SQL query from prompt + user input
 5. Show SQL query and ask user to confirm or edit if it is correct before executing it
 6. Execute SQL query and show results in a table
 
 ## Environment
 - DATABASE_URL: PostgreSQL database URL
-- OPENAI_TOKEN: OpenAI API token
+- OPENAI_TOKEN: OpenAI API token (Not nessessary, you can set it in the UI)
 - APP_PORT: Port to run the application (default: 5000)
 
 ## How to run
@@ -26,4 +27,4 @@ It also has a simple UI to show results in a table.
 4. Run the application using `python app.py` and open `http://localhost:5000` in your browser
 
 ## Contact 
-If you have any questions, please contact me at [t.me/define](https://t.me/define)
+If you have any questions, please contact me at [@define](https://t.me/define) in Telegram.
diff --git a/app.py b/app.py
@@ -18,10 +18,12 @@
 if not DATABASE_URL:
     print('Please set DATABASE_URL in .env file.')
     sys.exit(1)
-openai.api_key = os.getenv('OPENAI_TOKEN')
-if not os.getenv('OPENAI_TOKEN'):
-    print('Please set OPENAI_TOKEN in .env file.')
-    sys.exit(1)
+
+if os.getenv('OPENAI_TOKEN'):
+    openai.api_key = os.getenv('OPENAI_TOKEN')
+
+if not openai.api_key:
+    print('Please set OPENAI_TOKEN in .env file or set token in UI') # Not a critical error
 
 # Generate SQL Schema from PostgreSQL
 schema = Schema()
@@ -31,10 +33,10 @@
 @app.get('/')
 def index():
     """Show SQL Schema + prompt to ask GPT-3 to generate SQL queries"""
-    # Get JSON data (not escaped)
-    normalized_json_data = json.dumps(json_data);
+    normalized_json_data = json.dumps(json_data)
     return render_template(
         'index.html',
+        has_openai_key=bool(openai.api_key),
         sql_schema=sql_schema,
         json_data=normalized_json_data
     )
@@ -44,14 +46,25 @@ def generate():
     """Generate SQL query from prompt + user input"""
     try:
         content = request.json
-        print('Content:', content)
         user_input = content['query']
         query_temperture = content['temp']
         selected = content['selected']
         print('Selected tables:', selected)
         print('User input:', user_input)
         print('Query temperture:', query_temperture)
 
+        if not content['api_key'] and not openai.api_key:
+            return {
+                'success': False,
+                'error': 'Please set OPENAI_TOKEN in .env file or set token in UI'
+            }
+
+        if content['api_key'] and not openai.api_key:
+            openai.api_key = content['api_key'] # Inject API key from UI
+            print('API key was set from UI')
+        else:
+            print('API key was set from .env file')
+
         # Update prompt
         regen_schema = schema.regen(selected)
         new_prompt = f'Given an input question, respond with syntactically correct PostgreSQL. Be creative but the SQL must be correct, not nessesary to use all tables. {regen_schema}'
@@ -90,7 +103,7 @@ def generate():
         print(err)
         return {
             'success': False,
-            'sql_query': err
+            'error': str(err)
         }
 
 @app.post('/run')
@@ -123,13 +136,65 @@ def execute():
         print(err)
         return {
             'success': False,
-            'error': err
+            'error': str(err)
+        }
+    except Exception as err:
+        print(err)
+        return {
+            'success': False,
+            'error': str(err)
+        }
+
+@app.post('/generate_prompt')
+def generate_prompt():
+    """Generate prompt from selected tables"""
+    try:
+        content = request.json
+        selected = content['selected']
+        query_temperture = content['temp']
+
+        if not content['api_key'] and not openai.api_key:
+            return {
+                'success': False,
+                'error': 'Please set OPENAI_TOKEN in .env file or set token in UI'
+            }
+
+        if content['api_key'] and not openai.api_key:
+            openai.api_key = content['api_key']
+
+        # Update prompt
+        regen_schema = schema.regen(selected)
+        final_prompt = f'Your task is create creative prompt, using this scheme of database: {regen_schema}\n\nDo not generate SQL query, generate text based prompt:\n\n'
+
+        gpt_response = openai.Completion.create(
+            engine="text-davinci-003",
+            prompt=final_prompt,
+            temperature=float(query_temperture),
+            max_tokens=500,
+            top_p=1,
+            frequency_penalty=0,
+            presence_penalty=0,
+            stop=["\n\n"]
+        )
+
+        print(gpt_response)
+
+        used_tokens = gpt_response['usage']['total_tokens']
+
+        # Get SQL query
+        query = gpt_response['choices'][0]['text'].lstrip().rstrip()
+        print('Generated prompt:', query)
+
+        return {
+            'success': True,
+            'query': query,
+            'used_tokens': used_tokens,
         }
     except Exception as err:
         print(err)
         return {
             'success': False,
-            'error': err
+            'error': str(err)
         }
 
 
diff --git a/docs/demo.gif b/docs/demo.gif
diff --git a/schema.py b/schema.py
@@ -11,6 +11,7 @@
 
 class Schema:
     """Generate SQL Schema from PostgreSQL"""
+    
     def __init__(self, schema = 'public'):
         """Connect to PostgreSQL database"""
         self.schema = schema
@@ -25,6 +26,7 @@ def __init__(self, schema = 'public'):
         self.comments = []
         self.tables = []
         self.columns = []
+
     def get_tables(self):
         """Get list of tables"""
         self.cur.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = %s", (self.schema,))
diff --git a/tpl/index.html b/tpl/index.html
@@ -42,6 +42,31 @@
 	<div id="app">
 		<div class="container">
 			<div class="row">
+
+				<div class="col-md-8 col-md-offset-2" v-if="view === 'setkey'">
+					<div class="panel panel-default">
+						<div class="panel-heading">
+							<h3 class="panel-title">Set OpenAI Token</h3>
+						</div>
+						<div class="panel-body">
+							<p>Looks like API token is not set in environment variable. Please enter your OpenAI API token below and it will be saved to browser localStorage</p>
+							<fieldset>
+								<div class="form-group">
+									<input class="form-control" v-model="key" placeholder="Enter OpenAI API token"
+										name="key" type="text">
+								</div>
+								<div>
+									<input class="btn btn-lg btn-success btn-block" type="submit" value="Save to localStorage and continue"
+										@click="setKey">
+								</div>
+							</fieldset>
+						</div>
+						<div class="panel-footer">
+							<span class="text-muted">You can get your OpenAI API token <a href="https://beta.openai.com/account/api-keys" target="_blank">here</a>.</span>
+						</div>
+					</div>
+				</div>
+
 				<div class="col-md-8 col-md-offset-2" v-if="view === 'index'">
 					<div class="panel panel-default">
 						<div class="panel-heading">
@@ -61,8 +86,15 @@ <h3 class="panel-title">GPT SQL Query: Configuation</h3>
 										v-model="temp" max="1" step="0.1" value="0.5">
 								</div>
 								<div v-if="!isLoading">
-									<input class="btn btn-lg btn-success btn-block" type="submit" value="🧙‍♂️ Do Magic!"
-										@click="generate">
+									<div class="row">
+										<div class="col-md-6">
+											<input class="btn btn-lg btn-success btn-block" type="submit" value="🧙‍♂️ Do Magic!" @click="generate">
+										</div>
+										<div class="col-md-6">
+											<input class="btn btn-lg btn-success btn-block" type="submit" value="🤷‍♂️ I don't know what to ask" @click="generateRandom">
+										</div>
+									</div>
+									
 								</div>
 								<div v-else>
 									<button class="btn btn-lg btn-success btn-block" type="button" disabled>
@@ -122,13 +154,21 @@ <h3 class="panel-title">Approve generate SQL query<br /><br /></h3>
 						</div>
 
 						<div class="panel-body">
+							<div class="alert alert-warning" role="alert" v-if="sql_error">
+								<b>Error in generated query:</b> [[sql_error]]
+							</div>
 							<fieldset>
 								<div class="form-group">
-									<textarea class="form-control" name="query" type="text" rows="15" cols="90"
-										v-model="sql_query"></textarea>
+									<textarea
+										class="form-control"
+										name="query"
+										type="text" 
+										cols="80"
+										rows="12"
+										v-model="sql_query">
+									</textarea>
 								</div>
 								<div v-if="!isLoading">
-									<!-- row of two buttons - refresh & run -->
 									<div class="row">
 										<div class="col-md-6">
 											<input class="btn btn-lg btn-success btn-block" type="submit" :value="'🧙‍♂️ Regenerate (~' + (used_tokens) + ' tokens)'"
@@ -181,9 +221,6 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 								</tbody>
 							</table>
 						</div>
-
-						<!-- Execution time -->
-
 					</div>
 					<div class="panel-footer">
 						<div class="pull-right">
@@ -210,11 +247,12 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 	</div>
 
 	<!-- Vue -->
-	<script>
+	<script lang="javascript">
 		var app = new Vue({
 			el: '#app',
 			delimiters: ['[[', ']]'],
 			data: {
+				hasOpenAIKey: '{{has_openai_key}}',
 				sql_schema: '{{sql_schema}}',
 				in_schema: JSON.parse('{{json_data|safe}}'),
 				sql_query: '',
@@ -226,10 +264,16 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 				isLoading: false,
 				seconds_elapsed: 0,
 				used_tokens: 0,
+				key: '',
+				sql_error: null
 			},
 
 			created() {
-				this.schema = this.in_schema;
+				if(localStorage.getItem('openai_key'))
+					this.key = localStorage.getItem('openai_key')
+				if(this.hasOpenAIKey === 'False' && !this.key)
+					this.view = 'setkey'
+				this.schema = this.in_schema
 				for (var key in this.schema) {
 					if (localStorage.getItem(key) === 'true') {
 						this.schema[key].selected = true;
@@ -238,12 +282,19 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 					}
 				}
 
-				//this.view = 'generate';
-				//this.sql_query = "SELECT * FROM users WHERE username LIKE 'A%' ORDER BY created_at DESC LIMIT 10;"
+				this.view = 'generate';
+				this.sql_query = "SELECT * FROM usfers WHERE fusername LIKE 'A%' ORDER BY created_at DESC LIMIT 10;"
 			},
 
 			methods: {
 
+				setKey: function () {
+					if(!this.key.match(/^sk-([a-zA-Z0-9]{1,64})$/))
+						return alert('Invalid key, it should be in the format sk-[a-zA-Z0-9]{32}');
+					this.view = 'index';
+					localStorage.setItem('openai_key', this.key);
+				},
+
 				updateSchema: function (key) {
 					// save to localStorage
 					localStorage.setItem(key, this.schema[key].selected);
@@ -267,10 +318,37 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 					this.$forceUpdate();
 				},
 
+				generateRandom: async function () {
+					if(this.isLoading)
+						return;
+					this.isLoading = true;
+
+					const response = await fetch('/generate_prompt', {
+						method: 'POST',
+						headers: {
+							'Content-Type': 'application/json'
+						},
+						body: JSON.stringify({
+							temp: this.temp,
+							api_key: this.key,
+							selected: Object.keys(this.schema).filter(key => this.schema[key]
+								.selected)
+						})
+					});
+
+					const data = await response.json();
+					this.isLoading = false;
+					if(!data.success)
+						return alert(data.error);
+					this.query = data.query;
+				},
+
 				generate: async function () {
 					if(this.query === '' || this.isLoading)
 						return;
 					this.isLoading = true;
+					this.sql_error = null;
+
 					const response = await fetch('/generate', {
 						method: 'POST',
 						headers: {
@@ -279,10 +357,12 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 						body: JSON.stringify({
 							query: this.query,
 							temp: this.temp,
+							api_key: this.key,
 							selected: Object.keys(this.schema).filter(key => this.schema[key]
 								.selected)
 						})
 					});
+
 					const data = await response.json();
 					this.isLoading = false;
 					if(data.success) {
@@ -297,6 +377,7 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 				runQuery: async function () {
 					// fetch
 					this.isLoading = true;
+					this.sql_error = null;
 					const response = await fetch('/run', {
 						method: 'POST',
 						headers: {
@@ -314,7 +395,7 @@ <h3 class="panel-title">Results of SQL query<br /><br /></h3>
 						this.seconds_elapsed = data.seconds_elapsed;
 						this.view = 'run';
 					} else {
-						alert(data.error);
+						this.sql_error = data.error;
 					}
 				},
 
