Update mesage types

Author: BenConstable9

text_2_sql/autogen/Iteration 5 - Agentic Vector Based Text2SQL.ipynb

@@ -50,7 +50,7 @@
       "source": [
         "import dotenv\n",
         "import logging\n",
-        "from autogen_text_2_sql import AutoGenText2Sql, AgentRequestBody"
+        "from autogen_text_2_sql import AutoGenText2Sql, QuestionBody, QuestionPayload"
       ]
     },
     {
@@ -100,16 +100,9 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "async for message in agentic_text_2_sql.process_question(AgentRequestBody(question=\"What total number of orders in June 2008?\")):\n",
+        "async for message in agentic_text_2_sql.process_question(QuestionPayload(body=QuestionBody(question=\"What total number of orders in June 2008?\"))):\n",
         "    logging.info(\"Received %s Message from Text2SQL System\", message)"
       ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": []
     }
   ],
   "metadata": {

text_2_sql/autogen/src/autogen_text_2_sql/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from autogen_text_2_sql.autogen_text_2_sql import AutoGenText2Sql
-from text_2_sql_core.payloads.agent_request_response_pair import AgentRequestBody
+from text_2_sql_core.payloads.interaction_payloads import QuestionBody, QuestionPayload
 
-__all__ = ["AutoGenText2Sql", "AgentRequestBody"]
+__all__ = ["AutoGenText2Sql", "QuestionBody", "QuestionPayload"]

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -17,17 +17,16 @@
 import os
 from datetime import datetime
 
-from text_2_sql_core.payloads.agent_request_response_pair import (
-    AgentRequestResponsePair,
-    AgentRequestBody,
-    AnswerWithSources,
+from text_2_sql_core.payloads.interaction_payloads import (
+    QuestionPayload,
+    AnswerWithSourcesPayload,
+    AnswerWithSourcesBody,
     Source,
-    DismabiguationRequests,
-)
-from text_2_sql_core.payloads.chat_history import ChatHistoryItem
-from text_2_sql_core.payloads.processing_update import (
+    DismabiguationRequestPayload,
     ProcessingUpdateBody,
-    ProcessingUpdate,
+    ProcessingUpdatePayload,
+    InteractionPayload,
+    PayloadType,
 )
 from autogen_agentchat.base import TaskResult
 from typing import AsyncGenerator
@@ -108,17 +107,19 @@ def agentic_flow(self):
         )
         return flow
 
-    def extract_disambiguation_request(self, messages: list) -> DismabiguationRequests:
+    def extract_disambiguation_request(
+        self, messages: list
+    ) -> DismabiguationRequestPayload:
         """Extract the disambiguation request from the answer."""
 
         disambiguation_request = messages[-1].content
 
         # TODO: Properly extract the disambiguation request
-        return DismabiguationRequests(
+        return DismabiguationRequestPayload(
             disambiguation_request=disambiguation_request,
         )
 
-    def extract_sources(self, messages: list) -> AnswerWithSources:
+    def extract_sources(self, messages: list) -> AnswerWithSourcesPayload:
         """Extract the sources from the answer."""
 
         answer = messages[-1].content
@@ -152,16 +153,18 @@ def extract_sources(self, messages: list) -> AnswerWithSources:
             logging.error("Could not load message: %s", sql_query_results)
             raise ValueError("Could not load message")
 
-        return AnswerWithSources(
-            answer=answer,
-            sources=sources,
+        return AnswerWithSourcesPayload(
+            body=AnswerWithSourcesBody(
+                answer=answer,
+                sources=sources,
+            )
         )
 
     async def process_question(
         self,
-        request: AgentRequestBody,
-        chat_history: list[ChatHistoryItem] = None,
-    ) -> AsyncGenerator[AgentRequestResponsePair | ProcessingUpdate, None]:
+        question_payload: QuestionPayload,
+        chat_history: list[InteractionPayload] = None,
+    ) -> AsyncGenerator[InteractionPayload, None]:
         """Process the complete question through the unified system.
 
         Args:
@@ -174,23 +177,22 @@ async def process_question(
         -------
             dict: The response from the system.
         """
-        logging.info("Processing question: %s", request.question)
+        logging.info("Processing question: %s", question_payload.body.question)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "question": request.question,
+            "question": question_payload.body.question,
             "chat_history": {},
-            "injected_parameters": request.injected_parameters,
+            "injected_parameters": question_payload.body.injected_parameters,
         }
 
         if chat_history is not None:
             # Update input
             for idx, chat in enumerate(chat_history):
-                # For now only consider the user query
-                chat_history_key = f"chat_{idx}"
-                agent_input[
-                    chat_history_key
-                ] = chat.request_response_pair.request.question
+                if chat.root.payload_type == PayloadType.QUESTION:
+                    # For now only consider the user query
+                    chat_history_key = f"chat_{idx}"
+                    agent_input[chat_history_key] = chat.root.body.question
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
             logging.debug("Message: %s", message)
@@ -213,27 +215,24 @@ async def process_question(
                     )
 
                 if processing_update is not None:
-                    payload = ProcessingUpdate(
-                        processing_update=processing_update,
+                    payload = ProcessingUpdatePayload(
+                        body=processing_update,
                     )
 
             elif isinstance(message, TaskResult):
                 # Now we need to return the final answer or the disambiguation request
                 logging.info("TaskResult: %s", message)
 
-                response = None
                 if message.messages[-1].source == "answer_agent":
                     # If the message is from the answer_agent, we need to return the final answer
-                    response = self.extract_sources(message.messages)
+                    payload = self.extract_sources(message.messages)
                 elif message.messages[-1].source == "parallel_query_solving_agent":
                     # Load into disambiguation request
-                    response = self.extract_disambiguation_request(message.messages)
+                    payload = self.extract_disambiguation_request(message.messages)
                 else:
                     logging.error("Unexpected TaskResult: %s", message)
                     raise ValueError("Unexpected TaskResult")
 
-                payload = AgentRequestResponsePair(request=request, response=response)
-
             if payload is not None:
                 logging.debug("Payload: %s", payload)
                 yield payload