Fix disambiguation

Author: BenConstable9

text_2_sql/autogen/src/autogen_text_2_sql/autogen_text_2_sql.py

@@ -62,7 +62,7 @@ def termination_condition(self):
         termination = (
             TextMentionTermination("TERMINATE")
             | SourceMatchTermination("answer_agent")
-            | TextMentionTermination("requires_user_information_request")
+            | TextMentionTermination("contains_disambiguation_requests")
             | MaxMessageTermination(5)
         )
         return termination
@@ -131,20 +131,43 @@ def extract_decomposed_user_messages(self, messages: list) -> list[list[str]]:
         sub_message_results = self.parse_message_content(messages[1].content)
         logging.info("Decomposed Results: %s", sub_message_results)
 
-        return sub_message_results.get("decomposed_user_messages", [])
+        decomposed_user_messages = sub_message_results.get(
+            "decomposed_user_messages", []
+        )
+
+        logging.debug(
+            "Returning decomposed_user_messages: %s", decomposed_user_messages
+        )
+
+        return decomposed_user_messages
 
     def extract_disambiguation_request(
         self, messages: list
     ) -> DismabiguationRequestsPayload:
         """Extract the disambiguation request from the answer."""
-        disambiguation_request = messages[-1].content
+        all_disambiguation_requests = self.parse_message_content(messages[-1].content)
 
         decomposed_user_messages = self.extract_decomposed_user_messages(messages)
