Update answer

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -3,6 +3,7 @@
 from autogen_agentchat.conditions import (
     TextMentionTermination,
     MaxMessageTermination,
+    SourceMatchTermination,
 )
 from autogen_agentchat.teams import SelectorGroupChat
 from autogen_text_2_sql.creators.llm_model_creator import LLMModelCreator
@@ -11,9 +12,6 @@
 from autogen_text_2_sql.custom_agents.parallel_query_solving_agent import (
     ParallelQuerySolvingAgent,
 )
-from autogen_text_2_sql.custom_agents.sources_agent import (
-    SourcesAgent,
-)
 from autogen_agentchat.agents import UserProxyAgent
 from autogen_agentchat.messages import TextMessage
 import json
@@ -66,16 +64,14 @@ def get_all_agents(self):
 
         self.answer_agent = LLMAgentCreator.create("answer_agent")
 
-        self.sources_agent = SourcesAgent()
-
         # Auto-responding UserProxyAgent
         self.user_proxy = EmptyResponseUserProxyAgent(name="user_proxy")
 
         agents = [
             self.user_proxy,
             self.query_rewrite_agent,
             self.parallel_query_solving_agent,
-            self.sources_agent,
+            self.answer_agent,
         ]
 
         return agents
@@ -85,8 +81,9 @@ def termination_condition(self):
         """Define the termination condition for the chat."""
         termination = (
             TextMentionTermination("TERMINATE")
-            | (TextMentionTermination("answer") & TextMentionTermination("sources"))
-            | MaxMessageTermination(20)
+            | SourceMatchTermination("answer_agent")
+            | TextMentionTermination("requires_user_information_request")
+            | MaxMessageTermination(5)
         )
         return termination
 
@@ -105,8 +102,6 @@ def unified_selector(self, messages):
         # Handle transition after parallel query solving
         elif current_agent == "parallel_query_solving_agent":
             decision = "answer_agent"
-        elif current_agent == "answer_agent":
-            decision = "sources_agent"
 
         if decision:
             logging.info(f"Agent transition: {current_agent} -> {decision}")
@@ -127,6 +122,45 @@ def agentic_flow(self):
         )
         return flow
 
