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
+    POSTGRES_MAX_ROWS: int = 500
+    POSTGRES_MAX_DISTINCT: int = 100
+    POSTGRES_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/postgres_loader.py

@@ -1,9 +1,10 @@
 """PostgreSQL loader for loading database schemas into FalkorDB graphs."""
 
+import re
 import datetime
 import decimal
 import logging
-import re
+from api.config import Config
 from typing import Tuple, Dict, Any, List
 
 import psycopg2
@@ -233,6 +234,36 @@ def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
             if column_default:
                 description_parts.append(f"(Default: {column_default})")
 
+            # 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()
+            rows_count, distinct_count = output[0]
+
+            max_rows = Config.POSTGRES_MAX_ROWS
+            max_distinct = Config.POSTGRES_MAX_DISTINCT
+            uniqueness_threshold = Config.POSTGRES_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_values = [row[0] for row in cursor.fetchall() 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)):
+                            description_parts.append(
+                                f"(Optional values: {', '.join(f'({str(v)})' for v in distinct_values)})"
+                            )
+
             columns_info[col_name] = {
                 'type': data_type,
                 'null': is_nullable,
@@ -241,6 +272,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