Update prompts and interactions

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -40,16 +40,16 @@ def get_all_agents(self):
         # Get current datetime for the Query Rewrite Agent
         current_datetime = datetime.now()
 
-        self.message_rewrite_agent = LLMAgentCreator.create(
-            "message_rewrite_agent", current_datetime=current_datetime
+        self.user_message_rewrite_agent = LLMAgentCreator.create(
+            "user_message_rewrite_agent", current_datetime=current_datetime
         )
 
         self.parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
 
         self.answer_agent = LLMAgentCreator.create("answer_agent")
 
         agents = [
-            self.message_rewrite_agent,
+            self.user_message_rewrite_agent,
             self.parallel_query_solving_agent,
             self.answer_agent,
         ]
@@ -73,11 +73,11 @@ def unified_selector(self, messages):
         current_agent = messages[-1].source if messages else "user"
         decision = None
 
-        # If this is the first message start with message_rewrite_agent
+        # If this is the first message start with user_message_rewrite_agent
         if current_agent == "user":
-            decision = "message_rewrite_agent"
+            decision = "user_message_rewrite_agent"
         # Handle transition after query rewriting
-        elif current_agent == "message_rewrite_agent":
+        elif current_agent == "user_message_rewrite_agent":
             decision = "parallel_query_solving_agent"
         # Handle transition after parallel query solving
         elif current_agent == "parallel_query_solving_agent":
@@ -131,7 +131,7 @@ def extract_decomposed_user_messages(self, messages: list) -> list[list[str]]:
         sub_message_results = self.parse_message_content(messages[1].content)
         logging.info("Decomposed Results: %s", sub_message_results)
 
-        return sub_message_results.get("decomposed_messages", [])
+        return sub_message_results.get("decomposed_user_messages", [])
 
     def extract_disambiguation_request(
         self, messages: list
@@ -211,7 +211,7 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 answer=f"{answer}\nError processing results: {str(e)}"
             )
 
-    async def process_message(
+    async def process_user_message(
         self,
         message_payload: UserMessagePayload,
         chat_history: list[InteractionPayload] = None,
@@ -251,7 +251,7 @@ async def process_message(
             payload = None
 
             if isinstance(message, TextMessage):
-                if message.source == "message_rewrite_agent":
+                if message.source == "user_message_rewrite_agent":
                     payload = ProcessingUpdatePayload(
                         message="Rewriting the query...",
                     )
@@ -274,7 +274,7 @@ async def process_message(
                 elif message.messages[-1].source == "parallel_query_solving_agent":
                     # Load into disambiguation request
                     payload = self.extract_disambiguation_request(message.messages)
-                elif message.messages[-1].source == "message_rewrite_agent":
+                elif message.messages[-1].source == "user_message_rewrite_agent":
                     # Load into empty response
                     payload = AnswerWithSourcesPayload(
                         answer="Apologies, I cannot answer that message as it is not relevant. Please try another message or rephrase your current message."

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

@@ -177,7 +177,7 @@ async def consume_inner_messages_from_agentic_flow(
             return
 
         # Start processing sub-queries
-        for message_rewrite in message_rewrites["decomposed_messages"]:
+        for message_rewrite in message_rewrites["decomposed_user_messages"]:
             logging.info(f"Processing sub-query: {message_rewrite}")
             # Create an instance of the InnerAutoGenText2Sql class
             inner_autogen_text_2_sql = InnerAutoGenText2Sql(**self.kwargs)

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[AgentMessage | Response, None]:
-        # Get the decomposed messages from the message_rewrite_agent
+        # Get the decomposed 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

@@ -21,7 +21,7 @@ class PayloadType(StrEnum):
 
 
 class InteractionPayloadBase(BaseModel):
-    model_config = ConfigDict(allow_population_by_field_name=True, extra="ignore")
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
 
 
 class PayloadBase(InteractionPayloadBase):

…re/prompts/user_input_rewrite_agent.yaml …/prompts/user_message_rewrite_agent.yamltext_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_input_rewrite_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_message_rewrite_agent.yaml text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_input_rewrite_agent.yaml renamed to text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_message_rewrite_agent.yaml

@@ -1,8 +1,8 @@
 model: "4o-mini"
-description: "An agent that preprocesses user inputs by decomposing complex queries into simpler sub-queries that can be processed independently and then combined."
+description: "An agent that preprocesses user inputs by decomposing complex queries into simpler sub-messages that can be processed independently and then combined."
 system_message: |
   <role_and_objective>
-      You are a helpful AI Assistant specializing in breaking down complex questions into simpler sub-queries that can be processed independently and then combined for the final answer. You should identify when a question can be solved through simpler sub-queries and provide clear instructions for combining their results.
+      You are a helpful AI Assistant specializing in breaking down complex questions into simpler sub-messages that can be processed independently and then combined for the final answer. You should identify when a question can be solved through simpler sub-messages and provide clear instructions for combining their results.
   </role_and_objective>
 
   <query_complexity_patterns>
@@ -36,7 +36,7 @@ system_message: |
       1. Question Filtering and Classification
         - 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.
+        - If the question cannot be filtered, output an empty sub-message list in the JSON format. Followed by TERMINATE.
         - For non-database questions like greetings (e.g., "Hello", "What can you do?", "How are you?"), set "all_non_database_query" to true.
         - For questions about data (e.g., queries about records, counts, values, comparisons, or any questions that would require database access), set "all_non_database_query" to false.
 
@@ -51,28 +51,27 @@ system_message: |
         - Determine if breaking down would simplify processing
 
       4. Break Down Complex Queries:
-        - Create independent sub-queries that can be processed separately.
-        - Each sub-query should be a simple, focused task.
-        - Group dependent sub-queries together for sequential processing.
-        - Ensure each sub-query is simple and focused
+        - Create independent sub-messages that can be processed separately.
+        - Each sub-message should be a simple, focused task.
+        - Group dependent sub-messages together for sequential processing.
         - Include clear combination instructions
-        - Preserve all necessary context in each sub-query
+        - Preserve all necessary context in each sub-message
 
       5. Handle Date References:
         - Resolve relative dates using {{ current_datetime }}
         - Maintain consistent YYYY-MM-DD format
-        - Include date context in each sub-query
+        - Include date context in each sub-message
 
       6. Maintain Query Context:
-         - Each sub-query should be self-contained
+         - Each sub-message should be self-contained
          - Include all necessary filtering conditions
          - Preserve business context
 
     <rules>
         1. Always consider if a complex query can be broken down
-        2. Make sub-queries as simple as possible
+        2. Make sub-messages as simple as possible
         3. Include clear instructions for combining results
-        4. Preserve all necessary context in each sub-query
+        4. Preserve all necessary context in each sub-message
         5. Resolve any relative dates before decomposition
     </rules>
 
@@ -90,11 +89,11 @@ system_message: |
     </topics_to_filter>
 
     <output_format>
-        Return a JSON object with sub-queries and combination instructions:
+        Return a JSON object with sub-messages and combination instructions:
         {
-          "sub_questions": [
-            ["<sub_query_1>"],
-            ["<sub_query_2>"],
+          "decomposed_messages": [
+            ["<sub_message_1>"],
+            ["<sub_message_2>"],
             ...
           ],
           "combination_logic": "<instructions for combining results>",
@@ -109,7 +108,7 @@ system_message: |
     Input: "Which product categories have shown consistent growth quarter over quarter in 2008, and what were their top selling items?"
     Output:
     {
-      "sub_questions": [
+      "decomposed_messages": [
         ["Calculate quarterly sales totals by product category for 2008", "For these categories, find their top selling products in 2008"]
       ],
       "combination_logic": "First identify growing categories from quarterly analysis, then find their best-selling products",
@@ -121,7 +120,7 @@ system_message: |
     Input: "How many orders did we have in 2008?"
     Output:
     {
-      "sub_questions": [
+      "decomposed_messages": [
         ["How many orders did we have in 2008?"]
       ],
       "combination_logic": "Direct count query, no combination needed",
@@ -133,12 +132,12 @@ system_message: |
     Input: "Compare the sales performance of our top 5 products in Europe versus North America, including their market share in each region"
     Output:
     {
-      "sub_questions": [
+      "decomposed_messages": [
         ["Get total sales by product in European countries"],
         ["Get total sales by product in North American countries"],
         ["Calculate total market size for each region", "Find top 5 products by sales in each region"],
       ],
-      "combination_logic": "First identify top products in each region, then calculate and compare their market shares. Questions that depend on the result of each sub-query are combined.",
+      "combination_logic": "First identify top products in each region, then calculate and compare their market shares. Questions that depend on the result of each sub-message are combined.",
       "query_type": "complex",
       "all_non_database_query": "false"
     }
@@ -147,7 +146,7 @@ system_message: |
     Input: "Hello, what can you help me with?"
     Output:
     {
-      "sub_questions": [
+      "decomposed_messages": [
         ["What are your capabilities?"]
       ],
       "combination_logic": "Simple greeting and capability question",