Update data contract

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

text_2_sql/autogen/src/autogen_text_2_sql/__init__.py

@@ -1,3 +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_response import AgentRequestBody
 
-__all__ = ["AutoGenText2Sql"]
+__all__ = ["AutoGenText2Sql", "AgentRequestBody"]

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -18,11 +18,17 @@
 import os
 from datetime import datetime
 
-from text_2_sql_core.payloads import (
+from text_2_sql_core.payloads.agent_response import (
+    AgentResponse,
+    AgentRequestBody,
     AnswerWithSources,
-    UserInformationRequest,
+    Source,
+    DismabiguationRequests,
+)
+from text_2_sql_core.payloads.chat_history import ChatHistoryItem
+from text_2_sql_core.payloads.processing_update import (
+    ProcessingUpdateBody,
     ProcessingUpdate,
-    ChatHistoryItem,
 )
 from autogen_agentchat.base import Response, TaskResult
 from typing import AsyncGenerator
@@ -123,6 +129,16 @@ def agentic_flow(self):
         )
         return flow
 
+    def extract_disambiguation_request(self, messages: list) -> DismabiguationRequests:
+        """Extract the disambiguation request from the answer."""
+
+        disambiguation_request = messages[-1].content
+
+        # TODO: Properly extract the disambiguation request
+        return DismabiguationRequests(
+            disambiguation_request=disambiguation_request,
+        )
+
     def extract_sources(self, messages: list) -> AnswerWithSources:
         """Extract the sources from the answer."""
 
@@ -147,10 +163,10 @@ def extract_sources(self, messages: list) -> AnswerWithSources:
                 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"],
-                        }
+                        Source(
+                            sql_query=sql_query_result["sql_query"],
+                            sql_rows=sql_query_result["sql_rows"],
+                        )
                     )
 
         except json.JSONDecodeError:
@@ -164,10 +180,9 @@ def extract_sources(self, messages: list) -> AnswerWithSources:
 
     async def process_question(
         self,
-        question: str,
+        request: AgentRequestBody,
         chat_history: list[ChatHistoryItem] = None,
-        injected_parameters: dict = None,
-    ) -> AsyncGenerator[AnswerWithSources | UserInformationRequest, None]:
+    ) -> AsyncGenerator[AgentResponse | ProcessingUpdate, None]:
         """Process the complete question through the unified system.
 
         Args:
@@ -180,58 +195,63 @@ async def process_question(
         -------
             dict: The response from the system.
         """
-        logging.info("Processing question: %s", question)
+        logging.info("Processing question: %s", request.question)
         logging.info("Chat history: %s", chat_history)
 
         agent_input = {
-            "question": question,
+            "question": request.question,
             "chat_history": {},
-            "injected_parameters": injected_parameters,
+            "injected_parameters": request.injected_parameters,
         }
 
         if chat_history is not None:
             # Update input
             for idx, chat in enumerate(chat_history):
                 # For now only consider the user query
-                agent_input[f"chat_{idx}"] = chat.user_query
+                agent_input[f"chat_{idx}"] = chat.request.question
 
         async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
             logging.debug("Message: %s", message)
 
             payload = None
 
             if isinstance(message, TextMessage):
+                processing_update = None
                 if message.source == "query_rewrite_agent":
