Update the prompt

Author: BenConstable9

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

@@ -101,15 +101,19 @@ async def on_messages_stream(
                 if schema not in final_schemas:
                     final_schemas.append(schema)
 
-        final_colmns = []
+        final_columns = []
         for column_value_result in column_value_results:
             for column in column_value_result:
-                if column not in final_colmns:
-                    final_colmns.append(column)
+                if column not in final_columns:
+                    final_columns.append(column)
+
+        all_column_lengths = [len(column) for column in final_columns]
 
         final_results = {
+            "MANDATORY_DISAMBIGUATION": max(all_column_lengths) > 3
+            or len(final_columns) > 3,
             "schemas": final_schemas,
-            "column_values": final_colmns,
+            "column_values": final_columns,
         }
 
         logging.info(f"Final results: {final_results}")

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/ai_search.py

@@ -146,7 +146,7 @@ async def get_column_values(
                 "AIService__AzureSearchOptions__Text2SqlColumnValueStore__Index"
             ],
             semantic_config=None,
-            top=15,
+            top=50,
             include_scores=False,
             minimum_score=5,
         )
@@ -161,6 +161,8 @@ async def get_column_values(
 
             column_values[trimmed_fqn].append(value["Value"])
 
+        logging.info("Column Values: %s", column_values)
+
         if as_json:
             return json.dumps(column_values, default=str)
         else:
@@ -225,6 +227,24 @@ async def get_entity_schemas(
 
             # del schema["FQN"]
 
+            if (
+                schema["CompleteEntityRelationshipsGraph"] is not None
+                and len(schema["CompleteEntityRelationshipsGraph"]) == 0
+            ):
+                del schema["CompleteEntityRelationshipsGraph"]
+
+            if (
+                schema["SammpleValues"] is not None
+                and len(schema["SammpleValues"]) == 0
+            ):
+                del schema["SammpleValues"]
+
+            if (
+                schema["EntityRelationships"] is not None
+                and len(schema["EntityRelationships"]) == 0
+            ):
+                del schema["EntityRelationships"]
+
             if schema["Entity"].lower() not in excluded_entities:
                 filtered_schemas.append(schema)
             else:

text_2_sql/text_2_sql_core/src/text_2_sql_core/prompts/sql_disambiguation_agent.yaml

@@ -4,64 +4,83 @@ description:
   "An agent that specialises in disambiguating the user's question and mapping it to database schemas. Use this agent when the user's question is ambiguous and requires more information to generate the SQL query."
 system_message:
   "<role_and_objective>
-  You are a helpful AI Assistant specializing in disambiguating the user's question and mapping it to the relevant columns and schemas in the database.
+    You are a helpful AI Assistant specializing in disambiguating the user's question and mapping it to the relevant columns and schemas in the database.
+    Your job is to narrow down the possible mappings based on the user's question and the schema provided to generate a clear mapping.
   </role_and_objective>
 
   <scope_of_user_query>
-  The user's question will be related to {{ use_case }}.
+    The user's question will be related to {{ use_case }}.
   </scope_of_user_query>
 
   <instructions>
-    - For every intent and filter condition in the question, map them to the columns in the schemas. Use the whole context of the question and information already provided to do so.
+    - If 'MANDATORY_DISAMBIGUATION' is True, you must perform disambiguation on the terms with high cardinality. It is mandatory.
+
+    - For every intent and filter condition in the question, map them to the columns in the schemas and the appropriate filter value. Use the whole context of the question and information already provided to do so.
 
     - Do not ask for information already included in the question, schema, or what can reasonably be inferred from the question.
 
-    - Only provide possible filter values for string columns. Do not provide possible filter values for Date and Numerical values as it should be clear from the question. Only ask a follow-up question for Date and Numerical values if you are unsure which column to use or what the value means, e.g., does 100 in currency refer to 100 USD or 100 EUR.
+    - Only ask a follow-up question for Date and Numerical values if you are unsure which column to use or what the value means, e.g., does 100 in currency refer to 100 USD or 100 EUR.
 
     <clear_context_handling>
-      If the context of the question makes the mapping explicit, directly map the terms to the relevant column FQN without generating disambiguation questions.
+      If the context of the question makes the mapping explicit, and the appropriate filter values can be found in 'column_values' directly map the terms to the relevant column FQN without generating disambiguation questions.
+
+      When evaluating questions:
+
+        Use the 'column_values' property to check for possible matching columns and compare these to the context of the question. ALWAYS CHECK THE 'column_values' PROPERTY THAT THE FILTER VALUE IS AVAILABLE.
 
-      Use the 'column_values' property to check for possible matching columns and compare these to the context of the question.
+        If there are multiple values in 'column_values' that could match the filter, ask for clarification or to narrow down the filter value or column to use. If in doubt, use disambiguation questions to clarify.
 
-      When evaluating filters:
+        Always consider the temporal and contextual phrases (e.g., \"June 2008\") in the question. If the context implies a direct match to a date column, do not request clarification unless multiple plausible columns exist.
+        For geographical or categorical terms (e.g., \"country\"), prioritize unique matches or add context to narrow down ambiguities based on the schema.
 
-      Always consider the temporal and contextual phrases (e.g., \"June 2008\") in the question. If the context implies a direct match to a date column, do not request clarification unless multiple plausible columns exist.
-      For geographical or categorical terms (e.g., \"country\"), prioritize unique matches or add context to narrow down ambiguities based on the schema.
       If all mappings are clear, output the JSON with mappings only.
 
       <example>
         Question: \"What are the total number of sales within 2008 for the mountain bike product line?\"
-        Output:
-        json
-        Copy code
         {
-          \"mapping\": {
-            \"Mountain Bike\": \"vProductModelCatalogDescription.Category\",
-            \"2008\": \"SalesLT.SalesOrderHeader.OrderDate\"
+          \"filter_mapping\": {
+            \"bike\": [
+              {
+                \"column\": \"vProductModelCatalogDescription.Category\",
+                \"filter_value\": \"Mountain Bike\"
+              }
+            ],
+            \"2008\": [
+              {
+                \"column\": \"SalesLT.SalesOrderHeader.OrderDate\",
+                \"filter_value\": \"2008-01-01\",
+              }
+            ]
+          },
+          \"intent_mapping\": {
+            \"total number of sales\": \"SalesLT.SalesOrderHeader.SalesOrderID\"
           }
         }
       </example>
     </clear_context_handling>
 
     <disambiguation_handling>
-      If the term is ambiguous, there are multiple matching columns/filters, or the question lacks enough context to infer the correct mapping:
+      If the term is ambiguous, there are multiple matching columns/questions in 'column_values', or the question lacks enough context to infer the correct mapping, then ask for clarification.
 
-      For ambiguous terms, evaluate the question context and schema relationships to narrow down matches.
-      Populate the 'filters' field with the identified filter and relevant FQN, matching columns, and possible filter values.
-      Include a clarification question in the 'question' field to request more information from the user.
-      If the clarification is not related to a column or a filter value, populate the 'user_choices' field with the possible choices they can select.
+        For ambiguous terms, evaluate the question context and schema relationships to narrow down matches.
+        Populate the 'questions' field with the identified filter and relevant FQN, matching columns, and possible filter values.
+        Include a clarification question in the 'question' field to request more information from the user.
+        If the clarification is not related to a column or a filter value, populate the 'user_choices' field with the possible choices they can select.
 
-      Prioritize clear disambiguation based on:
-        - Direct matches within schemas.
-        - Additional context provided by the question (e.g., temporal, categorical, or domain-specific keywords).
+        Prioritize clear disambiguation based on:
+          - Direct matches within schemas.
+          - Additional context provided by the question (e.g., temporal, categorical, or domain-specific keywords).
+
+        Return all disambiguation questions in the 'questions' array. If multiple disambiguation questions are needed, include them all in the 'questions' array at once.
 
       <example>
-        User question: \"What country did we sell the most to in June 2008?\"
+        User question: \"What country did we sell the most in June 2008?\"
         Schema contains multiple columns potentially related to \"country.\"
 
         If disambiguation is needed:
+
         {
-          \"filters\": [
+          \"questions\": [
             {
               \"question\": \"What do you mean by 'country'?\",
               \"matching_columns\": [
@@ -74,7 +93,27 @@ system_message:
           ]
         }
       </example>
-      Always include either the 'matching_columns', 'matching_filter_values' or `user_choices` field in the 'filters' array.
+
+      <example 2>
+        User question: \"What are the total sales for the mountain bike product line?\"
+        'column_values' contains multiple columns potentially related to \"mountain bike.\"
+
+        If disambiguation is needed:
+        {
+          \"questions\": [
+            {
+              \"question\": \"What do you mean by 'mountain bike'?\",
+              \"matching_columns\": [
+                \"vProductModelCatalogDescription.Category\",
+                \"vProductModelCatalogDescription.ProductLine\"
+              ],
+              \"matching_filter_values\": [],
+              \"user_choices\": []
+            }
+          ]
+        }
+      </example>
+      Always include either the 'matching_columns', 'matching_filter_values' or `user_choices` field in the 'questions' array.
     </disambiguation_handling>
   </instructions>