dbagent-cli 0.11.2__tar.gz → 0.11.3__tar.gz

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.2
+Version: 0.11.3
 Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
 Author: Santhosh Gupta
 License: MIT

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/dbagent/__init__.py

@@ -2,5 +2,5 @@
 DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
 """
 
-__version__ = "0.11.2"
+__version__ = "0.11.3"
 __author__ = "Santhosh Gupta"

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/dbagent/agent/pipeline.py

@@ -134,8 +134,13 @@ class QueryPipeline:
         # Database identifier
         self.db_name = "database"
         if hasattr(connector, "engine") and connector.engine:
-            self.db_name = connector.engine.url.database or "database"
-            self.db_id = str(connector.engine.url)
+            try:
+                db_val = getattr(getattr(connector.engine, "url", None), "database", None)
+                if db_val and isinstance(db_val, str):
+                    self.db_name = db_val
+                self.db_id = str(getattr(connector.engine, "url", self.db_name))
+            except Exception:
+                self.db_id = self.db_name
         else:
             self.db_id = self.db_name
         self.cache_mgr = SchemaCacheManager()
@@ -513,6 +518,42 @@ class QueryPipeline:
             result.error = "LLM returned empty SQL."
             return result
 
+        # --- 5.5. Schema Ground-Truth Column & Key Pre-Validation & Auto-Correction ---
+        if sub_schema and sub_schema.tables:
+            for _ in range(self.MAX_RETRIES):
+                is_schema_valid, schema_errors = SchemaGroundTruthValidator.validate_sql_against_schema(
+                    result.sql, sub_schema
+                )
+                if is_schema_valid:
+                    break
+
+                # Hallucinated columns / keys detected -> Auto-correct using physical schema definitions
+                result.retries += 1
+                error_feedback = (
+                    "Schema Ground-Truth Column & Key Validation Discrepancy:\n"
+                    + "\n".join(f"- {err}" for err in schema_errors)
+                    + "\nPlease rewrite the query using ONLY the exact physical primary keys, foreign keys, and column names defined in the schema."
+                )
+                try:
+                    corrected_sql = self.generator.generate_retry_sql(
+                        schema=sub_schema,
+                        user_prompt=augmented_prompt,
+                        failed_sql=result.sql,
+                        error_message=error_feedback,
+                        model=self.model,
+                    )
+                    if corrected_sql and corrected_sql.strip():
+                        result.sql = corrected_sql
+                    else:
+                        break
+                except Exception:
+                    break
+
+            # Attach ground truth verification metadata
+            result.ground_truth = SchemaGroundTruthValidator.validate_response(
+                f"```sql\n{result.sql}\n```", sub_schema
+            )
+
         # --- 6. AST & Security Mode Validation ---
         result.query_type = SQLValidator.classify_query(result.sql)
         is_allowed, violations = SQLValidator.validate_for_mode(result.sql, mode=self.mode)

dbagent_cli-0.11.3/dbagent/schema/ground_truth.py

@@ -0,0 +1,337 @@
+"""
+Schema Ground-Truth Verification Engine.
+Fact-checks AI-generated architectural reports and relationship mappings against the physical database schema metadata:
+- Table existence verification
+- Column existence and data type verification
+- Relational foreign key constraint verification
+- Logical / semantic mapping verification
+- SQL query dialect validation
+"""
+
+import re
+from typing import List, Dict, Any, Optional, Tuple, Set
+from dataclasses import dataclass, field
+from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
+
+
+@dataclass
+class GroundTruthValidationResult:
+    """Result of validating AI architectural response against database schema ground truth."""
+    verified_tables: List[str] = field(default_factory=list)
+    missing_tables: List[str] = field(default_factory=list)
+    verified_columns: Dict[str, List[str]] = field(default_factory=dict)
+    missing_columns: Dict[str, List[str]] = field(default_factory=dict)
+    verified_fks: List[str] = field(default_factory=list)
+    sql_queries_checked: int = 0
+    is_valid: bool = True
+    score_pct: int = 100
+
+    def format_summary(self) -> str:
+        """Render a formatted Ground-Truth Verification Card for Rich Console."""
+        color = "green" if self.score_pct >= 90 else ("yellow" if self.score_pct >= 70 else "red")
+        lines = [
+            f"\n[bold {color}]🛡️ Schema Ground-Truth Verification ({self.score_pct}% Confirmed):[/bold {color}]"
+        ]
+        if self.verified_tables:
+            lines.append(f"  [bold green]✓[/bold green] [bold white]Tables Verified ({len(self.verified_tables)}):[/bold white] [cyan]{', '.join(self.verified_tables)}[/cyan]")
+        if self.verified_fks:
+            lines.append(f"  [bold green]✓[/bold green] [bold white]Relational Foreign Keys Confirmed:[/bold white]")
+            for fk in self.verified_fks[:4]:
+                lines.append(f"    [dim]•[/dim] {fk}")
+            if len(self.verified_fks) > 4:
+                lines.append(f"    [dim]• ... and {len(self.verified_fks) - 4} more[/dim]")
+        if self.sql_queries_checked > 0:
+            lines.append(f"  [bold green]✓[/bold green] [bold white]SQL Queries Verified:[/bold white] {self.sql_queries_checked} dialect statement(s) validated against schema")
+        for missing in self.missing_tables:
+            lines.append(f"  [bold yellow]⚠ Note:[/bold yellow] Query referenced [dim]{missing}[/dim]")
+        for tbl, cols in self.missing_columns.items():
+            lines.append(f"  [bold red]⚠ Warning:[/bold red] Column(s) [yellow]{', '.join(cols)}[/yellow] not found in table [cyan]{tbl}[/cyan]")
+        return "\n".join(lines)
+
+
+class SchemaGroundTruthValidator:
+    """Validates that all tables, columns, and relationships in an AI response match real database ground-truth."""
+
+    @classmethod
+    def validate_response(cls, response_text: str, schema: Optional[DatabaseSchema]) -> Optional[GroundTruthValidationResult]:
+        if not schema or not schema.tables:
+            return None
+
+        result = GroundTruthValidationResult()
+        schema_table_map = {t.name.lower(): t for t in schema.tables}
+        
+        # 1. Identify all database tables referenced in the response
+        resp_lower = response_text.lower()
+        for t_name, t_model in schema_table_map.items():
+            if re.search(r'\b' + re.escape(t_name) + r'\b', resp_lower):
+                result.verified_tables.append(t_model.name)
+                
+                # Check columns for this table
+                col_map = {c.name.lower(): c.name for c in t_model.columns}
+                found_cols = []
+                for c_lower, c_orig in col_map.items():
+                    if len(c_lower) > 2 and re.search(r'\b' + re.escape(c_lower) + r'\b', resp_lower):
+                        found_cols.append(c_orig)
+                if found_cols:
+                    result.verified_columns[t_model.name] = found_cols
+
+                # Check foreign keys
+                for fk in t_model.foreign_keys:
+                    ref_tbl = fk.referred_table
+                    constrained = ", ".join(fk.constrained_columns)
+                    ref_cols = ", ".join(fk.referred_columns)
+                    if ref_tbl and ref_tbl.lower() in resp_lower:
+                        fk_desc = f"[cyan]{t_model.name}[/cyan] ([dim]{constrained}[/dim]) ➔ [cyan]{ref_tbl}[/cyan] ([dim]{ref_cols}[/dim])"
+                        if fk_desc not in result.verified_fks:
+                            result.verified_fks.append(fk_desc)
+
+        # 2. Extract and validate only actual SQL queries against physical schema
+        sql_blocks = re.findall(r'```(?:sql)\s*([\s\S]*?)\s*```', response_text, re.IGNORECASE)
+        result.sql_queries_checked = len(sql_blocks)
+        sql_referenced_tables = set()
+        sql_ignore_keywords = {"via", "lateral", "on", "using", "where", "select", "table", "as", "and", "or", "values"}
+
+        for block in sql_blocks:
+            # Extract FROM and JOIN table references
+            matches = re.findall(r'\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)', block, re.IGNORECASE)
+            for m in matches:
+                m_clean = m.strip().lower()
+                if m_clean not in sql_ignore_keywords:
+                    sql_referenced_tables.add(m_clean)
+
+        for sql_tbl in sql_referenced_tables:
+            if sql_tbl in schema_table_map:
+                if schema_table_map[sql_tbl].name not in result.verified_tables:
+                    result.verified_tables.append(schema_table_map[sql_tbl].name)
+            else:
+                # Check if it's a common alias or slight naming divergence
+                found_match = None
+                for real_tbl in schema_table_map:
+                    if sql_tbl in real_tbl or real_tbl in sql_tbl:
+                        found_match = schema_table_map[real_tbl].name
+                        break
+                if found_match:
+                    result.missing_tables.append(f"{sql_tbl} ➔ (Actual: {found_match})")
+                else:
+                    result.missing_tables.append(sql_tbl)
+
+        total_checked = len(result.verified_tables)
+        if result.missing_tables:
+            result.score_pct = max(70, 100 - (len(result.missing_tables) * 15))
+        elif total_checked > 0:
+            result.score_pct = 100
+        else:
+            result.score_pct = 95
+
+        return result
+
+    @classmethod
+    def _clean_sql(cls, sql: str) -> str:
+        """Strip SQL comments and string literals to avoid false identifier matching."""
+        clean = re.sub(r"--[^\n]*\n?", " ", sql)
+        clean = re.sub(r"/\*.*?\*/", " ", clean, flags=re.DOTALL)
+        clean = re.sub(r"'(''|[^'])*'", "''", clean)
+        return clean
+
+    @classmethod
+    def _clean_identifier(cls, ident: str) -> str:
+        """Remove quotes, brackets, and schema qualifiers (e.g. `public."table"` -> `table`)."""
+        ident = ident.strip()
+        if "." in ident:
+            ident = ident.split(".")[-1]
+        ident = ident.strip('"`[] \t\n;')
+        return ident
+
+    @classmethod
+    def validate_sql_against_schema(
+        cls, sql: str, schema: Optional[DatabaseSchema]
+    ) -> Tuple[bool, List[str]]:
+        """
+        Validates that all tables, columns, primary keys, and foreign keys referenced
+        in an SQL statement or script exist in the physical database schema metadata.
+
+        Returns:
+            (is_valid, list_of_error_messages)
+        """
+        if not schema or not schema.tables:
+            return True, []
+
+        if not sql or not sql.strip():
+            return True, []
+
+        from dbagent.sql.validator import SQLValidator
+
+        statements = SQLValidator.split_statements(sql)
+        errors: List[str] = []
+
+        schema_table_map: Dict[str, TableModel] = {t.name.lower(): t for t in schema.tables}
+        sql_keywords = {
+            "select", "from", "where", "join", "inner", "left", "right", "full", "cross", "outer",
+            "on", "using", "group", "by", "order", "having", "limit", "offset", "insert", "into",
+            "values", "update", "set", "delete", "create", "alter", "drop", "table", "index", "view",
+            "and", "or", "not", "in", "is", "null", "like", "ilike", "between", "exists", "case",
+            "when", "then", "else", "end", "as", "distinct", "all", "any", "some", "union",
+            "intersect", "except", "with", "recursive", "returning", "asc", "desc", "true", "false",
+            "count", "sum", "avg", "min", "max", "coalesce", "cast", "extract", "now", "date",
+            "current_date", "current_timestamp", "upper", "lower", "trim", "concat", "row_number",
+        }
+
+        for stmt in statements:
+            clean_stmt = cls._clean_sql(stmt)
+            table_alias_map: Dict[str, TableModel] = {}
+            stmt_tables: List[TableModel] = []
+
+            # 1. Discover all tables and aliases in statement
+            table_ref_patterns = [
+                r"\bFROM\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
+                r"\bJOIN\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
+                r"\bUPDATE\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
+                r"\bINSERT\s+INTO\s+([a-zA-Z0-9_\"`.]+)",
+                r"\bDELETE\s+FROM\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
+                r"\bMERGE\s+INTO\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
+            ]
+
+            for pat in table_ref_patterns:
+                for match in re.finditer(pat, clean_stmt, re.IGNORECASE):
+                    raw_table = cls._clean_identifier(match.group(1))
+                    raw_alias = cls._clean_identifier(match.group(2)) if match.lastindex and match.lastindex >= 2 and match.group(2) else None
+                    if not raw_table or raw_table.lower() in sql_keywords:
+                        continue
+
+                    tbl_lower = raw_table.lower()
+                    if tbl_lower in schema_table_map:
+                        t_model = schema_table_map[tbl_lower]
+                        if t_model not in stmt_tables:
+                            stmt_tables.append(t_model)
+                        table_alias_map[tbl_lower] = t_model
+                        table_alias_map[t_model.name.lower()] = t_model
+                        if raw_alias:
+                            alias_clean = raw_alias.lower()
+                            if alias_clean not in sql_keywords:
+                                table_alias_map[alias_clean] = t_model
+                    else:
+                        # Table not in schema - check for near matches
+                        similar = [
+                            s_name for s_name in schema_table_map
+                            if raw_table.lower() in s_name or s_name in raw_table.lower()
+                        ]
+                        hint = f" (Did you mean: {', '.join(similar[:3])}?)" if similar else ""
+                        errors.append(f"Table '{raw_table}' does not exist in the database schema.{hint}")
+
+            if not stmt_tables:
+                continue
+
+            # 2. Check UPDATE statements for invalid column assignments in SET clause
+            update_match = re.search(
+                r"\bUPDATE\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?\s+SET\s+(.*?)(?:\s+WHERE|\s+RETURNING|;|\Z)",
+                clean_stmt,
+                re.IGNORECASE | re.DOTALL,
+            )
+            if update_match:
+                upd_table_name = cls._clean_identifier(update_match.group(1)).lower()
+                target_table = schema_table_map.get(upd_table_name) or (stmt_tables[0] if stmt_tables else None)
+                if target_table:
+                    set_clause = update_match.group(3)
+                    assigns = re.findall(r"([a-zA-Z0-9_\"`.]+)\s*=", set_clause)
+                    for col_expr in assigns:
+                        col_clean = cls._clean_identifier(col_expr).lower()
+                        col_model_map = {c.name.lower(): c for c in target_table.columns}
+                        if col_clean not in col_model_map and col_clean not in sql_keywords:
+                            avail = ", ".join(c.name for c in target_table.columns[:10])
+                            pk_str = ", ".join(target_table.primary_key) if isinstance(target_table.primary_key, list) else str(target_table.primary_key or "")
+                            pk_hint = f" Primary key is '{pk_str}'." if pk_str else ""
+                            fk_cols = [fk.constrained_columns[0] for fk in target_table.foreign_keys if fk.constrained_columns]
+                            fk_hint = f" Valid foreign keys: [{', '.join(fk_cols)}]." if fk_cols else ""
+                            errors.append(
+                                f"Column '{col_clean}' does not exist in table '{target_table.name}'. Available columns: [{avail}].{pk_hint}{fk_hint}"
+                            )
+
+            # 3. Check INSERT statements for invalid column names in column list
+            insert_match = re.search(
+                r"\bINSERT\s+INTO\s+([a-zA-Z0-9_\"`.]+)\s*\((.*?)\)",
+                clean_stmt,
+                re.IGNORECASE | re.DOTALL,
+            )
+            if insert_match:
+                ins_table_name = cls._clean_identifier(insert_match.group(1)).lower()
+                target_table = schema_table_map.get(ins_table_name)
+                if target_table:
+                    col_list_str = insert_match.group(2)
+                    ins_cols = [cls._clean_identifier(c) for c in col_list_str.split(",") if c.strip()]
+                    col_model_map = {c.name.lower(): c for c in target_table.columns}
+                    for ic in ins_cols:
+                        if ic.lower() not in col_model_map and ic.lower() not in sql_keywords:
+                            avail = ", ".join(c.name for c in target_table.columns[:10])
+                            errors.append(
+                                f"Column '{ic}' does not exist in table '{target_table.name}'. Available columns: [{avail}]."
+                            )
+
+            # 4. Check qualified column references (alias.column or table.column)
+            qualified_cols = re.findall(r"\b([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)\b", clean_stmt)
+            for tbl_or_alias, col_name in qualified_cols:
+                tbl_or_alias_clean = tbl_or_alias.lower()
+                col_clean = col_name.lower()
+                if tbl_or_alias_clean in table_alias_map and col_clean != "*" and col_clean not in sql_keywords:
+                    t_model = table_alias_map[tbl_or_alias_clean]
+                    col_map = {c.name.lower(): c for c in t_model.columns}
+                    if col_clean not in col_map:
+                        avail = ", ".join(c.name for c in t_model.columns[:10])
+                        pk_str = ", ".join(t_model.primary_key) if isinstance(t_model.primary_key, list) else str(t_model.primary_key or "")
+                        pk_hint = f" Primary key: '{pk_str}'." if pk_str else ""
+                        fk_cols = [fk.constrained_columns[0] for fk in t_model.foreign_keys if fk.constrained_columns]
+                        fk_hint = f" Foreign keys: [{', '.join(fk_cols)}]." if fk_cols else ""
+                        errors.append(
+                            f"Column '{col_name}' does not exist in table '{t_model.name}'. Available columns: [{avail}].{pk_hint}{fk_hint}"
+                        )
+
+            # 5. Check unqualified column references in WHERE clauses
+            where_match = re.search(
+                r"\bWHERE\s+(.*?)(?:\s+ORDER\s+BY|\s+GROUP\s+BY|\s+LIMIT|\s+RETURNING|;|\Z)",
+                clean_stmt,
+                re.IGNORECASE | re.DOTALL,
+            )
+            if where_match:
+                where_clause = where_match.group(1)
+                # Find identifiers on the LHS of comparisons: col = val, col IN (...), etc.
+                lhs_idents = re.findall(
+                    r"(?:^|\s+)([a-zA-Z0-9_]+)\s*(?:=|!=|<>|<=|>=|<|>|\bLIKE\b|\bILIKE\b|\bIS\b|\bIN\b|\bBETWEEN\b)",
+                    where_clause,
+                    re.IGNORECASE,
+                )
+                for ident in lhs_idents:
+                    ident_clean = ident.strip().lower()
+                    if ident_clean in sql_keywords or ident_clean.isdigit():
+                        continue
+                    if "." in ident:
+                        continue  # Handled by qualified check
+
+                    if len(stmt_tables) == 1:
+                        target_t = stmt_tables[0]
+                        col_map = {c.name.lower(): c for c in target_t.columns}
+                        if ident_clean not in col_map and ident_clean not in table_alias_map:
+                            avail = ", ".join(c.name for c in target_t.columns[:10])
+                            pk_str = ", ".join(target_t.primary_key) if isinstance(target_t.primary_key, list) else str(target_t.primary_key or "")
+                            pk_hint = f" Primary key is '{pk_str}'." if pk_str else ""
+                            errors.append(
+                                f"Column '{ident}' in WHERE clause does not exist in table '{target_t.name}'. Available columns: [{avail}].{pk_hint}"
+                            )
+                    elif len(stmt_tables) > 1:
+                        # Multi-table query: check if column exists in ANY of the statement's tables
+                        all_stmt_cols = set()
+                        for t in stmt_tables:
+                            for c in t.columns:
+                                all_stmt_cols.add(c.name.lower())
+                        if ident_clean not in all_stmt_cols and ident_clean not in table_alias_map:
+                            tables_str = ", ".join(t.name for t in stmt_tables)
+                            errors.append(
+                                f"Column '{ident}' in WHERE clause does not exist in any of the referenced tables: [{tables_str}]."
+                            )
+
+        # Deduplicate error messages while preserving order
+        deduped_errors: List[str] = []
+        for e in errors:
+            if e not in deduped_errors:
+                deduped_errors.append(e)
+
+        is_valid = len(deduped_errors) == 0
+        return is_valid, deduped_errors

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/dbagent_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dbagent-cli
-Version: 0.11.2
+Version: 0.11.3
 Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
 Author: Santhosh Gupta
 License: MIT

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/dbagent_cli.egg-info/SOURCES.txt

@@ -81,6 +81,7 @@ tests/test_doctor_and_viz.py
 tests/test_formats_and_exit_codes.py
 tests/test_generator.py
 tests/test_graph.py
+tests/test_ground_truth_sql_validation.py
 tests/test_knowledge.py
 tests/test_multi_turn_conversation.py
 tests/test_plugins.py

{dbagent_cli-0.11.2 → dbagent_cli-0.11.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dbagent-cli"
-version = "0.11.2"
+version = "0.11.3"
 description = "The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)"
 readme = "README.md"
 requires-python = ">=3.9"

dbagent_cli-0.11.3/tests/test_ground_truth_sql_validation.py

@@ -0,0 +1,150 @@
+"""
+Unit tests for Schema Ground-Truth SQL Validation and Auto-Correction.
+Verifies table, column, PK, and FK validation against physical database schema.
+"""
+
+import pytest
+from unittest.mock import MagicMock
+from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
+from dbagent.schema.ground_truth import SchemaGroundTruthValidator
+from dbagent.agent.pipeline import QueryPipeline, PipelineResult
+from dbagent.security.modes import OperationMode
+
+
+@pytest.fixture
+def sample_schema():
+    """Create a sample schema with realistic primary keys and foreign keys."""
+    ivis_site = TableModel(
+        name="ivis_site",
+        columns=[
+            ColumnModel(name="pk_site_id", data_type="INTEGER", is_primary_key=True),
+            ColumnModel(name="site_name", data_type="VARCHAR(255)"),
+            ColumnModel(name="fk_project_id", data_type="INTEGER"),
+            ColumnModel(name="fk_site_group_id", data_type="INTEGER"),
+            ColumnModel(name="is_active", data_type="BOOLEAN"),
+        ],
+        foreign_keys=[
+            ForeignKeyModel(
+                constrained_columns=["fk_project_id"],
+                referred_table="ivis_project",
+                referred_columns=["pk_project_id"],
+            ),
+            ForeignKeyModel(
+                constrained_columns=["fk_site_group_id"],
+                referred_table="ivis_site_group",
+                referred_columns=["pk_site_group_id"],
+            ),
+        ],
+        primary_key=["pk_site_id"],
+    )
+
+    ivis_project = TableModel(
+        name="ivis_project",
+        columns=[
+            ColumnModel(name="pk_project_id", data_type="INTEGER", is_primary_key=True),
+            ColumnModel(name="project_name", data_type="VARCHAR(255)"),
+            ColumnModel(name="customer_code", data_type="VARCHAR(50)"),
+        ],
+        foreign_keys=[],
+        primary_key=["pk_project_id"],
+    )
+
+    ivis_site_group = TableModel(
+        name="ivis_site_group",
+        columns=[
+            ColumnModel(name="pk_site_group_id", data_type="INTEGER", is_primary_key=True),
+            ColumnModel(name="group_name", data_type="VARCHAR(255)"),
+        ],
+        foreign_keys=[],
+        primary_key=["pk_site_group_id"],
+    )
+
+    return DatabaseSchema(
+        database_name="iportal",
+        dialect_name="postgresql",
+        tables=[ivis_site, ivis_project, ivis_site_group],
+    )
+
+
+def test_valid_select_query_passes_ground_truth(sample_schema):
+    sql = "SELECT s.pk_site_id, s.site_name, p.project_name FROM ivis_site s JOIN ivis_project p ON s.fk_project_id = p.pk_project_id WHERE s.is_active = true;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is True
+    assert len(errors) == 0
+
+
+def test_valid_update_query_passes_ground_truth(sample_schema):
+    sql = "UPDATE ivis_site SET fk_project_id = 10, fk_site_group_id = 20 WHERE pk_site_id = 100;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is True
+    assert len(errors) == 0
+
+
+def test_hallucinated_column_in_update_set_detected(sample_schema):
+    sql = "UPDATE ivis_site SET project = 5, site_group = 10 WHERE pk_site_id = 100;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is False
+    assert any("Column 'project' does not exist in table 'ivis_site'" in e for e in errors)
+    assert any("Column 'site_group' does not exist in table 'ivis_site'" in e for e in errors)
+    assert any("Valid foreign keys" in e for e in errors)
+
+
+def test_hallucinated_column_in_where_clause_detected(sample_schema):
+    sql = "UPDATE ivis_site SET fk_project_id = 5 WHERE customer_site_id = 100;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is False
+    assert any("Column 'customer_site_id' in WHERE clause does not exist in table 'ivis_site'" in e for e in errors)
+    assert any("Primary key is 'pk_site_id'" in e for e in errors)
+
+
+def test_hallucinated_table_name_detected(sample_schema):
+    sql = "SELECT * FROM non_existent_table WHERE id = 1;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is False
+    assert any("Table 'non_existent_table' does not exist" in e for e in errors)
+
+
+def test_hallucinated_qualified_column_detected(sample_schema):
+    sql = "SELECT s.pk_site_id, s.fake_column FROM ivis_site s;"
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is False
+    assert any("Column 'fake_column' does not exist in table 'ivis_site'" in e for e in errors)
+
+
+def test_multi_statement_validation(sample_schema):
+    sql = """
+    UPDATE ivis_site SET fk_project_id = 2 WHERE pk_site_id = 1;
+    UPDATE ivis_site SET bad_col = 3 WHERE pk_site_id = 2;
+    """
+    is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
+    assert is_valid is False
+    assert any("Column 'bad_col' does not exist in table 'ivis_site'" in e for e in errors)
+
+
+def test_pipeline_auto_corrects_hallucinated_columns(sample_schema):
+    mock_connector = MagicMock()
+    mock_connector.get_table_names.return_value = ["ivis_site", "ivis_project", "ivis_site_group"]
+    mock_connector.resolve_tables.return_value = (["ivis_site"], [])
+    mock_connector.inspect_schema.return_value = sample_schema
+    mock_connector.execute_query.return_value = (["pk_site_id", "fk_project_id"], [{"pk_site_id": 1, "fk_project_id": 2}], None)
+
+    mock_llm = MagicMock()
+    mock_llm.generate.side_effect = [
+        "UPDATE ivis_site SET project = 2 WHERE site_id = 1;",
+        "UPDATE ivis_site SET fk_project_id = 2 WHERE pk_site_id = 1;",
+    ]
+
+    pipeline = QueryPipeline(
+        connector=mock_connector,
+        llm=mock_llm,
+        mode=OperationMode.CONFIRM,
+        auto_execute=True,
+    )
+
+    result = pipeline.run("Change the project for site 1 to project 2", force=True)
+
+    assert result.success is True
+    assert "fk_project_id = 2" in result.sql
+    assert "pk_site_id = 1" in result.sql
+    assert result.retries >= 1
+    assert result.ground_truth is not None

dbagent_cli-0.11.2/dbagent/schema/ground_truth.py

@@ -1,126 +0,0 @@
-"""
-Schema Ground-Truth Verification Engine.
-Fact-checks AI-generated architectural reports and relationship mappings against the physical database schema metadata:
-- Table existence verification
-- Column existence and data type verification
-- Relational foreign key constraint verification
-- Logical / semantic mapping verification
-- SQL query dialect validation
-"""
-
-import re
-from typing import List, Dict, Any, Optional
-from dataclasses import dataclass, field
-from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
-
-
-@dataclass
-class GroundTruthValidationResult:
-    """Result of validating AI architectural response against database schema ground truth."""
-    verified_tables: List[str] = field(default_factory=list)
-    missing_tables: List[str] = field(default_factory=list)
-    verified_columns: Dict[str, List[str]] = field(default_factory=dict)
-    missing_columns: Dict[str, List[str]] = field(default_factory=dict)
-    verified_fks: List[str] = field(default_factory=list)
-    sql_queries_checked: int = 0
-    is_valid: bool = True
-    score_pct: int = 100
-
-    def format_summary(self) -> str:
-        """Render a formatted Ground-Truth Verification Card for Rich Console."""
-        color = "green" if self.score_pct >= 90 else ("yellow" if self.score_pct >= 70 else "red")
-        lines = [
-            f"\n[bold {color}]🛡️ Schema Ground-Truth Verification ({self.score_pct}% Confirmed):[/bold {color}]"
-        ]
-        if self.verified_tables:
-            lines.append(f"  [bold green]✓[/bold green] [bold white]Tables Verified ({len(self.verified_tables)}):[/bold white] [cyan]{', '.join(self.verified_tables)}[/cyan]")
-        if self.verified_fks:
-            lines.append(f"  [bold green]✓[/bold green] [bold white]Relational Foreign Keys Confirmed:[/bold white]")
-            for fk in self.verified_fks[:4]:
-                lines.append(f"    [dim]•[/dim] {fk}")
-            if len(self.verified_fks) > 4:
-                lines.append(f"    [dim]• ... and {len(self.verified_fks) - 4} more[/dim]")
-        if self.sql_queries_checked > 0:
-            lines.append(f"  [bold green]✓[/bold green] [bold white]SQL Queries Verified:[/bold white] {self.sql_queries_checked} dialect statement(s) validated against schema")
-        for missing in self.missing_tables:
-            lines.append(f"  [bold yellow]⚠ Note:[/bold yellow] Query referenced [dim]{missing}[/dim]")
-        for tbl, cols in self.missing_columns.items():
-            lines.append(f"  [bold red]⚠ Warning:[/bold red] Column(s) [yellow]{', '.join(cols)}[/yellow] not found in table [cyan]{tbl}[/cyan]")
-        return "\n".join(lines)
-
-
-class SchemaGroundTruthValidator:
-    """Validates that all tables, columns, and relationships in an AI response match real database ground-truth."""
-
-    @classmethod
-    def validate_response(cls, response_text: str, schema: Optional[DatabaseSchema]) -> Optional[GroundTruthValidationResult]:
-        if not schema or not schema.tables:
-            return None
-
-        result = GroundTruthValidationResult()
-        schema_table_map = {t.name.lower(): t for t in schema.tables}
-        
-        # 1. Identify all database tables referenced in the response
-        resp_lower = response_text.lower()
-        for t_name, t_model in schema_table_map.items():
-            if re.search(r'\b' + re.escape(t_name) + r'\b', resp_lower):
-                result.verified_tables.append(t_model.name)
-                
-                # Check columns for this table
-                col_map = {c.name.lower(): c.name for c in t_model.columns}
-                found_cols = []
-                for c_lower, c_orig in col_map.items():
-                    if len(c_lower) > 2 and re.search(r'\b' + re.escape(c_lower) + r'\b', resp_lower):
-                        found_cols.append(c_orig)
-                if found_cols:
-                    result.verified_columns[t_model.name] = found_cols
-
-                # Check foreign keys
-                for fk in t_model.foreign_keys:
-                    ref_tbl = fk.referred_table
-                    constrained = ", ".join(fk.constrained_columns)
-                    ref_cols = ", ".join(fk.referred_columns)
-                    if ref_tbl and ref_tbl.lower() in resp_lower:
-                        fk_desc = f"[cyan]{t_model.name}[/cyan] ([dim]{constrained}[/dim]) ➔ [cyan]{ref_tbl}[/cyan] ([dim]{ref_cols}[/dim])"
-                        if fk_desc not in result.verified_fks:
-                            result.verified_fks.append(fk_desc)
-
-        # 2. Extract and validate only actual SQL queries against physical schema
-        sql_blocks = re.findall(r'```(?:sql)\s*([\s\S]*?)\s*```', response_text, re.IGNORECASE)
-        result.sql_queries_checked = len(sql_blocks)
-        sql_referenced_tables = set()
-        sql_ignore_keywords = {"via", "lateral", "on", "using", "where", "select", "table", "as", "and", "or", "values"}
-
-        for block in sql_blocks:
-            # Extract FROM and JOIN table references
-            matches = re.findall(r'\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)', block, re.IGNORECASE)
-            for m in matches:
-                m_clean = m.strip().lower()
-                if m_clean not in sql_ignore_keywords:
-                    sql_referenced_tables.add(m_clean)
-
-        for sql_tbl in sql_referenced_tables:
-            if sql_tbl in schema_table_map:
-                if schema_table_map[sql_tbl].name not in result.verified_tables:
-                    result.verified_tables.append(schema_table_map[sql_tbl].name)
-            else:
-                # Check if it's a common alias or slight naming divergence
-                found_match = None
-                for real_tbl in schema_table_map:
-                    if sql_tbl in real_tbl or real_tbl in sql_tbl:
-                        found_match = schema_table_map[real_tbl].name
-                        break
-                if found_match:
-                    result.missing_tables.append(f"{sql_tbl} ➔ (Actual: {found_match})")
-                else:
-                    result.missing_tables.append(sql_tbl)
-
-        total_checked = len(result.verified_tables)
-        if result.missing_tables:
-            result.score_pct = max(70, 100 - (len(result.missing_tables) * 15))
-        elif total_checked > 0:
-            result.score_pct = 100
-        else:
-            result.score_pct = 95
-
-        return result