-                    # If the message is from the query_rewrite_agent, we need to update the chat history
-                    payload = ProcessingUpdate(
+                    processing_update = ProcessingUpdateBody(
                         message="Rewriting the query...",
                     )
                 elif message.source == "parallel_query_solving_agent":
-                    # If the message is from the parallel_query_solving_agent, we need to update the chat history
-                    payload = ProcessingUpdate(
+                    processing_update = ProcessingUpdateBody(
                         message="Solving the query...",
                     )
                 elif message.source == "answer_agent":
-                    # If the message is from the answer_agent, we need to update the chat history
-                    payload = ProcessingUpdate(
+                    processing_update = ProcessingUpdateBody(
                         message="Generating the answer...",
                     )
 
+                if processing_update is not None:
+                    payload = ProcessingUpdate(
+                        processing_update=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
-                    payload = self.extract_sources(message.messages)
+                    response = self.extract_sources(message.messages)
                 elif message.messages[-1].source == "parallel_query_solving_agent":
-                    payload = UserInformationRequest(
-                        **json.loads(message.messages[-1].content),
-                    )
+                    # Load into disambiguation request
+                    response = self.extract_disambiguation_request(message.messages)
                 else:
                     logging.error("Unexpected TaskResult: %s", message)
                     raise ValueError("Unexpected TaskResult")
 
+                payload = AgentResponse(request=request, response=response)
+
             if payload is not None:
                 logging.debug("Payload: %s", payload)
                 yield payload

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

@@ -44,10 +44,10 @@ async def on_messages_stream(
         last_response = messages[-1].content
         parameter_input = messages[0].content
         try:
-            user_parameters = json.loads(parameter_input)["parameters"]
+            injected_parameters = json.loads(parameter_input)["injected_parameters"]
         except json.JSONDecodeError:
             logging.error("Error decoding the user parameters.")
-            user_parameters = {}
+            injected_parameters = {}
 
         # Load the json of the last message to populate the final output object
         query_rewrites = json.loads(last_response)
@@ -117,7 +117,7 @@ async def consume_inner_messages_from_agentic_flow(
             inner_solving_generators.append(
                 consume_inner_messages_from_agentic_flow(
                     inner_autogen_text_2_sql.process_question(
-                        question=query_rewrite, parameters=user_parameters
+                        question=query_rewrite, injected_parameters=injected_parameters
                     ),
                     query_rewrite,
                     database_results,

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/__init__.py

@@ -1,12 +0,0 @@
-from text_2_sql_core.payloads.answer_with_sources import AnswerWithSources, Source
-from text_2_sql_core.payloads.user_information_request import UserInformationRequest
-from text_2_sql_core.payloads.processing_update import ProcessingUpdate
-from text_2_sql_core.payloads.chat_history import ChatHistoryItem
-
-__all__ = [
-    "AnswerWithSources",
-    "Source",
-    "UserInformationRequest",
-    "ProcessingUpdate",
-    "ChatHistoryItem",
-]

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/agent_response.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pydantic import BaseModel, RootModel, Field
+from enum import StrEnum
+
+from typing import Literal
+from datetime import datetime, timezone
+
+
+class AgentResponseHeader(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    timestamp: datetime = Field(
+        ...,
+        description="Timestamp in UTC",
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+
+
+class AgentResponseType(StrEnum):
+    ANSWER_WITH_SOURCES = "answer_with_sources"
+    DISAMBIGUATION = "disambiguation"
+
+
+class DismabiguationRequest(BaseModel):
+    question: str
+    matching_columns: list[str]
+    matching_filter_values: list[str]
+    other_user_choices: list[str]
+
+
+class DismabiguationRequests(BaseModel):
+    response_type: Literal[AgentResponseType.DISAMBIGUATION] = Field(
+        default=AgentResponseType.DISAMBIGUATION
+    )
+    requests: list[DismabiguationRequest]
+
+
+class Source(BaseModel):
+    sql_query: str
+    sql_rows: list[dict]
+
+
+class AnswerWithSources(BaseModel):
+    response_type: Literal[AgentResponseType.ANSWER_WITH_SOURCES] = Field(
+        default=AgentResponseType.ANSWER_WITH_SOURCES
+    )
+    answer: str
+    sources: list[Source] = Field(default_factory=list)
+
+
+class AgentResponseBody(RootModel):
+    root: DismabiguationRequests | AnswerWithSources = Field(
+        ..., discriminator="response_type"
+    )
+
+
+class AgentRequestBody(BaseModel):
+    question: str
+    injected_parameters: dict = Field(default_factory=dict)
+
+
+class AgentResponse(BaseModel):
+    header: AgentResponseHeader | None = Field(default=None)
+    request: AgentRequestBody
+    response: AgentResponseBody

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

@@ -1,9 +1,16 @@
-from pydantic import BaseModel
-from text_2_sql_core.payloads.answer_with_sources import AnswerWithSources
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from pydantic import BaseModel, Field
+from text_2_sql_core.payloads.agent_response import AgentResponse
+from datetime import datetime, timezone
 
 
 class ChatHistoryItem(BaseModel):
     """Chat history item with user message and agent response."""
 
-    user_query: str
-    agent_response: AnswerWithSources
+    timestamp: datetime = Field(
+        ...,
+        description="Timestamp in UTC",
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+    agent_response: AgentResponse

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/chat_history.py

@@ -1,6 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
 from pydantic import BaseModel, Field
+from datetime import datetime, timezone
 
 
-class ProcessingUpdate(BaseModel):
+class ProcessingUpdateHeader(BaseModel):
+    timestamp: datetime = Field(
+        ...,
+        description="Timestamp in UTC",
+        default_factory=lambda: datetime.now(timezone.utc),
+    )
+
+
+class ProcessingUpdateBody(BaseModel):
     title: str | None = Field(default="Processing...")
     message: str | None = Field(default="Processing...")
+
+
+class ProcessingUpdate(BaseModel):
+    header: ProcessingUpdateHeader | None = Field(
+        default_factory=ProcessingUpdateHeader
+    )
+    processing_update: ProcessingUpdateBody