Update entry points

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, QuestionPayload"
+        "from autogen_text_2_sql import AutoGenText2Sql, UserMessagePayload"
       ]
     },
     {
@@ -100,7 +100,7 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "async for message in agentic_text_2_sql.process_question(QuestionPayload(question=\"What is the total number of sales?\")):\n",
+        "async for message in agentic_text_2_sql.process_user_message(UserMessagePayload(user_message=\"What is the total number of sales?\")):\n",
         "    logging.info(\"Received %s Message from Text2SQL System\", message)"
       ]
     },

text_2_sql/autogen/evaluate_autogen_text2sql.ipynb

@@ -68,7 +68,7 @@
         "# Add the src directory to the path\n",
         "sys.path.append(str(notebook_dir / \"src\"))\n",
         "\n",
-        "from autogen_text_2_sql import AutoGenText2Sql, QuestionPayload\n",
+        "from autogen_text_2_sql import AutoGenText2Sql, UserMessagePayload\n",
         "from autogen_text_2_sql.evaluation_utils import get_final_sql_query\n",
         "\n",
         "# Configure logging\n",
@@ -127,7 +127,7 @@
         "    all_queries = []\n",
         "    final_query = None\n",
         "    \n",
-        "    async for message in autogen_text2sql.process_question(QuestionPayload(question=question)):\n",
+        "    async for message in autogen_text2sql.process_user_message(UserMessagePayload(user_message=question)):\n",
         "        if message.payload_type == \"answer_with_sources\":\n",
         "            # Extract from results\n",
         "            if hasattr(message.body, 'results'):\n",

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.interaction_payloads import UserInputPayload
+from text_2_sql_core.payloads.interaction_payloads import UserMessagePayload
 
-__all__ = ["AutoGenText2Sql", "UserInputPayload"]
+__all__ = ["AutoGenText2Sql", "UserMessagePayload"]

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -19,7 +19,7 @@
 import re
 
 from text_2_sql_core.payloads.interaction_payloads import (
-    UserInputPayload,
+    UserMessagePayload,
     AnswerWithSourcesPayload,
     DismabiguationRequestsPayload,
     ProcessingUpdatePayload,
@@ -102,15 +102,6 @@ def agentic_flow(self):
         )
         return flow
 
-    def extract_disambiguation_request(
-        self, messages: list
-    ) -> DismabiguationRequestsPayload:
-        """Extract the disambiguation request from the answer."""
-        disambiguation_request = messages[-1].content
-        return DismabiguationRequestsPayload(
-            disambiguation_request=disambiguation_request,
-        )
-
     def parse_message_content(self, content):
         """Parse different message content formats into a dictionary."""
         if isinstance(content, (list, dict)):
@@ -134,6 +125,26 @@ def parse_message_content(self, content):
         # If all parsing attempts fail, return the content as-is
         return content
 
+    def extract_decomposed_user_messages(self, messages: list) -> list[list[str]]:
+        """Extract the decomposed messages from the answer."""
+        # Only load sub-message results if we have a database result
+        sub_message_results = self.parse_message_content(messages[1].content)
+        logging.info("Decomposed Results: %s", sub_message_results)
+
+        return sub_message_results.get("decomposed_messages", [])
+
+    def extract_disambiguation_request(
+        self, messages: list
+    ) -> DismabiguationRequestsPayload:
+        """Extract the disambiguation request from the answer."""
+        disambiguation_request = messages[-1].content
+
+        decomposed_user_messages = self.extract_decomposed_user_messages(messages)
+        return DismabiguationRequestsPayload(
+            disambiguation_request=disambiguation_request,
+            decomposed_user_messages=decomposed_user_messages,
+        )
+
     def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
         """Extract the sources from the answer."""
         answer = messages[-1].content
@@ -145,24 +156,13 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
         except json.JSONDecodeError:
             logging.warning("Unable to read SQL query results: %s", sql_query_results)
             sql_query_results = {}
-            sub_message_results = {}
-        else:
-            # Only load sub-message results if we have a database result
-            sub_message_results = self.parse_message_content(messages[1].content)
-            logging.info("Sub-message Results: %s", sub_message_results)
 
         try:
-            decomposed_messages = [
-                sub_message
-                for sub_message_group in sub_message_results.get(
-                    "decomposed_messages", []
-                )
-                for sub_message in sub_message_group
-            ]
+            decomposed_user_messages = self.extract_decomposed_user_messages(messages)
 
             logging.info("SQL Query Results: %s", sql_query_results)
             payload = AnswerWithSourcesPayload(
-                answer=answer, decomposed_messages=decomposed_messages
+                answer=answer, decomposed_user_messages=decomposed_user_messages
             )
 
             if not isinstance(sql_query_results, dict):
