Update README, enhance SQL formatting, and improve schema handling

eosho · eosho · commit a3813800503c · 2025-12-22T20:06:29.000-05:00
- Added description of Azure AD support in README.
- Enhanced SQL formatting in output.py with dialect support.
- Implemented dynamic schema retrieval in data_nodes.py.
- Cleaned up conversation history handling in multiple files.
- Removed deprecated table schema definitions from amex.yaml.
- Updated SQLValidator to improve error handling and formatting.
- Refined response handling in app.py for better clarity.
diff --git a/README.md b/README.md
@@ -20,15 +20,15 @@
 
 A natural language to SQL (NL2SQL) platform built on LangGraph and Azure OpenAI. This multi-agent system automatically routes user questions to the appropriate database backend and generates optimized SQL queries and results.
 
+Built on top of LangChain's [`SQLDatabase`](https://docs.langchain.com/oss/python/langchain/sql-agent) with extended support for Azure AD authentication, Cosmos DB, and built-in dialect validation.
+
 ## Features
 
 - **Multi-Database Support**: PostgreSQL, Azure SQL, Azure Synapse, Azure Cosmos DB, Databricks SQL, and Google BigQuery
 - **Intent Detection**: Automatically routes queries to the correct data agent based on question context
 - **Multi-Turn Conversations**: Follow-up questions with context awareness (e.g., "What's the average?" after a query)
 - **SQL Validation**: Safe query execution with sqlglot-based validation across all dialects
 - **Configurable Agents**: YAML-based configuration for adding new data sources
-- **Conversation Persistence**: Built-in checkpointer for maintaining conversation state
-- **Async Architecture**: Fully asynchronous for high-performance query execution
 - **A2A Protocol**: Agent-to-Agent interoperability for integration with other A2A-compliant systems
 
 ## Architecture
diff --git a/src/data_agent/agent.py b/src/data_agent/agent.py
@@ -275,12 +275,11 @@ def intent_detection_node(state: AgentState) -> dict[str, Any]:
 
             system_content = intent_system_prompt.format(agent_descriptions=agent_list)
 
-            # Get recent conversation history for multi-turn context
             history = get_recent_history(state.get("messages", []), max_messages=4)
 
             messages = [
                 SystemMessage(content=system_content),
-                *history,  # Previous turns help detect "more of the same" patterns
+                *history,
                 HumanMessage(content=question),
             ]
 
@@ -356,10 +355,8 @@ def query_rewrite_node(state: AgentState) -> dict[str, Any]:
 
             agent_desc = agent_descriptions[datasource]
 
-            # Get recent conversation history for context
             history = get_recent_history(state.get("messages", []), max_messages=4)
 
-            # Build conversation context summary for the rewriter
             conversation_context = ""
             if history:
                 conversation_context = "\n## Conversation History (for context)\n"
diff --git a/src/data_agent/cli/output.py b/src/data_agent/cli/output.py
@@ -4,16 +4,22 @@
 from rich.syntax import Syntax
 
 from data_agent.cli.console import console, err_console
+from data_agent.validators import SQLValidator
 
 
-def print_sql(sql: str, title: str = "Generated SQL") -> None:
+def print_sql(
+    sql: str, title: str = "Generated SQL", dialect: str | None = None
+) -> None:
     """Print SQL with syntax highlighting in a panel.
 
     Args:
         sql: The SQL query string to display.
         title: The panel title.
+        dialect: SQL dialect for formatting.
     """
-    syntax = Syntax(sql.strip(), "sql", theme="monokai", line_numbers=False)
+    validator = SQLValidator(dialect=dialect or "postgres")
+    formatted_sql = validator.format_sql(sql)
+    syntax = Syntax(formatted_sql, "sql", theme="monokai", line_numbers=False)
     panel = Panel(syntax, title=f"[sql]{title}[/sql]", border_style="green")
     console.print(panel)
 
diff --git a/src/data_agent/config/amex.yaml b/src/data_agent/config/amex.yaml
@@ -99,228 +99,7 @@ data_agents:
       - Format dates in readable format (e.g., Dec 16, 2025)
       - Limit tables to 20 rows max; if more rows exist, show first 20 and note "... and X more rows"
       - After the table, provide a brief summary or insight about the data
