Add aiostream

Author: BenConstable9

text_2_sql/autogen/pyproject.toml

@@ -5,6 +5,7 @@ description = "AutoGen Based Implementation"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "aiostream>=0.6.4",
     "autogen-agentchat==0.4.0.dev11",
     "autogen-core==0.4.0.dev11",
     "autogen-ext[azure,openai]==0.4.0.dev11",

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -11,7 +11,7 @@
 from autogen_text_2_sql.custom_agents.parallel_query_solving_agent import (
     ParallelQuerySolvingAgent,
 )
-from text_2_sql.autogen.src.autogen_text_2_sql.custom_agents.sources_agent import (
+from autogen_text_2_sql.custom_agents.sources_agent import (
     SourcesAgent,
 )
 from autogen_agentchat.agents import UserProxyAgent
@@ -120,7 +120,7 @@ def agentic_flow(self):
         )
         return flow
 
-    async def process_question(
+    def process_question(
         self,
         question: str,
         chat_history: list[str] = None,

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

@@ -5,17 +5,17 @@
 from autogen_agentchat.agents import BaseChatAgent
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import (
-    AgentEvent,
     AgentMessage,
     ChatMessage,
     TextMessage,
 )
 from autogen_core import CancellationToken
 import json
 import logging
-import asyncio
 from autogen_text_2_sql.inner_autogen_text_2_sql import InnerAutoGenText2Sql
 
+from aiostream import stream
+
 
 class ParallelQuerySolvingAgent(BaseChatAgent):
     def __init__(self, engine_specific_rules: str, **kwargs: dict):
@@ -45,11 +45,10 @@ async def on_messages(
     async def on_messages_stream(
         self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
     ) -> AsyncGenerator[AgentMessage | Response, None]:
-        inner_messages: List[AgentEvent | ChatMessage] = []
+        inner_messages: List[AgentMessage | ChatMessage] = []
 
         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:
@@ -61,31 +60,71 @@ async def on_messages_stream(
 
         logging.info(f"Query Rewrite: {query_rewrites}")
 
-        inner_solving_tasks = []
+        inner_solving_generators = []
+
+        async def consume_inner_messages_from_agentic_flow(
+            agentic_flow, identifier, complete_inner_messages
+        ):
+            """
+            Consume the inner messages and append them to the specified list.
+
+            Args:
+            ----
+                agentic_flow: The async generator to consume messages from.
+                messages_list: The list to which messages should be added.
+            """
+            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)
+
+                yield {"source": identifier, "message": inner_message}
 
+        complete_inner_messages = {}
+
+        # Start processing sub-queries
         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(
+            # Launch tasks for each sub-query
+            inner_solving_generators.append(
                 inner_autogen_text_2_sql.process_question(
                     question=query_rewrite, parameters=user_parameters
                 )
             )
 
-        # Wait for all the inner solving tasks to complete
-        inner_solving_results = await asyncio.gather(*inner_solving_tasks)
+        combined_message_streams = stream.merge(*inner_solving_generators)
+
+        async with combined_message_streams.stream() as streamer:
+            async for inner_message in streamer:
+                print(inner_message)
+                yield inner_message
+
+        # # Process the results as they are yielded
+        # for completed in asyncio.as_completed(inner_solving_generators):
+        #     async for inner_message in completed:
+        #         # Yield the result as soon as it's available
+        #         yield inner_message
+
+        # # Wait for all tasks to complete
+        # await asyncio.gather(*inner_solving_generators, return_exceptions=True)
+
+        # # Log final results for debugging or auditing
+        # logging.info(f"Formatted Results: {complete_inner_messages}")
 
-        logging.info(f"Inner Solving Results: {inner_solving_results}")
+        # TODO: Trim out unnecessary information from the final response
 
+        # Final response
         yield Response(
             chat_message=TextMessage(
-                content=json.dumps(inner_solving_results), source=self.name
+                content=json.dumps(complete_inner_messages), source=self.name
             ),
-            inner_messages=inner_messages,
+            inner_messages=complete_inner_messages,
         )
 
-        async def on_reset(self, cancellation_token: CancellationToken) -> None:
-            pass
+    async def on_reset(self, cancellation_token: CancellationToken) -> None:
+        pass

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -188,7 +188,7 @@ def agentic_flow(self):
         )
         return flow
 
-    async def process_question(
+    def process_question(
         self,
         question: str,
         parameters: dict = None,

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/answer_agent.yaml

@@ -0,0 +1,6 @@
+model: "4o-mini"
+description: "An agent that generates a response to a user's question."
+system_message: |
+  <role_and_objective>
+      You are a helpful AI Assistant specializing in answering a user's question.
+  </role_and_objective>