Update steps

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -51,18 +51,18 @@ def __init__(self, state_store: StateStore, **kwargs):
     def get_all_agents(self):
         """Get all agents for the complete flow."""
 
-        self.user_message_rewrite_agent = LLMAgentCreator.create(
+        user_message_rewrite_agent = LLMAgentCreator.create(
             "user_message_rewrite_agent", **self.kwargs
         )
 
-        self.parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
+        parallel_query_solving_agent = ParallelQuerySolvingAgent(**self.kwargs)
 
-        self.answer_agent = LLMAgentCreator.create("answer_agent", **self.kwargs)
+        answer_agent = LLMAgentCreator.create("answer_agent", **self.kwargs)
 
         agents = [
-            self.user_message_rewrite_agent,
-            self.parallel_query_solving_agent,
-            self.answer_agent,
+            user_message_rewrite_agent,
+            parallel_query_solving_agent,
+            answer_agent,
         ]
 
         return agents
@@ -71,8 +71,11 @@ def get_all_agents(self):
     def termination_condition(self):
         """Define the termination condition for the chat."""
         termination = (
-            TextMentionTermination("TERMINATE")
-            | SourceMatchTermination("answer_agent")
+            SourceMatchTermination("answer_agent")
+            # | TextMentionTermination(
+            #     "[]",
+            #     sources=["user_message_rewrite_agent"],
+            # )
             | TextMentionTermination(
                 "contains_disambiguation_requests",
                 sources=["parallel_query_solving_agent"],

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/user_message_rewrite_agent.yaml

@@ -42,7 +42,6 @@ system_message: |
           - Consider whether the question relates to **data analysis** or is **possibly related** to `{{ use_case }}`.
           - If unsure whether a question is relevant to the use case, **do not filter it out**.
           - If the question **is disallowed**, return an **empty list of steps** in the JSON output.
-          - `"TERMINATE"` is an **internal rule**—it is **not** included in the JSON output.
           - **Set `"requires_sql_queries": true`** if the question requires database access to answer.
             **Set `"requires_sql_queries": false`** if it does not (e.g., "What can you help me with?").
 

text_2_sql/text_2_sql_core/src/text_2_sql_core/structured_outputs/answer_agent_with_follow_up_questions.py

@@ -4,5 +4,7 @@
 
 
 class AnswerAgentWithFollowUpQuestionsAgentOutput(BaseModel):
+    """The output of the answer agent with follow up questions."""
+
     answer: str
     follow_up_questions: list[str]