Update sql

Author: BenConstable9

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

@@ -101,19 +101,29 @@ async def on_messages_stream(
                 if schema not in final_schemas:
                     final_schemas.append(schema)
 
-        final_columns = []
-        for column_value_result in column_value_results:
+        columns_for_filter = {}
+        values_for_filter = {}
+        for filter, column_value_result in zip(
+            loaded_entity_result["filter_conditions"], column_value_results
+        ):
+            columns_for_filter[filter] = []
+            values_for_filter[filter] = []
             for column in column_value_result:
-                if column not in final_columns:
-                    final_columns.append(column)
+                if column["Column"] not in columns_for_filter[filter]:
+                    columns_for_filter[filter].append(column["Column"])
 
-        all_column_lengths = [len(column) for column in final_columns]
+                if column["Value"] not in values_for_filter[filter]:
+                    values_for_filter[filter].append(column["Value"])
+
+        num_all_values = [len(filter) for filter in values_for_filter]
+        num_all_columns = [len(filter) for filter in columns_for_filter]
 
         final_results = {
-            "MANDATORY_DISAMBIGUATION": max(all_column_lengths) > 3
-            or len(final_columns) > 3,
-            "schemas": final_schemas,
-            "column_values": final_columns,
+            "MANDATORY_DISAMBIGUATION": max(num_all_values) > 3
+            or max(num_all_columns) > 3,
+            "COLUMN_OPTIONS_FOR_FILTERS": columns_for_filter,
+            "VALUE_OPTIONS_FOR_FILTERS": values_for_filter,
+            "SCHEMA_OPTIONS": final_schemas,
         }
 
         logging.info(f"Final results: {final_results}")