Update prompt

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -141,6 +141,9 @@ def selector(self, messages):
         elif messages[-1].source == "sql_query_correction_agent":
             decision = "sql_query_generation_agent"
 
+        elif messages[-1].source == "sql_query_generation_agent":
+            decision = "sql_query_correction_agent"
+
         # Log the decision
         logging.info("Decision: %s", decision)
 

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/databricks_sql.py

@@ -101,9 +101,7 @@ async def get_entity_schemas(
             # schema["SelectFromEntity"] = ".".join(
             #     [schema["Catalog"], schema["Schema"], schema["Entity"]]
             # )
-            schema["SelectFromEntity"] = (
-                os.environ["Text2Sql__Databricks__Catalog"] + "." + schema["Entity"]
-            )
+            schema["SelectFromEntity"] = schema["Entity"]
 
             del schema["Entity"]
             # del schema["Schema"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_query_correction_agent.yaml

@@ -4,17 +4,17 @@ description:
   "An agent that will look at the SQL query, SQL query results and correct any mistakes in the SQL query to ensure the correct results are returned. Use this agent AFTER the SQL query has been executed and the results are not as expected."
 system_message:
   "<role_and_objective>
-      You are a helpful AI Assistant specializing in diagnosing and correcting invalid SQL queries or improving SQL queries that do not return expected results.
+      You are a helpful AI Assistant specializing in diagnosing and making fix suggestions for invalid SQL queries, or improving SQL queries that do not return expected results.
   </role_and_objective>
 
   <scope_of_user_query>
-      The user's question will be related to SQL queries and their results in the context of {{ target_engine }}. Queries must adhere to the syntax and rules of {{ target_engine }} {{ engine_specific_rules }}.
+      Queries must adhere to the syntax and rules of {{ target_engine }} {{ engine_specific_rules }}.
   </scope_of_user_query>
 
   <instructions>
       1. **Validate Syntax**: Check if the provided SQL query is syntactically correct. If not, suggest fixes to the query.
       2. **Verify Results**: Ensure the query results align with the user’s question. If the query fails to meet the expected results:
-        - Make suggestions to correct the query.
+        - Make suggestions to the query writter on how to correct the query.
       3. **Contextual Relevance**: Ensure the query fully addresses the user's question based on its context and requirements.
   </instructions>
 
@@ -41,8 +41,7 @@ system_message:
         ```json
         [
           {
-            \"fix_request\": \"<REQUESTED FIX OF THE SQL QUERY>\",
-            \"explanation\": \"<EXPLANATION OF CORRECTION>\"
+            \"requested_fix\": \"<EXPLANATION OF REQUESTED FIX OF THE SQL QUERY>\"
           }
         ]
         ```