+    def extract_sources(self, messages: list) -> AnswerWithSources:
+        """Extract the sources from the answer."""
+
+        answer = messages[-1].content
+
+        sql_query_results = messages[-2].content
+
+        try:
+            sql_query_results = json.loads(sql_query_results)
+
+            logging.info("SQL Query Results: %s", sql_query_results)
+
+            sources = []
+
+            for question, sql_query_result_list in sql_query_results["results"].items():
+                logging.info(
+                    "SQL Query Result for question '%s': %s",
+                    question,
+                    sql_query_result_list,
+                )
+
+                for sql_query_result in sql_query_result_list:
+                    logging.info("SQL Query Result: %s", sql_query_result)
+                    sources.append(
+                        {
+                            "sql_query": sql_query_result["sql_query"],
+                            "sql_rows": sql_query_result["sql_rows"],
+                        }
+                    )
+
+        except json.JSONDecodeError:
+            logging.error("Could not load message: %s", sql_query_results)
+            raise ValueError("Could not load message")
+
+        return AnswerWithSources(
+            answer=answer,
+            sources=sources,
+        )
+
     async def process_question(
         self,
         question: str,
@@ -160,8 +194,7 @@ async def process_question(
                 agent_input[f"chat_{idx}"] = chat
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
-            logging.info("Message: %s", message)
-            logging.info("Message type: %s", type(message))
+            logging.debug("Message: %s", message)
 
             payload = None
 
@@ -184,17 +217,19 @@ async def process_question(
 
             elif isinstance(message, TaskResult):
                 # Now we need to return the final answer or the disambiguation request
+                logging.info("TaskResult: %s", message)
 
-                if message.source == "answer_agent":
+                if message.messages[-1].source == "answer_agent":
                     # If the message is from the answer_agent, we need to return the final answer
-                    payload = AnswerWithSources(
-                        **json.loads(message.content),
-                    )
-                else:
+                    payload = self.extract_sources(message.messages)
+                elif message.messages[-1].source == "parallel_query_solving_agent":
                     payload = UserInformationRequest(
-                        **json.loads(message.content),
+                        **json.loads(message.messages[-1].content),
                     )
+                else:
+                    logging.error("Unexpected TaskResult: %s", message)
+                    raise ValueError("Unexpected TaskResult")
 
             if payload is not None:
-                logging.info("Payload: %s", payload)
+                logging.debug("Payload: %s", payload)
                 yield payload

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

@@ -3,18 +3,14 @@
 from typing import AsyncGenerator, List, Sequence
 
 from autogen_agentchat.agents import BaseChatAgent
-from autogen_agentchat.base import Response
-from autogen_agentchat.messages import (
-    AgentMessage,
-    ChatMessage,
-    TextMessage,
-)
+from autogen_agentchat.base import Response, TaskResult
+from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage
 from autogen_core import CancellationToken
 import json
 import logging
 from autogen_text_2_sql.inner_autogen_text_2_sql import InnerAutoGenText2Sql
-
 from aiostream import stream
+from json import JSONDecodeError
 
 
 class ParallelQuerySolvingAgent(BaseChatAgent):
@@ -59,7 +55,7 @@ async def on_messages_stream(
         logging.info(f"Query Rewrites: {query_rewrites}")
 
         async def consume_inner_messages_from_agentic_flow(
-            agentic_flow, identifier, complete_inner_messages
+            agentic_flow, identifier, database_results
         ):
             """
             Consume the inner messages and append them to the specified list.
@@ -71,14 +67,43 @@ async def consume_inner_messages_from_agentic_flow(
             """
             async for inner_message in agentic_flow:
                 # Add message to results dictionary, tagged by the function name
-                if identifier not in complete_inner_messages:
-                    complete_inner_messages[identifier] = []
-                complete_inner_messages[identifier].append(inner_message)
+                if identifier not in database_results:
+                    database_results[identifier] = []
+
+                logging.info(f"Checking Inner Message: {inner_message}")
+
+                if isinstance(inner_message, TaskResult) is False:
+                    try:
+                        inner_message = json.loads(inner_message.content)
+                        logging.info(f"Loaded: {inner_message}")
+
+                        # Search for specific message types and add them to the final output object
+                        if (
+                            "type" in inner_message
+                            and inner_message["type"] == "query_execution_with_limit"
+                        ):
+                            database_results[identifier].append(
+                                {
+                                    "sql_query": inner_message["sql_query"].replace(
+                                        "\n", " "
+                                    ),
+                                    "sql_rows": inner_message["sql_rows"],
+                                }
+                            )
+
+                    except (JSONDecodeError, TypeError) as e:
+                        logging.error("Could not load message: %s", inner_message)
+                        logging.warning(f"Error processing message: {e}")
+
+                    except Exception as e:
+                        logging.error("Could not load message: %s", inner_message)
+                        logging.error(f"Error processing message: {e}")
+                        raise e
 
                 yield inner_message
 
         inner_solving_generators = []
-        complete_inner_messages = {}
+        database_results = {}
 
         # Start processing sub-queries
         for query_rewrite in query_rewrites["sub_queries"]:
@@ -95,32 +120,33 @@ async def consume_inner_messages_from_agentic_flow(
                         question=query_rewrite, parameters=user_parameters
                     ),
                     query_rewrite,
-                    complete_inner_messages,
+                    database_results,
                 )
             )
 
-        logging.info("Created %i Inner Solving Generators", inner_solving_generators)
+        logging.info(
+            "Created %i Inner Solving Generators", len(inner_solving_generators)
+        )
         logging.info("Starting Inner Solving Generators")
         combined_message_streams = stream.merge(*inner_solving_generators)
 
         async with combined_message_streams.stream() as streamer:
             async for inner_message in streamer:
-                logging.info(f"Inner Solving Message: {inner_message}")
-                yield inner_message
+                if isinstance(inner_message, TextMessage):
+                    logging.debug(f"Inner Solving Message: {inner_message}")
+                    yield inner_message
 
         # Log final results for debugging or auditing
-        logging.info(f"Formatted Results: {complete_inner_messages}")
+        logging.info(f"Database Results: {database_results}")
 
-        # TODO: Trim out unnecessary information from the final response
         # Final response
         yield Response(
             chat_message=TextMessage(
-                content=json.dumps(complete_inner_messages), source=self.name
+                content=json.dumps(
+                    {"contains_results": True, "results": database_results}
+                ),
+                source=self.name,
             ),
-            inner_messages=[
-                complete_inner_message["message"]
-                for complete_inner_message in complete_inner_messages
-            ],
         )
 
     async def on_reset(self, cancellation_token: CancellationToken) -> None:

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

@@ -150,7 +150,7 @@ def unified_selector(self, messages):
                 correction_result = json.loads(messages[-1].content)
                 if isinstance(correction_result, dict):
                     if "answer" in correction_result and "sources" in correction_result:
-                        decision = "answer_and_sources_agent"
+                        decision = "user_proxy"
                     elif "corrected_query" in correction_result:
                         if correction_result.get("executing", False):
                             decision = "sql_query_correction_agent"
@@ -166,8 +166,6 @@ def unified_selector(self, messages):
                     decision = "sql_query_generation_agent"
             except json.JSONDecodeError:
                 decision = "sql_query_generation_agent"
-        elif current_agent == "answer_and_sources_agent":
-            decision = "user_proxy"  # Let user_proxy send TERMINATE
 
         if decision:
             logging.info(f"Agent transition: {current_agent} -> {decision}")

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -4,9 +4,8 @@
 class Source(BaseModel):
     sql_query: str
     sql_rows: list[dict]
-    markdown_table: str
 
 
 class AnswerWithSources(BaseModel):
     answer: str
-    sources: list[str] = Field(default_factory=list)
+    sources: list[Source] = Field(default_factory=list)

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/answer_with_sources.py

@@ -4,3 +4,9 @@ system_message: |
   <role_and_objective>
       You are a helpful AI Assistant specializing in answering a user's question.
   </role_and_objective>
+
+  Use the information obtained to generate a response to the user's question. The question has been broken down into a series of SQL queries and you need to generate a response based on the results of these queries.
+
+  Do not use any external resources to generate the response. The response should be based solely on the information provided in the SQL queries and their results.
+
+  You can use Markdown and Markdown tables to format the response.