Update query validation

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -16,11 +16,18 @@
 )
 from autogen_agentchat.agents import UserProxyAgent
 from autogen_agentchat.messages import TextMessage
-from autogen_agentchat.base import Response
 import json
 import os
 from datetime import datetime
 
+from text_2_sql_core.payloads import (
+    AnswerWithSources,
+    UserInformationRequest,
+    ProcessingUpdate,
+)
+from autogen_agentchat.base import Response, TaskResult
+from asyncio import AsyncGenerator
+
 
 class EmptyResponseUserProxyAgent(UserProxyAgent):
     """UserProxyAgent that automatically responds with empty messages."""
@@ -120,12 +127,12 @@ def agentic_flow(self):
         )
         return flow
 
-    def process_question(
+    async def process_question(
         self,
         question: str,
         chat_history: list[str] = None,
         parameters: dict = None,
-    ):
+    ) -> AsyncGenerator[AnswerWithSources | UserInformationRequest]:
         """Process the complete question through the unified system.
 
         Args:
@@ -152,4 +159,42 @@ def process_question(
             for idx, chat in enumerate(chat_history):
                 agent_input[f"chat_{idx}"] = chat
 
-        return self.agentic_flow.run_stream(task=json.dumps(agent_input))
+        async for message in self.agentic_flow.run_stream(task=json.dumps(agent_input)):
+            logging.info("Message: %s", message)
+            logging.info("Message type: %s", type(message))
+
+            payload = None
+
+            if isinstance(message, TextMessage):
+                if message.source == "query_rewrite_agent":
+                    # If the message is from the query_rewrite_agent, we need to update the chat history
+                    payload = ProcessingUpdate(
+                        title="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(
+                        title="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(
+                        title="Generating the answer...",
+                    )
+
+            elif isinstance(message, TaskResult):
+                # Now we need to return the final answer or the disambiguation request
+
+                if message.task == "answer_agent":
+                    # If the message is from the answer_agent, we need to return the final answer
+                    payload = AnswerWithSources(
+                        **json.loads(message.content),
+                    )
+                else:
+                    payload = UserInformationRequest(
+                        **json.loads(message.content),
+                    )
+
+            if payload is not None:
+                logging.info("Payload: %s", payload)
+                yield payload

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sql.py

@@ -143,7 +143,7 @@ async def query_validation(
             sqlglot.transpile(
                 sql_query,
                 read=self.database_engine.value.lower(),
-                error_level=sqlglot.ErrorLevel.ERROR,
+                error_level=sqlglot.ErrorLevel.RAISE,
             )
         except sqlglot.errors.ParseError as e:
             logging.error("SQL Query is invalid: %s", e.errors)

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

@@ -1,4 +1,5 @@
 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
 
-__all__ = ["AnswerWithSources", "Source", "UserInformationRequest"]
+__all__ = ["AnswerWithSources", "Source", "UserInformationRequest", "ProcessingUpdate"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/processing_update.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class ProcessingUpdate(BaseModel):
+    title: str
+    message: str

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/user_information_request.py

@@ -9,18 +9,18 @@ class RequestType(StrEnum):
 
 
 class ClarificationRequest(BaseModel):
-    type: Literal[RequestType.CLARIFICATION]
+    request_type: Literal[RequestType.CLARIFICATION]
     question: str
     other_user_choices: list[str]
 
 
 class DismabiguationRequest(BaseModel):
-    type: Literal[RequestType.DISAMBIGUATION]
+    request_type: Literal[RequestType.DISAMBIGUATION]
     question: str
     matching_columns: list[str]
     matching_filter_values: list[str]
     other_user_choices: list[str]
 
 
 class UserInformationRequest(RootModel):
-    root: DismabiguationRequest = Field(..., discriminator="type")
+    root: DismabiguationRequest = Field(..., discriminator="request_type")