Update prompts

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -73,7 +73,10 @@ def termination_condition(self):
         termination = (
             TextMentionTermination("TERMINATE")
             | SourceMatchTermination("answer_agent")
-            | TextMentionTermination("contains_disambiguation_requests")
+            | TextMentionTermination(
+                "contains_disambiguation_requests",
+                sources=["parallel_query_solving_agent"],
+            )
             | MaxMessageTermination(5)
         )
         return termination
@@ -142,32 +145,26 @@ def parse_message_content(self, content):
         # If all parsing attempts fail, return the content as-is
         return content
 
-    def extract_decomposed_user_messages(self, messages: list) -> list[list[str]]:
-        """Extract the decomposed messages from the answer."""
+    def extract_steps(self, messages: list) -> list[list[str]]:
+        """Extract the steps messages from the answer."""
         # Only load sub-message results if we have a database result
         sub_message_results = self.parse_message_content(messages[1].content)
-        logging.info("Decomposed Results: %s", sub_message_results)
+        logging.info("Steps Results: %s", sub_message_results)
 
-        decomposed_user_messages = sub_message_results.get(
-            "decomposed_user_messages", []
-        )
+        steps = sub_message_results.get("steps", [])
 
-        logging.debug(
-            "Returning decomposed_user_messages: %s", decomposed_user_messages
-        )
+        logging.debug("Returning steps: %s", steps)
 
-        return decomposed_user_messages
+        return steps
 
     def extract_disambiguation_request(
         self, messages: list
     ) -> DismabiguationRequestsPayload:
         """Extract the disambiguation request from the answer."""
         all_disambiguation_requests = self.parse_message_content(messages[-1].content)
 
