Separate out sources and answer agent

Author: BenConstable9

text_2_sql/autogen/src/__init__.py

@@ -11,8 +11,8 @@
 from autogen_text_2_sql.custom_agents.parallel_query_solving_agent import (
     ParallelQuerySolvingAgent,
 )
-from autogen_text_2_sql.custom_agents.answer_and_sources_agent import (
-    AnswerAndSourcesAgent,
+from text_2_sql.autogen.src.autogen_text_2_sql.custom_agents.sources_agent import (
+    SourcesAgent,
 )
 from autogen_agentchat.agents import UserProxyAgent
 from autogen_agentchat.messages import TextMessage
@@ -57,7 +57,9 @@ def get_all_agents(self):
             engine_specific_rules=self.engine_specific_rules, **self.kwargs
         )
 
-        self.answer_and_sources_agent = AnswerAndSourcesAgent()
+        self.answer_agent = LLMAgentCreator.create("answer_agent")
+
+        self.sources_agent = SourcesAgent()
 
         # Auto-responding UserProxyAgent
         self.user_proxy = EmptyResponseUserProxyAgent(name="user_proxy")
@@ -66,7 +68,7 @@ def get_all_agents(self):
             self.user_proxy,
             self.query_rewrite_agent,
             self.parallel_query_solving_agent,
-            self.answer_and_sources_agent,
+            self.sources_agent,
         ]
 
         return agents
@@ -95,7 +97,9 @@ def unified_selector(self, messages):
             decision = "parallel_query_solving_agent"
         # Handle transition after parallel query solving
         elif current_agent == "parallel_query_solving_agent":
-            decision = "answer_and_sources_agent"
+            decision = "answer_agent"
+        elif current_agent == "answer_agent":
+            decision = "sources_agent"
 
         if decision:
             logging.info(f"Agent transition: {current_agent} -> {decision}")
@@ -118,7 +122,7 @@ def agentic_flow(self):
 
     async def process_question(
         self,
-        task: str,
+        question: str,
         chat_history: list[str] = None,
         parameters: dict = None,
     ):
@@ -134,11 +138,11 @@ async def process_question(
         -------
             dict: The response from the system.
         """
-        logging.info("Processing question: %s", task)
+        logging.info("Processing question: %s", question)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "user_question": task,
+            "question": question,
             "chat_history": {},
             "parameters": parameters,
         }

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -41,6 +41,13 @@ async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | Response, None]:
         last_response = messages[-1].content
+        parameter_input = messages[0].content
+        last_response = messages[-1].content
+        try:
+            user_parameters = json.loads(parameter_input)["parameters"]
+        except json.JSONDecodeError:
+            logging.error("Error decoding the user parameters.")
+            user_parameters = {}
 
         # Load the json of the last message to populate the final output object
         query_rewrites = json.loads(last_response)
@@ -49,14 +56,16 @@ async def on_messages_stream(
 
         inner_solving_tasks = []
 
-        for query_rewrite in query_rewrites:
+        for query_rewrite in query_rewrites["sub_queries"]:
             # Create an instance of the InnerAutoGenText2Sql class
             inner_autogen_text_2_sql = InnerAutoGenText2Sql(
                 self.engine_specific_rules, **self.kwargs
             )
 
             inner_solving_tasks.append(
-                inner_autogen_text_2_sql.run_stream(task=query_rewrite)
+                inner_autogen_text_2_sql.process_question(
+                    question=query_rewrite, parameters=user_parameters
+                )
             )
 
         # Wait for all the inner solving tasks to complete

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

@@ -12,10 +12,10 @@
 import pandas as pd
 
 
-class AnswerAndSourcesAgent(BaseChatAgent):
+class SourcesAgent(BaseChatAgent):
     def __init__(self):
         super().__init__(
-            "answer_and_sources_agent",
+            "sources_agent",
             "An agent that formats the final answer and sources.",
         )
 

…ustom_agents/answer_and_sources_agent.py …ext_2_sql/custom_agents/sources_agent.pytext_2_sql/autogen/src/autogen_text_2_sql/custom_agents/answer_and_sources_agent.py renamed to text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/sources_agent.py text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/answer_and_sources_agent.py renamed to text_2_sql/autogen/src/autogen_text_2_sql/custom_agents/sources_agent.py

@@ -187,3 +187,29 @@ def agentic_flow(self):
             selector_func=self.unified_selector,
         )
         return flow
+
+    async def process_question(
+        self,
+        question: str,
+        parameters: dict = None,
+    ):
+        """Process the complete question through the unified system.
+
+        Args:
+        ----
+            task (str): The user question to process.
+            parameters (dict, optional): Parameters to pass to agents. Defaults to None.
+
+        Returns:
+        -------
+            dict: The response from the system.
+        """
+        logging.info("Processing question: %s", question)
+
+        agent_input = {
+            "question": question,
+            "chat_history": {},
+            "parameters": parameters,
+        }
+
+        return self.agentic_flow.run_stream(task=json.dumps(agent_input))