Update openai_integration.py

openai_integration.py

@@ -48,18 +48,18 @@ def reformulate_query(nl_query, schema_description):
         examples = """
         Example 1:
         Input: "I'm trying to figure out which agents are the busiest. Can you show me like the top few agents who have a lot on their plate?"
-        Reformulated: "Show me the top 5 agents with the highest workload, including their names, emails, and specializations."
+        Reformulated: "Show me the top 5 agents with the highest workload, including their names, emails, and specializations. Use Tables: agents"
 
         Example 2:
         Input: "I need to get an overview of our classes. Could you pull up a list that shows what each class is about and which client it's for?"
-        Reformulated: "List all classes with their subjects, client names, and the number of learners in each class, sorted by the number of learners in descending order."
+        Reformulated: "List all classes with their subjects, client names, and the number of learners in each class, sorted by the number of learners in descending order. Use Tables: clients, learners, classes"
         """
 
         prompt = (
             f"Database Schema:\n{schema_description}\n\n"
-            f"Using the database schema, reformulate the following natural language query to be more precise and clear for generating an SQL query:\n{nl_query}\n\n"
+            f"Using the database schema, reformulate the following natural language query to be more precise and in line with the database schema:\n{nl_query}\n\n"
             f"Here are some examples of how to reformulate the natural language query:\n{examples}\n\n"
-            f"Output only the natural language query string using the provided schema, don't prefix with anything and don't use markdown"
+            f"Output the query string, list the DB tables that will be used seperated by commas, don't prefix with anything and don't use markdown or quotation marks."
         )
 
         logging.info("Sending reformulation request to OpenAI...")
@@ -133,14 +133,14 @@ def generate_sql_from_reformulated(reformulated_query, schema_description):
             f"Examples:\n{examples}\n\n"
             f"Convert the natural language query into an SQL query that matches the schema.\n\n"
             f"Use the Examples as a guide on what I expect you to do.\n\n"
-            f"Only return the RAW SQL that can be directly executed, don't prefix it with sql and don't end the sql with Quotation marks etc"
+            f"Only return the RAW SQL that can be directly executed, don't prefix it with sql, don't use quotation marks and don't use markdown."
         )
 
         logging.info("Sending SQL generation request to OpenAI...")
         response = client.chat.completions.create(
             model=GPT_MODEL,
             messages=[
-                {"role": "system", "content": "You are an assistant that converts natural language to SQL using a databse shema."},
+                {"role": "system", "content": "You are an assistant that converts natural language to SQL using a provided databse shema."},
                 {"role": "user", "content": prompt}
             ],
             max_tokens=300,