Update payloads

Author: BenConstable9

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

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 (
-    QuestionPayload,
+    UserInputPayload,
     AnswerWithSourcesPayload,
     DismabiguationRequestsPayload,
     ProcessingUpdatePayload,
@@ -211,7 +211,7 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
 
     async def process_question(
         self,
-        question_payload: QuestionPayload,
+        question_payload: UserInputPayload,
         chat_history: list[InteractionPayload] = None,
     ) -> AsyncGenerator[InteractionPayload, None]:
         """Process the complete question through the unified system.
@@ -238,7 +238,7 @@ async def process_question(
         if chat_history is not None:
             # Update input
             for idx, chat in enumerate(chat_history):
-                if chat.root.payload_type == PayloadType.QUESTION:
+                if chat.root.payload_type == PayloadType.USER_INPUT:
                     # For now only consider the user query
                     chat_history_key = f"chat_{idx}"
                     agent_input[chat_history_key] = chat.root.body.question

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

@@ -29,10 +29,11 @@ class PayloadType(StrEnum):
     ANSWER_WITH_SOURCES = "answer_with_sources"
     DISAMBIGUATION_REQUEST = "disambiguation_request"
     PROCESSING_UPDATE = "processing_update"
-    QUESTION = "question"
+    USER_INPUT = "user_input"
 
 
 class ColumnFilterPair(BaseModel):
+    fqn: str
     column: str
     filter_value: str | None = Field(default=None)
 
@@ -57,6 +58,32 @@ def __init__(self, **kwargs):
         self.body = self.Body(**kwargs)
 
 
+request = DismabiguationRequestsPayload(
+    disambiguation_requests=[
+        {
+            "question": "Which of the following do you mean?",
+            "choices": [
+                {"fqn": "<fqn>", "column": "product_name", "filter_value": "Road Bike"},
+                {
+                    "fqn": "<fqn>",
+                    "column": "product_name",
+                    "filter_value": "Mountain Bike",
+                },
+            ],
+        },
+        {
+            "question": "Do you mean total sales by volume or number of customers?",
+            "choices": [
+                {"fqn": "<fqn>", "column": "sales_volume", "filter_value": None},
+                {"fqn": "<fqn>", "column": "customer_count", "filter_value": None},
+            ],
+        },
+    ]
+)
+
+print(request.model_dump())
+
+
 class AnswerWithSourcesPayload(PayloadBase):
     class Body(BaseModel):
         class Source(BaseModel):
@@ -94,7 +121,7 @@ def __init__(self, **kwargs):
         self.body = self.Body(**kwargs)
 
 
-class QuestionPayload(PayloadBase):
+class UserInputPayload(PayloadBase):
     class Body(BaseModel):
         question: str
         injected_parameters: dict = Field(default_factory=dict)
@@ -111,7 +138,7 @@ def add_defaults(cls, values):
             values["injected_parameters"] = {**defaults, **injected}
             return values
 
-    payload_type: Literal[PayloadType.QUESTION] = PayloadType.QUESTION
+    payload_type: Literal[PayloadType.USER_INPUT] = PayloadType.USER_INPUT
     payload_source: Literal[PayloadSource.USER] = PayloadSource.USER
     body: Body | None = Field(default=None)
 
@@ -122,6 +149,6 @@ def __init__(self, **kwargs):
 
 
 class InteractionPayload(RootModel):
-    root: QuestionPayload | ProcessingUpdatePayload | DismabiguationRequestsPayload | AnswerWithSourcesPayload = Field(
+    root: UserInputPayload | ProcessingUpdatePayload | DismabiguationRequestsPayload | AnswerWithSourcesPayload = Field(
         ..., discriminator="payload_type"
     )