@@ -213,7 +213,7 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
 
     async def process_message(
         self,
-        message_payload: UserInputPayload,
+        message_payload: UserMessagePayload,
         chat_history: list[InteractionPayload] = None,
     ) -> AsyncGenerator[InteractionPayload, None]:
         """Process the complete message through the unified system.
@@ -228,22 +228,22 @@ async def process_message(
         -------
             dict: The response from the system.
         """
-        logging.info("Processing message: %s", message_payload.body.message)
+        logging.info("Processing message: %s", message_payload.body.user_message)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "message": message_payload.body.message,
+            "message": message_payload.body.user_message,
             "chat_history": {},
             "injected_parameters": message_payload.body.injected_parameters,
         }
 
         if chat_history is not None:
             # Update input
             for idx, chat in enumerate(chat_history):
-                if chat.root.payload_type == PayloadType.message:
+                if chat.root.payload_type == PayloadType.USER_MESSAGE:
                     # For now only consider the user query
                     chat_history_key = f"chat_{idx}"
-                    agent_input[chat_history_key] = chat.root.body.message
+                    agent_input[chat_history_key] = chat.root.body.user_message
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
             logging.debug("Message: %s", message)

text_2_sql/autogen/src/autogen_text_2_sql/evaluation_utils.py

@@ -1,3 +1,5 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
 import re
 from typing import Optional, List, Dict, Any
 

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -169,9 +169,9 @@ def agentic_flow(self):
         )
         return flow
 
-    def process_question(
+    def process_user_message(
         self,
-        question: str,
+        user_message: str,
         injected_parameters: dict = None,
     ):
         """Process the complete question through the unified system.
@@ -185,15 +185,14 @@ def process_question(
         -------
             dict: The response from the system.
         """
-        logging.info("Processing question: %s", question)
+        logging.info("Processing question: %s", user_message)
 
         # Update environment with injected parameters
         self._update_environment(injected_parameters)
 
         try:
             agent_input = {
-                "question": question,
-                "chat_history": {},
+                "user_message": user_message,
                 "injected_parameters": injected_parameters,
             }
 

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -17,7 +17,7 @@ class PayloadType(StrEnum):
     ANSWER_WITH_SOURCES = "answer_with_sources"
     DISAMBIGUATION_REQUESTS = "disambiguation_requests"
     PROCESSING_UPDATE = "processing_update"
-    USER_INPUT = "user_input"
+    USER_MESSAGE = "user_message"
 
 
 class InteractionPayloadBase(BaseModel):
@@ -51,6 +51,9 @@ class DismabiguationRequest(InteractionPayloadBase):
         disambiguation_requests: list[DismabiguationRequest] = Field(
             alias="disambiguationRequests"
         )
+        decomposed_user_messages: list[list[str]] = Field(
+            default_factory=list, alias="decomposedUserMessages"
+        )
 
     payload_type: Literal[PayloadType.DISAMBIGUATION_REQUESTS] = Field(
         PayloadType.DISAMBIGUATION_REQUESTS, alias="payloadType"
@@ -73,7 +76,7 @@ class Source(InteractionPayloadBase):
             sql_rows: list[dict] = Field(default_factory=list, alias="sqlRows")
 
         answer: str
-        decomposed_user_messages: list[str] = Field(
+        decomposed_user_messages: list[list[str]] = Field(
             default_factory=list, alias="decomposedUserMessages"
         )
         sources: list[Source] = Field(default_factory=list)
@@ -111,7 +114,7 @@ def __init__(self, **kwargs):
         self.body = self.Body(**kwargs)
 
 
-class UserInputPayload(InteractionPayloadBase):
+class UserMessagePayload(InteractionPayloadBase):
     class Body(InteractionPayloadBase):
         user_message: str = Field(..., alias="userMessage")
         injected_parameters: dict = Field(
@@ -130,8 +133,8 @@ def add_defaults(cls, values):
             values["injected_parameters"] = {**defaults, **injected}
             return values
 
-    payload_type: Literal[PayloadType.USER_INPUT] = Field(
-        PayloadType.USER_INPUT, alias="payloadType"
+    payload_type: Literal[PayloadType.USER_MESSAGE] = Field(
+        PayloadType.USER_MESSAGE, alias="payloadType"
     )
     payload_source: Literal[PayloadSource.USER] = Field(
         PayloadSource.USER, alias="payloadSource"
@@ -145,6 +148,6 @@ def __init__(self, **kwargs):
 
 
 class InteractionPayload(RootModel):
-    root: UserInputPayload | ProcessingUpdatePayload | DismabiguationRequestsPayload | AnswerWithSourcesPayload = Field(
+    root: UserMessagePayload | ProcessingUpdatePayload | DismabiguationRequestsPayload | AnswerWithSourcesPayload = Field(
         ..., discriminator="payload_type"
     )