Add Taxonomy to postgres Schema extraction

api/agents/taxonomy_agent.py

@@ -17,9 +17,15 @@
 clarification in that way he have the relevent information to answer.
 * When you ask the user to confirm a value, please provide the value in your answer.
 * Mention only question about values and dont mention the SQL query or the tables in your answer.
+* Use the explanation to understand what was unclear or ambiguous in the user's request.
+
+** Understand what the SQL does and ask the user if this is what they need with a clarification question. **
 
 Please create the clarification question step by step.
 
+Last agent explanation:
+{EXPLANATION}
+
 Question:
 {QUESTION}
 
@@ -42,11 +48,14 @@ class TaxonomyAgent:
     def __init__(self):
         """Initialize the taxonomy agent."""
 
-    def get_answer(self, question: str, sql: str) -> str:
+    def get_answer(self, question: str, sql: str, explanation: str) -> str:
         """Get taxonomy classification for a question and SQL pair."""
         messages = [
             {
-                "content": TAXONOMY_PROMPT.format(QUESTION=question, SQL=sql),
+                "content": TAXONOMY_PROMPT.format(QUESTION=question,
+                                                  SQL=sql,
+                                                  EXPLANATION=explanation.replace("{", "{{").replace("}", "}}")
+                                                  ),
                 "role": "user",
             }
         ]

api/config.py

@@ -1,7 +1,10 @@
+
 """
 This module contains the configuration for the text2sql module.
 """
 
+
+
 import dataclasses
 from typing import Union
 
@@ -54,6 +57,9 @@ class Config:
     COMPLETION_MODEL = "azure/gpt-4.1"
     VALIDATOR_MODEL = "azure/gpt-4.1"
     TEMPERATURE = 0
+    DB_MAX_ROWS: int = 500
+    DB_MAX_DISTINCT: int = 100
+    DB_UNIQUENESS_THRESHOLD: float = 0.5
     # client = boto3.client('sts')
     # AWS_PROFILE = os.getenv("aws_profile_name")
     # AWS_REGION = os.getenv("aws_region_name")

api/loaders/base_loader.py

@@ -1,7 +1,8 @@
 """Base loader module providing abstract base class for data loaders."""
 
 from abc import ABC
-from typing import Tuple
+from typing import Tuple, List
+from api.config import Config
 
 
 class BaseLoader(ABC):
@@ -14,3 +15,62 @@ def load(_graph_id: str, _data) -> Tuple[bool, str]:
         This method must be implemented by any subclass.
         """
         return False, "Not implemented"
+
+    @staticmethod
+    def extract_distinct_values_for_column(cursor, table_name: str, col_name: str) -> List[str]:
+        """
+        Extract distinct values for a column if it meets the criteria for inclusion.
+
+        Args:
+            cursor: Database cursor
+            table_name: Name of the table
+            col_name: Name of the column
+
+        Returns:
+            List of formatted distinct values to add to description, or empty list
+        """
+        # Count rows and distinct values for the column
+        cursor.execute(
+            f"""
+            SELECT COUNT(*) AS total_count,
+                   COUNT(DISTINCT {col_name}) AS distinct_count
+            FROM {table_name};
+            """
+        )
+        output = cursor.fetchall()
+
+        # Auto-detect result format by checking the first result
+        first_result = output[0]
+        if isinstance(first_result, dict):
+            # MySQL/dictionary-style results
+            rows_count, distinct_count = first_result['total_count'], first_result['distinct_count']
+        else:
+            # PostgreSQL/tuple-style results
+            rows_count, distinct_count = first_result
+
+        max_rows = Config.DB_MAX_ROWS
+        max_distinct = Config.DB_MAX_DISTINCT
+        uniqueness_threshold = Config.DB_UNIQUENESS_THRESHOLD
+
+        if 0 < rows_count < max_rows and distinct_count < max_distinct:
+            uniqueness_value = distinct_count / rows_count
+            if uniqueness_value < uniqueness_threshold:
+                cursor.execute(
+                    f"SELECT DISTINCT {col_name} FROM {table_name};"
+                )
+
+                distinct_results = cursor.fetchall()
+                if distinct_results and isinstance(distinct_results[0], dict):
+                    # MySQL/dictionary-style results
+                    distinct_values = [row[col_name] for row in distinct_results if row[col_name] is not None]
+                else:
+                    # PostgreSQL/tuple-style results
+                    distinct_values = [row[0] for row in distinct_results if row[0] is not None]
+
+                if distinct_values:
+                    # Check first value type to avoid objects like dict/bytes
+                    first_val = distinct_values[0]
+                    if isinstance(first_val, (str, int, float)):
+                        return [f"(Optional values: {', '.join(f'({str(v)})' for v in distinct_values)})"]
+
+        return []

api/loaders/mysql_loader.py

@@ -278,6 +278,12 @@ def extract_columns_info(cursor, db_name: str, table_name: str) -> Dict[str, Any
             if column_default is not None:
                 description_parts.append(f"(Default: {column_default})")
 
+            # Add distinct values if applicable
+            distinct_values_desc = MySQLLoader.extract_distinct_values_for_column(
+                cursor, table_name, col_name
+            )
+            description_parts.extend(distinct_values_desc)
+
             columns_info[col_name] = {
                 'type': data_type,
                 'null': is_nullable,

api/loaders/postgres_loader.py

@@ -1,9 +1,9 @@
 """PostgreSQL loader for loading database schemas into FalkorDB graphs."""
 
+import re
 import datetime
 import decimal
 import logging
-import re
 from typing import Tuple, Dict, Any, List
 
 import psycopg2
@@ -233,6 +233,12 @@ def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
             if column_default:
                 description_parts.append(f"(Default: {column_default})")
 
+            # Add distinct values if applicable
+            distinct_values_desc = PostgresLoader.extract_distinct_values_for_column(
+                cursor, table_name, col_name
+            )
+            description_parts.extend(distinct_values_desc)
+
             columns_info[col_name] = {
                 'type': data_type,
                 'null': is_nullable,
@@ -241,6 +247,7 @@ def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
                 'default': column_default
             }
 
+
         return columns_info
 
     @staticmethod

tests/test_postgres_loader.py

@@ -65,10 +65,19 @@ def test_extract_columns_info(self):
         """Test column information extraction"""
         # Mock cursor with column data
         mock_cursor = Mock()
-        mock_cursor.fetchall.return_value = [
-            ("id", "integer", "NO", None, "PRIMARY KEY", "User ID"),
-            ("name", "varchar", "NO", None, "NONE", "User name"),
-            ("email", "varchar", "YES", None, "NONE", "User email address"),
+        mock_cursor.fetchall.side_effect = [
+            # First call: column metadata
+            [
+                ("id", "integer", "NO", None, "PRIMARY KEY", "User ID"),
+                ("name", "varchar", "NO", None, "NONE", "User name"),
+                ("email", "varchar", "YES", None, "NONE", "User email address"),
+            ],
+            # Second call: row count for 'id' column
+            [(100, 100)],
+            # Third call: row count for 'name' column  
+            [(100, 50)],
+            # Fourth call: row count for 'email' column
+            [(100, 80)]
         ]
 
         # Test the method