-        return DismabiguationRequestsPayload(
-            disambiguation_request=disambiguation_request,
-            decomposed_user_messages=decomposed_user_messages,
+        request_payload = DismabiguationRequestsPayload(
+            decomposed_user_messages=decomposed_user_messages
         )
 
+        for per_question_disambiguation_request in all_disambiguation_requests[
+            "disambiguation_requests"
+        ].values():
+            for disambiguation_request in per_question_disambiguation_request:
+                logging.info(
+                    "Disambiguation Request Identified: %s", disambiguation_request
+                )
+
+                request = DismabiguationRequestsPayload.Body.DismabiguationRequest(
+                    agent_question=disambiguation_request["agent_question"],
+                    user_choices=disambiguation_request["user_choices"],
+                )
+                request_payload.body.disambiguation_requests.append(request)
+
+        return request_payload
+
     def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
         """Extract the sources from the answer."""
         answer = messages[-1].content
@@ -169,11 +192,13 @@ def extract_answer_payload(self, messages: list) -> AnswerWithSourcesPayload:
                 logging.error(f"Expected dict, got {type(sql_query_results)}")
                 return payload
 
-            if "results" not in sql_query_results:
+            if "database_results" not in sql_query_results:
                 logging.error("No 'results' key in sql_query_results")
                 return payload
 
-            for message, sql_query_result_list in sql_query_results["results"].items():
+            for message, sql_query_result_list in sql_query_results[
+                "database_results"
+            ].items():
                 if not sql_query_result_list:  # Check if list is empty
                     logging.warning(f"No results for message: {message}")
                     continue

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

@@ -18,6 +18,18 @@
 from json import JSONDecodeError
 import re
 import os
+from pydantic import BaseModel, Field
+
+
+class FilteredParallelMessagesCollection(BaseModel):
+    database_results: dict[str, list] = Field(default_factory=dict)
+    disambiguation_requests: dict[str, list] = Field(default_factory=dict)
+
+    def add_identifier(self, identifier):
+        if identifier not in self.database_results:
+            self.database_results[identifier] = []
+        if identifier not in self.disambiguation_requests:
+            self.disambiguation_requests[identifier] = []
 
 
 class ParallelQuerySolvingAgent(BaseChatAgent):
@@ -89,7 +101,7 @@ async def on_messages_stream(
         logging.info(f"Query Rewrites: {message_rewrites}")
 
         async def consume_inner_messages_from_agentic_flow(
-            agentic_flow, identifier, database_results
+            agentic_flow, identifier, filtered_parallel_messages
         ):
             """
             Consume the inner messages and append them to the specified list.
@@ -101,8 +113,7 @@ async def consume_inner_messages_from_agentic_flow(
             """
             async for inner_message in agentic_flow:
                 # Add message to results dictionary, tagged by the function name
-                if identifier not in database_results:
-                    database_results[identifier] = []
+                filtered_parallel_messages.add_identifier(identifier)
 
                 logging.info(f"Checking Inner Message: {inner_message}")
 
@@ -122,7 +133,9 @@ async def consume_inner_messages_from_agentic_flow(
                                     == "query_execution_with_limit"
                                 ):
                                     logging.info("Contains query results")
-                                    database_results[identifier].append(
+                                    filtered_parallel_messages.database_results[
+                                        identifier
+                                    ].append(
                                         {
                                             "sql_query": parsed_message[
                                                 "sql_query"
@@ -138,29 +151,43 @@ async def consume_inner_messages_from_agentic_flow(
 
                         # Search for specific message types and add them to the final output object
                         if isinstance(parsed_message, dict):
+                            # Check if the message contains pre-run results
                             if ("contains_pre_run_results" in parsed_message) and (
                                 parsed_message["contains_pre_run_results"] is True
                             ):
                                 logging.info("Contains pre-run results")
                                 for pre_run_sql_query, pre_run_result in parsed_message[
                                     "cached_messages_and_schemas"
                                 ].items():
-                                    database_results[identifier].append(
+                                    filtered_parallel_messages.database_results[
+                                        identifier
+                                    ].append(
                                         {
                                             "sql_query": pre_run_sql_query.replace(
                                                 "\n", " "
                                             ),
                                             "sql_rows": pre_run_result["sql_rows"],
                                         }
                                     )
+                            # Check if disambiguation is required
+                            elif ("disambiguation_requests" in parsed_message) and (
+                                parsed_message["disambiguation_requests"]
+                            ):
+                                logging.info("Contains disambiguation requests")
+                                for disambiguation_request in parsed_message[
+                                    "disambiguation_requests"
+                                ]:
+                                    filtered_parallel_messages.disambiguation_requests[
+                                        identifier
+                                    ].append(disambiguation_request)
 
                 except Exception as e:
                     logging.warning(f"Error processing message: {e}")
 
                 yield inner_message
 
         inner_solving_generators = []
-        database_results = {}
+        filtered_parallel_messages = FilteredParallelMessagesCollection()
 
         # Convert all_non_database_query to lowercase string and compare
         all_non_database_query = str(
@@ -201,7 +228,7 @@ async def consume_inner_messages_from_agentic_flow(
                         injected_parameters=query_params,
                     ),
                     identifier,
-                    database_results,
+                    filtered_parallel_messages,
                 )
             )
 
@@ -218,17 +245,43 @@ async def consume_inner_messages_from_agentic_flow(
                     yield inner_message
 
         # Log final results for debugging or auditing
-        logging.info(f"Database Results: {database_results}")
+        logging.info(
+            "Database Results: %s", filtered_parallel_messages.database_results
+        )
+        logging.info(
+            "Disambiguation Requests: %s",
+            filtered_parallel_messages.disambiguation_requests,
+        )
 
-        # Final response
-        yield Response(
-            chat_message=TextMessage(
-                content=json.dumps(
-                    {"contains_results": True, "results": database_results}
+        if (
+            max(map(len, filtered_parallel_messages.disambiguation_requests.values()))
+            > 0
+        ):
+            # Final response
+            yield Response(
+                chat_message=TextMessage(
+                    content=json.dumps(
+                        {
+                            "contains_disambiguation_requests": True,
+                            "disambiguation_requests": filtered_parallel_messages.disambiguation_requests,
+                        }
+                    ),
+                    source=self.name,
                 ),
-                source=self.name,
-            ),
-        )
+            )
+        else:
+            # Final response
+            yield Response(
+                chat_message=TextMessage(
+                    content=json.dumps(
+                        {
+                            "contains_database_results": True,
+                            "database_results": filtered_parallel_messages.database_results,
+                        }
+                    ),
+                    source=self.name,
+                ),
+            )
 
     async def on_reset(self, cancellation_token: CancellationToken) -> None:
         pass

text_2_sql/autogen/src/autogen_text_2_sql/evaluation_utils.py

@@ -16,7 +16,7 @@ def extract_sql_queries_from_results(results: Dict[str, Any]) -> List[str]:
     """
     queries = []
 
-    if results.get("contains_results") and results.get("results"):
+    if results.get("contains_database_results") and results.get("results"):
         for question_results in results["results"].values():
             for result in question_results:
                 if isinstance(result, dict) and "sql_query" in result:

text_2_sql/autogen/src/autogen_text_2_sql/inner_autogen_text_2_sql.py

@@ -124,7 +124,11 @@ def get_all_agents(self):
     @property
     def termination_condition(self):
         """Define the termination condition for the chat."""
-        termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(10)
+        termination = (
+            TextMentionTermination("TERMINATE")
+            | MaxMessageTermination(10)
+            | TextMentionTermination("disambiguation_request")
+        )
         return termination
 
     def unified_selector(self, messages):

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

@@ -187,15 +187,17 @@ async def query_execution_with_limit(
                     "type": "query_execution_with_limit",
                     "sql_query": sql_query,
                     "sql_rows": result,
-                }
+                },
+                default=str,
             )
         else:
             return json.dumps(
                 {
                     "type": "errored_query_execution_with_limit",
                     "sql_query": sql_query,
                     "errors": validation_result,
-                }
+                },
+                default=str,
             )
 
     async def query_validation(

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

@@ -42,8 +42,8 @@ class DismabiguationRequest(InteractionPayloadBase):
             agent_question: str | None = Field(..., alias="agentQuestion")
             user_choices: list[str] | None = Field(default=None, alias="userChoices")
 
-        disambiguation_requests: list[DismabiguationRequest] = Field(
-            alias="disambiguationRequests"
+        disambiguation_requests: list[DismabiguationRequest] | None = Field(
+            default_factory=list, alias="disambiguationRequests"
         )
         decomposed_user_messages: list[list[str]] = Field(
             default_factory=list, alias="decomposedUserMessages"

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/disambiguation_and_sql_query_generation_agent.yaml

@@ -170,12 +170,16 @@ system_message:
 
     If disambiguation needed:
     {
-      \"disambiguation\": [{
-        \"question\": \"<specific_question>\",
-        \"matching_columns\": [\"<column1>\", \"<column2>\"],
-        \"matching_filter_values\": [\"<value1>\", \"<value2>\"],
-        \"other_user_choices\": [\"<choice1>\", \"<choice2>\"]
-      }]
+      \"disambiguation_requests\": [
+        {
+          \"agent_question\": \"<specific_question>\",
+          \"user_choices\": [\"<choice1>\", \"<choice2>\"]
+        },
+        {
+          \"agent_question\": \"<specific_question>\",
+          \"user_choices\": [\"<choice1>\", \"<choice2>\"]
+        }
+      ]
     }
     TERMINATE
     </output_format>