-    table_schemas:
-      - table_name: customers
-        table_description: Customer profiles with segmentation and risk scoring
-        columns:
-          - column_name: customer_id
-            data_type: STRING
-            description: Unique customer identifier (e.g., CUST-001)
-          - column_name: first_name
-            data_type: STRING
-            description: Customer first name
-          - column_name: last_name
-            data_type: STRING
-            description: Customer last name
-          - column_name: email
-            data_type: STRING
-            description: Customer email address
-          - column_name: phone
-            data_type: STRING
-            description: Customer phone number
-          - column_name: date_of_birth
-            data_type: DATE
-            description: Customer date of birth
-          - column_name: registration_date
-            data_type: DATE
-            description: Date customer registered
-          - column_name: customer_segment
-            data_type: STRING
-            description: Customer tier/segment
-            allowed_values:
-              Standard: Regular customers
-              Premium: High-value customers with enhanced benefits
-              VIP: Top-tier customers with exclusive services
-          - column_name: risk_score
-            data_type: FLOAT64
-            description: Customer risk score (0.0 to 1.0, higher = riskier)
-          - column_name: is_active
-            data_type: BOOL
-            description: Whether the customer account is active
-      - table_name: accounts
-        table_description: Customer bank accounts including checking, savings, credit, and investment
-        columns:
-          - column_name: account_id
-            data_type: STRING
-            description: Unique account identifier (e.g., ACC-001-CHK)
-          - column_name: customer_id
-            data_type: STRING
-            description: Customer who owns the account
-          - column_name: account_type
-            data_type: STRING
-            description: Type of account
-            allowed_values:
-              Checking: Standard checking account
-              Savings: Savings account with interest
-              Credit: Credit card account
-              Investment: Investment/brokerage account
-          - column_name: account_status
-            data_type: STRING
-            description: Current status of the account
-            allowed_values:
-              Active: Account is open and operational
-              Frozen: Account is temporarily frozen
-              Closed: Account has been closed
-          - column_name: opened_date
-            data_type: DATE
-            description: Date the account was opened
-          - column_name: currency
-            data_type: STRING
-            description: Account currency (e.g., USD)
-          - column_name: current_balance
-            data_type: NUMERIC
-            description: Current account balance (negative for credit owed)
-          - column_name: credit_limit
-            data_type: NUMERIC
-            description: Credit limit for credit accounts (NULL for others)
-          - column_name: interest_rate
-            data_type: FLOAT64
-            description: Annual interest rate percentage
-      - table_name: transactions
-        table_description: Financial transactions including deposits, withdrawals, transfers, and payments
-        columns:
-          - column_name: transaction_id
-            data_type: STRING
-            description: Unique transaction identifier
-          - column_name: account_id
-            data_type: STRING
-            description: Account involved in the transaction
-          - column_name: transaction_timestamp
-            data_type: TIMESTAMP
-            description: Exact time of the transaction
-          - column_name: transaction_date
-            data_type: DATE
-            description: Date of the transaction
-          - column_name: transaction_type
-            data_type: STRING
-            description: Type of transaction
-            allowed_values:
-              Deposit: Money added to account
-              Withdrawal: Cash withdrawn from account
-              Transfer: Money transferred between accounts
-              Payment: Payment to merchant or bill
-              Refund: Refund received
-          - column_name: amount
-            data_type: NUMERIC
-            description: Transaction amount (negative for outgoing)
-          - column_name: currency
-            data_type: STRING
-            description: Transaction currency
-          - column_name: merchant_name
-            data_type: STRING
-            description: Merchant or payee name (NULL for deposits/transfers)
-          - column_name: merchant_category
-            data_type: STRING
-            description: Category of merchant
-            allowed_values:
-              Grocery: Grocery stores and supermarkets
-              Gas: Gas stations
-              Entertainment: Entertainment and streaming
-              Travel: Airlines, hotels, travel services
-              Electronics: Electronics and tech stores
-              Shopping: General retail shopping
-              Food: Restaurants and food delivery
-              Transportation: Ride services and public transit
-              Housing: Rent and housing payments
-              Other: Other/uncategorized
-          - column_name: channel
-            data_type: STRING
-            description: Transaction channel
-            allowed_values:
-              Branch: In-person at bank branch
-              ATM: ATM transaction
-              Online: Online/web transaction
-              Mobile: Mobile app transaction
-              POS: Point of sale terminal
-          - column_name: status
-            data_type: STRING
-            description: Transaction status
-            allowed_values:
-              Completed: Transaction completed successfully
-              Pending: Transaction is pending
-              Reversed: Transaction was reversed/cancelled
-              Failed: Transaction failed
-          - column_name: reference_id
-            data_type: STRING
-            description: External reference identifier
-          - column_name: description
-            data_type: STRING
-            description: Transaction description or notes
-      - table_name: fraud_alerts
-        table_description: Fraud detection alerts and investigation status
-        columns:
-          - column_name: alert_id
-            data_type: STRING
-            description: Unique alert identifier
-          - column_name: transaction_id
-            data_type: STRING
-            description: Transaction that triggered the alert
-          - column_name: alert_timestamp
-            data_type: TIMESTAMP
-            description: When the alert was generated
-          - column_name: alert_type
-            data_type: STRING
-            description: Type of fraud alert
-            allowed_values:
-              Suspicious Amount: Unusually large transaction amount
-              Velocity: Too many transactions in short time
-              Unusual Location: Transaction from unusual location
-              Pattern: Unusual spending pattern detected
-          - column_name: severity
-            data_type: STRING
-            description: Alert severity level
-            allowed_values:
-              Low: Minor concern, low risk
-              Medium: Moderate concern, investigation recommended
-              High: Serious concern, immediate attention needed
-              Critical: Urgent, possible active fraud
-          - column_name: status
-            data_type: STRING
-            description: Investigation status
-            allowed_values:
-              Investigating: Under active investigation
-              Confirmed Fraud: Confirmed as fraudulent
-              False Positive: Determined to be legitimate
-              Resolved: Issue has been resolved
-          - column_name: risk_score
-            data_type: FLOAT64
-            description: Calculated risk score (0.0 to 1.0)
-          - column_name: notes
-            data_type: STRING
-            description: Investigation notes
-      - table_name: monthly_summaries
-        table_description: Monthly account summaries with aggregated transaction data
-        columns:
-          - column_name: summary_id
-            data_type: STRING
-            description: Unique summary identifier
-          - column_name: account_id
-            data_type: STRING
-            description: Account the summary is for
-          - column_name: year_month
-            data_type: STRING
-            description: Summary period in YYYY-MM format
-          - column_name: opening_balance
-            data_type: NUMERIC
-            description: Balance at start of month
-          - column_name: closing_balance
-            data_type: NUMERIC
-            description: Balance at end of month
-          - column_name: total_deposits
-            data_type: NUMERIC
-            description: Total deposits during month
-          - column_name: total_withdrawals
-            data_type: NUMERIC
-            description: Total withdrawals during month
-          - column_name: transaction_count
-            data_type: INT64
-            description: Number of transactions during month
-          - column_name: avg_transaction_amount
-            data_type: NUMERIC
-            description: Average transaction amount
-          - column_name: largest_transaction
-            data_type: NUMERIC
-            description: Largest single transaction amount
+    # table_schemas omitted - will use dynamic schema discovery from BigQuery
     few_shot_examples:
       - question: What are the total deposits by customer segment this month?
         answer: This query shows total deposits grouped by customer segment (VIP, Premium, Standard) for the current month.
