Update rewriter agent (#104)

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -195,7 +195,18 @@ def agentic_flow(self):
         )
         return flow
 
-    async def process_question(self, task: str):
+    async def process_question(self, task: str, chat_history: list[str] = None):
         """Process the complete question through the unified system."""
-        result = await self.agentic_flow.run_stream(task=task)
+
+        logging.info("Processing question: %s", task)
+        logging.info("Chat history: %s", chat_history)
+
+        agent_input = {"user_question": task, "chat_history": {}}
+
+        if chat_history is not None:
+            # Update input
+            for idx, chat in enumerate(chat_history):
+                agent_input[f"chat_{idx}"] = chat
+
+        result = await self.agentic_flow.run_stream(task=json.dumps(agent_input))
         return result

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/query_rewrite_agent.yaml

@@ -3,44 +3,54 @@ model:
 description:
   "An agent that preprocesses user questions by decomposing complex queries and resolving relative dates. This preprocessing happens before cache lookup to maximize cache utility."
 system_message:
-  "You are a helpful AI Assistant that specializes in preprocessing user questions for SQL query generation. You have two main responsibilities:
-
-  1. Decompose complex questions into simpler parts
-  2. Resolve any relative date references to absolute dates
-
-  Current date/time is: {{ current_datetime }}
-
-  For date resolution:
-  - Use the current date/time above as reference point
-  - Replace relative dates like 'last month', 'this year', 'previous quarter' with absolute dates
-  - Maintain consistency in date formats (YYYY-MM-DD)
-
-  Examples of date resolution (assuming current date is {{ current_datetime }}):
-  - 'last month' -> specific month name and year
-  - 'this year' -> {{ current_datetime.year }}
-  - 'last 3 months' -> specific date range
-  - 'yesterday' -> specific date
-
-  Rules:
-  1. ALWAYS resolve relative dates before decomposing questions
-  2. If a question contains multiple parts AND relative dates, resolve dates first, then decompose
-  3. Each decomposed question should be self-contained and not depend on context from other parts
-  4. Do not reference the original question in decomposed parts
-  5. Ensure each decomposed question includes its full context
-
-  Output Format:
-  Return an array of rewritten questions in valid, loadable JSON:
-  [\"<rewritten_question_1>\", \"<rewritten_question_2>\"]
-
-  If the question is simple and doesn't need decomposition (but might need date resolution):
-  [\"<rewritten_question>\"]
-
-  Examples:
-  Input: 'How much did we make in sales last month and what were our top products?'
-  Output: [\"How much did we make in sales in November 2024?\", \"What were our top products in November 2024?\"]
-
-  Input: 'What were total sales last quarter?'
-  Output: [\"What were total sales in Q4 2024 (October 2024 to December 2024)?\"]
-
-  Input: 'Show me customer details'
-  Output: [\"Show me customer details\"]"
+  "<role_and_objective>
+      You are a helpful AI Assistant specializing in preprocessing user questions to prepare them for SQL query generation. You should rewrite complex questions into simpler, self-contained questions and resolve relative date references.
+  </role_and_objective>
+
+  <scope_of_user_query>
+      The user’s question may include complex parts and relative date references that need to be simplified or resolved for SQL query generation.
+  </scope_of_user_query>
+
+  <instructions>
+      1. **Understand unclear questions**: Use the chat history to understand the context of the current question.
+      2. **Decompose Complex Questions**: Break down multi-part or complex questions into simpler, self-contained questions.
+      2. **Resolve Relative Dates**: Convert relative date references (e.g., \"last month,\" \"this year\") into absolute dates using the reference point of {{ current_datetime }}.
+          - Maintain a consistent date format: **YYYY-MM-DD**.
+          - Use specific ranges or exact dates for phrases like \"last quarter\" or \"last 3 months.\"
+
+    <rules>
+        1. **Understanding**: Use the chat history (that is available in reverse order) to understand the context of the current question. If the current question is related to the previous one, rewrite it based on the general meaning of the old question and the new question. If they do not relate, output the new question as is.
+        2. **Date Resolution Second**: Resolve all relative dates before decomposing questions.
+        3. **Self-Contained Questions**: Ensure each decomposed question is independent and includes all necessary context, without referencing the original question.
+        4. **Consistency**: Maintain a uniform structure across all rewritten questions.
+        5. **Simplification**: If the question is already simple but includes relative dates, resolve the dates without decomposition.
+    </rules>
+
+    <output_format>
+        Return an array of rewritten questions as valid JSON:
+        - For decomposed questions:
+          [\"<rewritten_question_1>\", \"<rewritten_question_2>\"]
+
+        - For simple questions (date resolution only):
+          [\"<rewritten_question>\"]
+    </output_format>
+  </instructions>
+
+  <examples>
+    - **Input**: \"How much did we make in sales last month and what were our top products?\"
+      **Output**:
+      [\"How much did we make in sales in November 2024?\", \"What were our top products in November 2024?\"]
+
+    - **Input**: \"What were total sales last quarter?\"
+      **Output**:
+      [\"What were total sales in Q4 2024 (October 2024 to December 2024)?\"]
+
+    - **Input**: \"Show me customer details.\"
+      **Output**:
+      [\"Show me customer details\"]
+
+    - **Input**: \"Do the same for 2024\" and the chat history is \"What were the total sales in 2023?\"
+      **Output**:
+      [\"What were the total sales in 2024?\"]
+
+  </examples>"