-        decomposed_user_messages = self.extract_decomposed_user_messages(messages)
-        request_payload = DismabiguationRequestsPayload(
-            decomposed_user_messages=decomposed_user_messages
-        )
+        steps = self.extract_steps(messages)
+        request_payload = DismabiguationRequestsPayload(steps=steps)
 
         for per_question_disambiguation_request in all_disambiguation_requests[
             "disambiguation_requests"
@@ -198,12 +195,10 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
             sql_query_results = {}
 
         try:
-            decomposed_user_messages = self.extract_decomposed_user_messages(messages)
+            steps = self.extract_steps(messages)
 
             logging.info("SQL Query Results: %s", sql_query_results)
-            payload = AnswerWithSourcesPayload(
-                answer=answer, decomposed_user_messages=decomposed_user_messages
-            )
+            payload = AnswerWithSourcesPayload(answer=answer, steps=steps)
 
             if not isinstance(sql_query_results, dict):
                 logging.error(f"Expected dict, got {type(sql_query_results)}")

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

@@ -96,9 +96,9 @@ async def on_messages_stream(
             injected_parameters = {}
 
         # Load the json of the last message to populate the final output object
-        sequential_rounds = json.loads(last_response)
+        sequential_steps = json.loads(last_response)
 
-        logging.info(f"Query Rewrites: {sequential_rounds}")
+        logging.info(f"Query Rewrites: {sequential_steps}")
 
         async def consume_inner_messages_from_agentic_flow(
             agentic_flow, identifier, filtered_parallel_messages
@@ -195,12 +195,12 @@ async def consume_inner_messages_from_agentic_flow(
         inner_solving_generators = []
         filtered_parallel_messages = FilteredParallelMessagesCollection()
 
-        # Convert all_non_database_query to lowercase string and compare
-        all_non_database_query = str(
-            sequential_rounds.get("all_non_database_query", "false")
+        # Convert requires_sql_queries to lowercase string and compare
+        requires_sql_queries = str(
+            sequential_steps.get("requires_sql_queries", "false")
         ).lower()
 
-        if all_non_database_query == "true":
+        if requires_sql_queries == "false":
             yield Response(
                 chat_message=TextMessage(
                     content="All queries are non-database queries. Nothing to process.",
@@ -210,7 +210,7 @@ async def consume_inner_messages_from_agentic_flow(
             return
 
         # Start processing sub-queries
-        for sequential_round in sequential_rounds["decomposed_user_messages"]:
+        for sequential_round in sequential_steps["steps"]:
             logging.info(f"Processing round: {sequential_round}")
 
             for parallel_message in sequential_round:

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

@@ -40,7 +40,7 @@ async def on_messages(
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentEvent | Response, None]:
-        # Get the decomposed messages from the user_message_rewrite_agent
+        # Get the steps messages from the user_message_rewrite_agent
         try:
             request_details = json.loads(messages[0].content)
             injected_parameters = request_details["injected_parameters"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -54,9 +54,7 @@ class DismabiguationRequest(InteractionPayloadBase):
         disambiguation_requests: list[DismabiguationRequest] | None = Field(
             default_factory=list, alias="disambiguationRequests"
         )
-        decomposed_user_messages: list[list[str]] = Field(
-            default_factory=list, alias="decomposedUserMessages"
-        )
+        steps: list[list[str]] = Field(default_factory=list, alias="Steps")
 
     payload_type: Literal[PayloadType.DISAMBIGUATION_REQUESTS] = Field(
         PayloadType.DISAMBIGUATION_REQUESTS, alias="payloadType"
@@ -81,9 +79,7 @@ class Source(InteractionPayloadBase):
             sql_rows: list[dict] = Field(default_factory=list, alias="sqlRows")
 
         answer: str
-        decomposed_user_messages: list[list[str]] = Field(
-            default_factory=list, alias="decomposedUserMessages"
-        )
+        steps: list[list[str]] = Field(default_factory=list, alias="Steps")
         sources: list[Source] = Field(default_factory=list)
 
     payload_type: Literal[PayloadType.ANSWER_WITH_SOURCES] = Field(

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

@@ -2,7 +2,7 @@ model: "4o-mini"
 description: "An agent that generates a response to a user's question."
 system_message: |
   <role_and_objective>
-    You are Senior Data Analystm, specializing in providing data driven answers to a user's question. Use the general business use case of '{{ use_case }}' to aid understanding of the user's question. You should provide a clear and concise response based on the information obtained from the SQL queries and their results. Adopt a data-driven approach to generate the response.
+    You are Senior Data Analyst, specializing in providing data driven answers to a user's question. Use the general business use case of '{{ use_case }}' to aid understanding of the user's question. You should provide a clear and concise response based on the information obtained from the SQL queries and their results. Adopt a data-driven approach to generate the response.
   </role_and_objective>
 
   <system_information>

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/disambiguation_and_sql_query_generation_agent.yaml

@@ -2,8 +2,8 @@ model:
   4o-mini
 description:
   "An agent that specialises in disambiguating the user's question and mapping it to database schemas for {{ use_case }}."
-system_message:
-  "<role_and_objective>
+system_message: |
+  <role_and_objective>
     You are Senior Data Engineer specializing in disambiguating questions, mapping them to the relevant columns and schemas in the database and finally generating SQL queries.
     Use the general business use case of '{{ use_case }}' to aid understanding of the user's question.
     Your job is to create clear mappings between the user's intent and the available database schema.
@@ -152,8 +152,10 @@ system_message:
     </sql_query_generation_rules>
 
     <disambiguation_rules>
-      BEFORE CARRY OUT DISAMBIGUATION, ENSURE THAT YOU HAVE CHECKED ALL AVAILABLE DATABASE SCHEMAS AND FILTERS FOR A MOST PROBABLE MAPPING. YOU WILL NEED TO THINK THROUGH THE SCHEMAS AND CONSIDER SCHEMAS / COLUMNS THAT ARE SPELT DIFFERENTLY, BUT ARE LIKELY TO MEAN THE SAME THING.
-      ALWAYS PRIORITIZE CLEAR MAPPINGS OVER DISAMBIGUATION REQUESTS.
+      **Important**:
+        Before carrying out disambiguation, ensure that you have checked all available database schemas and filters for a most probable mapping. You will need to think through the schemas and consider schemas / columns that are spelt differently, but are likely to mean the same thing.
+
+        You must never ask for information that is already available in the user's message. e.g. if the user asks for the average age of students, and the schema has a column named 'age' in the 'student' table, you should not ask the user to clarify the column name. Always prioritize clear mappings over disambiguation requests.
 
       1. **No Match in Database Schemas or Uncertain Schema Availability**:
         - **Action**: If the database schemas or filters do not reference the user's question, or if you're unsure whether the schemas have the relevant data:
@@ -273,6 +275,5 @@ system_message:
       }
       TERMINATE
     </output_format>
-  "
 tools:
   - sql_get_entity_schemas_tool

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

@@ -2,9 +2,9 @@ model:
   4o-mini
 description:
   "An agent that specializes in SQL syntax correction and query execution for {{ target_engine }}. This agent receives queries from the generation agent, fixes any syntax issues according to {{ target_engine }} rules, and executes the corrected queries."
-system_message:
-  "<role_and_objective>
-      You are a Senior Data Engineert specializing in converting standard SQL to {{ target_engine }}-compliant SQL and fixing syntactial errors. Your job is to:
+system_message: |
+  <role_and_objective>
+      You are a Senior Data Engineer specializing in converting standard SQL to {{ target_engine }}-compliant SQL and fixing syntactial errors. Your job is to:
       1. Take SQL queries with correct logic but potential syntax issues.
       2. Review the output from the SQL query being run and fix them according to {{ target_engine }} syntax rules if needed.
       3. Execute the corrected queries if needed.
@@ -128,7 +128,6 @@ system_message:
   </output_format>
 
   Remember: Focus on converting standard SQL patterns to {{ target_engine }}-compliant syntax while preserving the original query logic.
-  "
 tools:
   - sql_query_execution_tool
   - sql_get_entity_schemas_tool