diff --git a/src/data_agent/nodes/data_nodes.py b/src/data_agent/nodes/data_nodes.py
@@ -122,13 +122,40 @@ def _check_is_cosmos(datasource: Any) -> bool:
         # Check by class name to avoid import issues
         return type(datasource).__name__ == "CosmosAdapter"
 
+    def _get_schema_context(self) -> str:
+        """Get schema context from config or dynamically from database.
+
+        If table_schemas is defined in config, uses the static schema.
+        Otherwise, fetches schema dynamically from the database using
+        SQLDatabase.get_table_info().
+
+        Returns:
+            Schema context string for the LLM prompt.
+        """
+        if self._config.table_schemas:
+            return SchemaFormatter.format_schema_context(self._config)
+
+        if not self._is_cosmos and isinstance(self._datasource, SQLDatabase):
+            try:
+                table_info = self._datasource.get_table_info()
+                if table_info:
+                    logger.debug(
+                        "Using dynamic schema from database. Available tables: %s",
+                        table_info,
+                    )
+                    return f"Available tables and their schemas:\n\n{table_info}"
+            except Exception as e:
+                logger.warning("Failed to fetch dynamic schema: %s", e)
+
+        return ""
+
     def _build_prompt(self) -> str:
         """Build system prompt, adding Cosmos constraints if needed.
 
         Returns:
             Formatted system prompt with schema context and date.
         """
-        schema_context = SchemaFormatter.format_schema_context(self._config)
+        schema_context = self._get_schema_context()
         few_shot = SchemaFormatter.format_few_shot_examples(self._config)
         base_prompt = self._config.system_prompt or DEFAULT_SQL_PROMPT
 
diff --git a/src/data_agent/nodes/response.py b/src/data_agent/nodes/response.py
@@ -68,12 +68,11 @@ def generate_response(self, state: "AgentState") -> dict[str, Any]:
         sql = state.get("generated_sql", "")
         result = state.get("result", {})
 
-        # Get recent conversation history for conversational continuity
         history = get_recent_history(state.get("messages", []), max_messages=4)
 
         messages = [
             SystemMessage(content=prompt),
-            *history,  # Previous turns for conversational context
+            *history,
             HumanMessage(
                 content=(
                     f"Question: {question}\n\nSQL Query: {sql}\n\nResults: {result}"
diff --git a/src/data_agent/ui/app.py b/src/data_agent/ui/app.py
diff --git a/src/data_agent/validators/sql_validator.py b/src/data_agent/validators/sql_validator.py
diff --git a/uv.lock b/uv.lock
