Update agentic flow

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -86,11 +86,11 @@ def termination_condition(self):
     def unified_selector(self, messages):
         """Unified selector for the complete flow."""
         logging.info("Messages: %s", messages)
-        current_agent = messages[-1].source if messages else "start"
+        current_agent = messages[-1].source if messages else "user"
         decision = None
 
         # If this is the first message start with query_rewrite_agent
-        if current_agent == "start":
+        if current_agent == "user":
             decision = "query_rewrite_agent"
         # Handle transition after query rewriting
         elif current_agent == "query_rewrite_agent":

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

@@ -75,7 +75,10 @@ async def consume_inner_messages_from_agentic_flow(
                     complete_inner_messages[identifier] = []
                 complete_inner_messages[identifier].append(inner_message)
 
-                yield {"source": identifier, "message": inner_message}
+                yield TextMessage(
+                    content=json.dumps(inner_message),
+                    source=f"{self.name}-{identifier}",
+                )
 
         inner_solving_generators = []
         complete_inner_messages = {}

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -126,10 +126,10 @@ def unified_selector(self, messages):
         """Unified selector for the complete flow."""
         """Unified selector for the complete flow."""
         logging.info("Messages: %s", messages)
-        current_agent = messages[-1].source if messages else "start"
+        current_agent = messages[-1].source if messages else "user"
         decision = None
 
-        if current_agent == "start":
+        if current_agent == "user":
             decision = (
                 "sql_query_cache_agent"
                 if self.use_query_cache