Chooses best query from chatgpt

app.py

@@ -131,6 +131,11 @@ def extract_db_code(text):
     matches = re.findall(pattern, text, re.DOTALL)
     return [match.strip() for match in matches]
 
+def extract_from_code_block(text):
+    pattern = r'```(?:\w+)?\s?(.*?)```'
+    match = re.search(pattern, text, re.DOTALL)
+    return match.group(1).strip() if match else ''
+
 def generate_dummy_db(db_info, question):
     pre_prompt = """
     Generate a SQLite database with dummy data for this database from the DB Layout. Your task is to generate just a database, no queries. For each input do the following:
@@ -188,6 +193,36 @@ def test_query_on_dummy_db(db_code, query):
         print(f"Query: {query}\tError encountered: {e}")
         return False
 
+def choose_best_query(queries, question):
+    pre_prompt = """
+    Given a list of queries. Your task is to choose just a single query which satisfies the question the most with the least amount of filters, groupings, and conditions. For each input do the following:
+        1. Breakdown the list of queries into small pieces and explain what each query is doing.
+        2. Explain why each query is relevant to the question.
+        3. Choose the most relevant query from your explanation that aligns to the question best with the least amount of unnecessary filters or conditions. Output the best query in a single code block ``````.
+    """
+    prompt = pre_prompt + "\n\nQuestion: " + question + "\n\nQueries:" + "\n\n".join(queries)
+
+    while True:
+        try:
+            response = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {"role": "user", "content": prompt}
+                ],
+                #temperature=0.7,
+            )
+            response_text = response['choices'][0]['message']['content']
+            print(response_text)
+            
+            query = extract_from_code_block(response_text)
+
+            return query
+            
+        except Exception as e:
+            print(f'Error occurred: {str(e)}')
+            print('Waiting for 10 seconds before retrying...')
+            time.sleep(10)
+
 
 def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, log=False, num_return_sequences=1, num_beams=1, do_sample=False):
     if num_return_sequences > num_beams:
@@ -246,15 +281,18 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
             query = query.replace("\n", " ").replace("\t", " ").strip()
             # Test against dummy database
             success = test_query_on_dummy_db(db_code, query)
-            # Format again
-            query = format(query) if format_sql else query
+            
             if success:
                 responses.append(query)
         else:
             responses.append(query)
             
-    # Choose a random response from responses
-    output = responses[0] if len(responses) > 0 else "###"
+    # Choose the best query if num_return_sequences > 1
+    if num_return_sequences > 1:
+        query = choose_best_query(responses, input_message)
+        # Format again
+        query = format(query) if format_sql else query
+        responses = [query]
 
     if log:
         # Log the request to Firestore