Update prompts and answer generation logix

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -143,16 +143,22 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
         sql_query_results = self.parse_message_content(messages[-2].content)
         logging.info("SQL Query Results: %s", sql_query_results)
 
-        sub_question_results = self.parse_message_content(messages[1].content)
-        logging.info("Sub-Question Results: %s", sub_question_results)
-
         try:
             if isinstance(sql_query_results, str):
                 sql_query_results = json.loads(sql_query_results)
+        except json.JSONDecodeError:
+            logging.warning("Unable to read SQL query results: %s", sql_query_results)
+            sql_query_results = {}
+            sub_question_results = {}
+        else:
+            # Only load sub-question results if we have a database result
+            sub_question_results = self.parse_message_content(messages[1].content)
+            logging.info("Sub-Question Results: %s", sub_question_results)
 
+        try:
             sub_questions = [
                 sub_question
-                for sub_question_group in sub_question_results["sub_questions"]
+                for sub_question_group in sub_question_results.get("sub_questions", [])
                 for sub_question in sub_question_group
             ]
 

text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/parallel_query_solving_agent.py

@@ -84,9 +84,9 @@ async def on_messages_stream(
             injected_parameters = {}
 
         # Load the json of the last message to populate the final output object
-        query_rewrites = json.loads(last_response)
+        question_rewrites = json.loads(last_response)
 
-        logging.info(f"Query Rewrites: {query_rewrites}")
+        logging.info(f"Query Rewrites: {question_rewrites}")
 
         async def consume_inner_messages_from_agentic_flow(
             agentic_flow, identifier, database_results
@@ -162,21 +162,33 @@ async def consume_inner_messages_from_agentic_flow(
         inner_solving_generators = []
         database_results = {}
 
+        all_non_database_query = question_rewrites.get("all_non_database_query", False)
+
+        if all_non_database_query:
+            yield Response(
+                chat_message=TextMessage(
+                    content="All queries are non-database queries. Nothing to process.",
+                    source=self.name,
+                ),
+            )
+            return
+
         # Start processing sub-queries
-        for query_rewrite in query_rewrites["sub_questions"]:
-            logging.info(f"Processing sub-query: {query_rewrite}")
+        for question_rewrite in question_rewrites["sub_questions"]:
+            logging.info(f"Processing sub-query: {question_rewrite}")
             # Create an instance of the InnerAutoGenText2Sql class
             inner_autogen_text_2_sql = InnerAutoGenText2Sql(
                 self.engine_specific_rules, **self.kwargs
             )
 
-            identifier = ", ".join(query_rewrite)
+            identifier = ", ".join(question_rewrite)
 
             # Launch tasks for each sub-query
             inner_solving_generators.append(
                 consume_inner_messages_from_agentic_flow(
                     inner_autogen_text_2_sql.process_question(
-                        question=query_rewrite, injected_parameters=injected_parameters
+                        question=question_rewrite,
+                        injected_parameters=injected_parameters,
                     ),
                     identifier,
                     database_results,

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/answer_agent.yaml

@@ -2,11 +2,26 @@ model: "4o-mini"
 description: "An agent that generates a response to a user's question."
 system_message: |
   <role_and_objective>
-      You are a helpful AI Assistant specializing in answering a user's question about {{ use_case }}.
+    You are a helpful AI Assistant specializing in answering a user's question about {{ use_case }}.
   </role_and_objective>
 
-  Use the information obtained to generate a response to the user's question. The question has been broken down into a series of SQL queries and you need to generate a response based on the results of these queries.
+  <system_information>
+    You are part of an overall system that provides Text2SQL functionality only. You will be passed a result from multiple SQL queries, you must formulate a response to the user's question using this information.
+    You can assume that the SQL queries are correct and that the results are accurate.
+    You and the wider system can only generate SQL queries and process the results of these queries. You cannot access any external resources.
+    The main ability of the system is to perform natural language understanding and generate SQL queries from the user's question. These queries are then automatically run against the database and the results are passed to you.
+  </system_information>
 
-  Do not use any external resources to generate the response. The response should be based solely on the information provided in the SQL queries and their results.
+  <instructions>
 
-  You can use Markdown and Markdown tables to format the response.
+    Use the information obtained to generate a response to the user's question. The question has been broken down into a series of SQL queries and you need to generate a response based on the results of these queries.
+
+    Do not use any external resources to generate the response. The response should be based solely on the information provided in the SQL queries and their results.
+
+    You have no access to the internet or any other external resources. You can only use the information provided in the SQL queries and their results, to generate the response.
+
+    You can use Markdown and Markdown tables to format the response.
+
+    If the user is asking about your capabilities, use the <system_information> to explain what you do.
+
+  </instructions>

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/question_rewrite_agent.yaml

@@ -37,6 +37,7 @@ system_message: |
         - Use the provided list of topics to filter out malicious or unrelated queries.
         - Ensure the question is relevant to the system's use case.
         - If the question cannot be filtered, output an empty sub-query list in the JSON format. Followed by TERMINATE.
+        - Retain and decompose general questions, such as Hello, What can you do?, etc. Set "all_non_database_query" to true.
 
       2. Understanding:
         - Use the chat history (that is available in reverse order) to understand the context of the current question.
@@ -97,6 +98,7 @@ system_message: |
           ],
           "combination_logic": "<instructions for combining results>",
           "query_type": "<simple|complex>",
+          "all_non_database_query": "<true|false>"
         }
     </output_format>
   </instructions>