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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (99) hide show
  1. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/PKG-INFO +1 -1
  2. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/__init__.py +1 -1
  3. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/generator.py +2 -2
  4. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/pipeline.py +44 -3
  5. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/cli.py +3 -3
  6. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/base.py +2 -2
  7. dbagent_cli-0.11.4/dbagent/schema/ground_truth.py +337 -0
  8. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/PKG-INFO +1 -1
  9. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/SOURCES.txt +1 -0
  10. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/pyproject.toml +1 -1
  11. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_direct_sql.py +1 -1
  12. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_generator.py +1 -1
  13. dbagent_cli-0.11.4/tests/test_ground_truth_sql_validation.py +150 -0
  14. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_knowledge.py +8 -8
  15. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_prompt_queue.py +2 -2
  16. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_prompt_queue_and_sql.py +7 -7
  17. dbagent_cli-0.11.2/dbagent/schema/ground_truth.py +0 -126
  18. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/README.md +0 -0
  19. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/confidence.py +0 -0
  20. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/conversation.py +0 -0
  21. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/doctor.py +0 -0
  22. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/error_classifier.py +0 -0
  23. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/recommender.py +0 -0
  24. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/validator.py +0 -0
  25. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/benchmark.py +0 -0
  26. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/config.py +0 -0
  27. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/factory.py +0 -0
  28. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/mongo.py +0 -0
  29. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/relational.py +0 -0
  30. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/errors.py +0 -0
  31. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/__init__.py +0 -0
  32. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/cache.py +0 -0
  33. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/changeset.py +0 -0
  34. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/cost.py +0 -0
  35. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/prompt_queue.py +0 -0
  36. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/recovery.py +0 -0
  37. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/sandbox.py +0 -0
  38. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/streaming.py +0 -0
  39. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/anthropic_provider.py +0 -0
  40. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/auto_setup.py +0 -0
  41. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/base.py +0 -0
  42. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/deepseek_provider.py +0 -0
  43. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/factory.py +0 -0
  44. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/gemini_provider.py +0 -0
  45. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/groq_provider.py +0 -0
  46. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/mistral_provider.py +0 -0
  47. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/mock_provider.py +0 -0
  48. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/ollama_provider.py +0 -0
  49. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/openai_provider.py +0 -0
  50. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/openrouter_provider.py +0 -0
  51. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/plugins.py +0 -0
  52. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sample_db.py +0 -0
  53. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/cache.py +0 -0
  54. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/formatter.py +0 -0
  55. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/graph.py +0 -0
  56. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/models.py +0 -0
  57. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/retriever.py +0 -0
  58. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/selector.py +0 -0
  59. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/__init__.py +0 -0
  60. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/audit.py +0 -0
  61. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/environment.py +0 -0
  62. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/masking.py +0 -0
  63. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/modes.py +0 -0
  64. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/rbac.py +0 -0
  65. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/__init__.py +0 -0
  66. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/dictionary.py +0 -0
  67. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/knowledge.py +0 -0
  68. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/templates.py +0 -0
  69. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/__init__.py +0 -0
  70. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/optimizer.py +0 -0
  71. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/validator.py +0 -0
  72. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/chat_screen.py +0 -0
  73. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/console.py +0 -0
  74. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/grid_window.py +0 -0
  75. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/live_streamer.py +0 -0
  76. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/viewer.py +0 -0
  77. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/visualizer.py +0 -0
  78. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/dependency_links.txt +0 -0
  79. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/entry_points.txt +0 -0
  80. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/requires.txt +0 -0
  81. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/top_level.txt +0 -0
  82. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/setup.cfg +0 -0
  83. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_benchmarks.py +0 -0
  84. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_cache.py +0 -0
  85. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_cli.py +0 -0
  86. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_connectors.py +0 -0
  87. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_demo_and_sample_db.py +0 -0
  88. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_doctor_and_viz.py +0 -0
  89. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_formats_and_exit_codes.py +0 -0
  90. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_graph.py +0 -0
  91. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_multi_turn_conversation.py +0 -0
  92. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_plugins.py +0 -0
  93. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_providers.py +0 -0
  94. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_recovery.py +0 -0
  95. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_safety.py +0 -0
  96. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_sandbox.py +0 -0
  97. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_schema.py +0 -0
  98. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_security.py +0 -0
  99. {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_semantic.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.11.2
3
+ Version: 0.11.4
4
4
  Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
5
5
  Author: Santhosh Gupta
6
6
  License: MIT
@@ -2,5 +2,5 @@
2
2
  DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
3
3
  """
4
4
 
5
- __version__ = "0.11.2"
5
+ __version__ = "0.11.4"
6
6
  __author__ = "Santhosh Gupta"
@@ -60,7 +60,7 @@ Generate ONLY the raw executable SQL query for the given database. No markdown f
60
60
  - If the column is VARCHAR(1) or CHAR(1) or text, check if values use 'Y'/'N', 'YES'/'NO', '1'/'0', 'ACTIVE'/'INACTIVE', 'TRUE'/'FALSE'.
61
61
  - If sample values show 'Y'/'N' (or if it's a single-character VARCHAR/CHAR column), use `active = 'Y'` or `active IN ('Y', '1', 'true', 'TRUE')` for active/enabled status.
62
62
  10. For Deletions & Cascading Mapping Cleanup (when deleting a record/user without foreign key constraint violations):
63
- - For PostgreSQL: Use a single Common Table Expression (CTE) with data-modifying statements deleting from all child, mapping, activity, profile, and preference tables first, and finally deleting from the parent table (e.g. `WITH target AS (SELECT pk_user_id FROM ivis_user WHERE ...), del_act AS (DELETE FROM user_activity WHERE user_id IN (SELECT pk_user_id FROM target)), del_map AS (DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target)) DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target);`).
63
+ - For PostgreSQL: Use a single Common Table Expression (CTE) with data-modifying statements deleting from all child, mapping, activity, profile, and preference tables first, and finally deleting from the parent table (e.g. `WITH target AS (SELECT id FROM users WHERE status = 'inactive'), del_orders AS (DELETE FROM orders WHERE user_id IN (SELECT id FROM target)), del_logs AS (DELETE FROM user_logs WHERE user_id IN (SELECT id FROM target)) DELETE FROM users WHERE id IN (SELECT id FROM target);`).
64
64
  - For MySQL / SQLite: Generate clean, ordered DELETE statements deleting referencing child records before the parent record.
65
65
  - Always delete from all referencing tables present in the schema to avoid foreign key violations.
66
66
  11. Multi-Turn Session Memory & Query Refinement:
@@ -201,7 +201,7 @@ class ScriptGenerator:
201
201
  if re.search(r'\b(?:is|are)\b.*\b(?:same|different|difference|related|equal|linked)\b', prompt_lower):
202
202
  return "analysis"
203
203
 
204
- # Contextual clarification for analysis (e.g. "i mean ivis_scope", "i mean table_x")
204
+ # Contextual clarification for analysis (e.g. "i mean user_profiles", "i mean table_x")
205
205
  if re.search(r'\bi mean\b', prompt_lower) and not re.search(r'\b(insert|update|delete|drop)\b', prompt_lower):
206
206
  return "analysis"
207
207
 
@@ -134,8 +134,13 @@ class QueryPipeline:
134
134
  # Database identifier
135
135
  self.db_name = "database"
136
136
  if hasattr(connector, "engine") and connector.engine:
137
- self.db_name = connector.engine.url.database or "database"
138
- self.db_id = str(connector.engine.url)
137
+ try:
138
+ db_val = getattr(getattr(connector.engine, "url", None), "database", None)
139
+ if db_val and isinstance(db_val, str):
140
+ self.db_name = db_val
141
+ self.db_id = str(getattr(connector.engine, "url", self.db_name))
142
+ except Exception:
143
+ self.db_id = self.db_name
139
144
  else:
140
145
  self.db_id = self.db_name
141
146
  self.cache_mgr = SchemaCacheManager()
@@ -171,7 +176,7 @@ class QueryPipeline:
171
176
  all_table_names = self.connector.get_table_names()
172
177
  stems = set()
173
178
  for t in target_names:
174
- unprefixed = re.sub(r'^(?:ivis_|tbl_|table_|app_)', '', t.lower())
179
+ unprefixed = re.sub(r'^(?:tbl_|table_|app_|dim_|fact_|rel_|raw_)', '', t.lower())
175
180
  unprefixed = re.sub(r'(?:s|es)$', '', unprefixed)
176
181
  if len(unprefixed) >= 3:
177
182
  stems.add(unprefixed)
@@ -513,6 +518,42 @@ class QueryPipeline:
513
518
  result.error = "LLM returned empty SQL."
514
519
  return result
515
520
 
521
+ # --- 5.5. Schema Ground-Truth Column & Key Pre-Validation & Auto-Correction ---
522
+ if sub_schema and sub_schema.tables:
523
+ for _ in range(self.MAX_RETRIES):
524
+ is_schema_valid, schema_errors = SchemaGroundTruthValidator.validate_sql_against_schema(
525
+ result.sql, sub_schema
526
+ )
527
+ if is_schema_valid:
528
+ break
529
+
530
+ # Hallucinated columns / keys detected -> Auto-correct using physical schema definitions
531
+ result.retries += 1
532
+ error_feedback = (
533
+ "Schema Ground-Truth Column & Key Validation Discrepancy:\n"
534
+ + "\n".join(f"- {err}" for err in schema_errors)
535
+ + "\nPlease rewrite the query using ONLY the exact physical primary keys, foreign keys, and column names defined in the schema."
536
+ )
537
+ try:
538
+ corrected_sql = self.generator.generate_retry_sql(
539
+ schema=sub_schema,
540
+ user_prompt=augmented_prompt,
541
+ failed_sql=result.sql,
542
+ error_message=error_feedback,
543
+ model=self.model,
544
+ )
545
+ if corrected_sql and corrected_sql.strip():
546
+ result.sql = corrected_sql
547
+ else:
548
+ break
549
+ except Exception:
550
+ break
551
+
552
+ # Attach ground truth verification metadata
553
+ result.ground_truth = SchemaGroundTruthValidator.validate_response(
554
+ f"```sql\n{result.sql}\n```", sub_schema
555
+ )
556
+
516
557
  # --- 6. AST & Security Mode Validation ---
517
558
  result.query_type = SQLValidator.classify_query(result.sql)
518
559
  is_allowed, violations = SQLValidator.validate_for_mode(result.sql, mode=self.mode)
@@ -808,7 +808,7 @@ def chat_command(
808
808
 
809
809
  if cmd_name == ":table":
810
810
  if not cmd_arg:
811
- print_warning("Usage: :table <table_name> (e.g. :table ivis_users)")
811
+ print_warning("Usage: :table <table_name> (e.g. :table users)")
812
812
  continue
813
813
  exact, fuzzy = connector.resolve_tables(cmd_arg)
814
814
  target = exact[0] if exact else (fuzzy[0] if fuzzy else cmd_arg)
@@ -823,7 +823,7 @@ def chat_command(
823
823
 
824
824
  if cmd_name in [":fk", ":relations", ":mapped"]:
825
825
  if not cmd_arg:
826
- print_warning("Usage: :fk <table_name> (e.g. :fk ivis_users)")
826
+ print_warning("Usage: :fk <table_name> (e.g. :fk users)")
827
827
  continue
828
828
  exact, fuzzy = connector.resolve_tables(cmd_arg)
829
829
  target = exact[0] if exact else (fuzzy[0] if fuzzy else cmd_arg)
@@ -931,7 +931,7 @@ def chat_command(
931
931
 
932
932
  if cmd_name in [":learn", ":remember"]:
933
933
  if not cmd_arg or "=" not in cmd_arg:
934
- print_warning("Usage: :learn <topic> = <description> (e.g. :learn site_user_map = ivis_site maps to ivis_user via site_group -> group)")
934
+ print_warning("Usage: :learn <topic> = <description> (e.g. :learn user_roles = users map to roles via user_role_map)")
935
935
  continue
936
936
  parts = cmd_arg.split("=", 1)
937
937
  topic = parts[0].strip()
@@ -87,8 +87,8 @@ class BaseConnector(ABC):
87
87
  exact.append(t)
88
88
  continue
89
89
 
90
- # 2. Match without common prefixes (e.g. "ivis_scope" -> "scope", "scopes")
91
- unprefixed = re.sub(r'^(?:ivis_|tbl_|table_|app_)', '', t_lower)
90
+ # 2. Match without common database table prefixes (e.g. "tbl_orders" -> "orders")
91
+ unprefixed = re.sub(r'^(?:tbl_|table_|app_|dim_|fact_|rel_|raw_)', '', t_lower)
92
92
  if unprefixed and len(unprefixed) >= 3 and unprefixed not in stop_words:
93
93
  if (unprefixed in tokens or
94
94
  f"{unprefixed}s" in tokens or
@@ -0,0 +1,337 @@
1
+ """
2
+ Schema Ground-Truth Verification Engine.
3
+ Fact-checks AI-generated architectural reports and relationship mappings against the physical database schema metadata:
4
+ - Table existence verification
5
+ - Column existence and data type verification
6
+ - Relational foreign key constraint verification
7
+ - Logical / semantic mapping verification
8
+ - SQL query dialect validation
9
+ """
10
+
11
+ import re
12
+ from typing import List, Dict, Any, Optional, Tuple, Set
13
+ from dataclasses import dataclass, field
14
+ from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
15
+
16
+
17
+ @dataclass
18
+ class GroundTruthValidationResult:
19
+ """Result of validating AI architectural response against database schema ground truth."""
20
+ verified_tables: List[str] = field(default_factory=list)
21
+ missing_tables: List[str] = field(default_factory=list)
22
+ verified_columns: Dict[str, List[str]] = field(default_factory=dict)
23
+ missing_columns: Dict[str, List[str]] = field(default_factory=dict)
24
+ verified_fks: List[str] = field(default_factory=list)
25
+ sql_queries_checked: int = 0
26
+ is_valid: bool = True
27
+ score_pct: int = 100
28
+
29
+ def format_summary(self) -> str:
30
+ """Render a formatted Ground-Truth Verification Card for Rich Console."""
31
+ color = "green" if self.score_pct >= 90 else ("yellow" if self.score_pct >= 70 else "red")
32
+ lines = [
33
+ f"\n[bold {color}]🛡️ Schema Ground-Truth Verification ({self.score_pct}% Confirmed):[/bold {color}]"
34
+ ]
35
+ if self.verified_tables:
36
+ lines.append(f" [bold green]✓[/bold green] [bold white]Tables Verified ({len(self.verified_tables)}):[/bold white] [cyan]{', '.join(self.verified_tables)}[/cyan]")
37
+ if self.verified_fks:
38
+ lines.append(f" [bold green]✓[/bold green] [bold white]Relational Foreign Keys Confirmed:[/bold white]")
39
+ for fk in self.verified_fks[:4]:
40
+ lines.append(f" [dim]•[/dim] {fk}")
41
+ if len(self.verified_fks) > 4:
42
+ lines.append(f" [dim]• ... and {len(self.verified_fks) - 4} more[/dim]")
43
+ if self.sql_queries_checked > 0:
44
+ lines.append(f" [bold green]✓[/bold green] [bold white]SQL Queries Verified:[/bold white] {self.sql_queries_checked} dialect statement(s) validated against schema")
45
+ for missing in self.missing_tables:
46
+ lines.append(f" [bold yellow]⚠ Note:[/bold yellow] Query referenced [dim]{missing}[/dim]")
47
+ for tbl, cols in self.missing_columns.items():
48
+ lines.append(f" [bold red]⚠ Warning:[/bold red] Column(s) [yellow]{', '.join(cols)}[/yellow] not found in table [cyan]{tbl}[/cyan]")
49
+ return "\n".join(lines)
50
+
51
+
52
+ class SchemaGroundTruthValidator:
53
+ """Validates that all tables, columns, and relationships in an AI response match real database ground-truth."""
54
+
55
+ @classmethod
56
+ def validate_response(cls, response_text: str, schema: Optional[DatabaseSchema]) -> Optional[GroundTruthValidationResult]:
57
+ if not schema or not schema.tables:
58
+ return None
59
+
60
+ result = GroundTruthValidationResult()
61
+ schema_table_map = {t.name.lower(): t for t in schema.tables}
62
+
63
+ # 1. Identify all database tables referenced in the response
64
+ resp_lower = response_text.lower()
65
+ for t_name, t_model in schema_table_map.items():
66
+ if re.search(r'\b' + re.escape(t_name) + r'\b', resp_lower):
67
+ result.verified_tables.append(t_model.name)
68
+
69
+ # Check columns for this table
70
+ col_map = {c.name.lower(): c.name for c in t_model.columns}
71
+ found_cols = []
72
+ for c_lower, c_orig in col_map.items():
73
+ if len(c_lower) > 2 and re.search(r'\b' + re.escape(c_lower) + r'\b', resp_lower):
74
+ found_cols.append(c_orig)
75
+ if found_cols:
76
+ result.verified_columns[t_model.name] = found_cols
77
+
78
+ # Check foreign keys
79
+ for fk in t_model.foreign_keys:
80
+ ref_tbl = fk.referred_table
81
+ constrained = ", ".join(fk.constrained_columns)
82
+ ref_cols = ", ".join(fk.referred_columns)
83
+ if ref_tbl and ref_tbl.lower() in resp_lower:
84
+ fk_desc = f"[cyan]{t_model.name}[/cyan] ([dim]{constrained}[/dim]) ➔ [cyan]{ref_tbl}[/cyan] ([dim]{ref_cols}[/dim])"
85
+ if fk_desc not in result.verified_fks:
86
+ result.verified_fks.append(fk_desc)
87
+
88
+ # 2. Extract and validate only actual SQL queries against physical schema
89
+ sql_blocks = re.findall(r'```(?:sql)\s*([\s\S]*?)\s*```', response_text, re.IGNORECASE)
90
+ result.sql_queries_checked = len(sql_blocks)
91
+ sql_referenced_tables = set()
92
+ sql_ignore_keywords = {"via", "lateral", "on", "using", "where", "select", "table", "as", "and", "or", "values"}
93
+
94
+ for block in sql_blocks:
95
+ # Extract FROM and JOIN table references
96
+ matches = re.findall(r'\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)', block, re.IGNORECASE)
97
+ for m in matches:
98
+ m_clean = m.strip().lower()
99
+ if m_clean not in sql_ignore_keywords:
100
+ sql_referenced_tables.add(m_clean)
101
+
102
+ for sql_tbl in sql_referenced_tables:
103
+ if sql_tbl in schema_table_map:
104
+ if schema_table_map[sql_tbl].name not in result.verified_tables:
105
+ result.verified_tables.append(schema_table_map[sql_tbl].name)
106
+ else:
107
+ # Check if it's a common alias or slight naming divergence
108
+ found_match = None
109
+ for real_tbl in schema_table_map:
110
+ if sql_tbl in real_tbl or real_tbl in sql_tbl:
111
+ found_match = schema_table_map[real_tbl].name
112
+ break
113
+ if found_match:
114
+ result.missing_tables.append(f"{sql_tbl} ➔ (Actual: {found_match})")
115
+ else:
116
+ result.missing_tables.append(sql_tbl)
117
+
118
+ total_checked = len(result.verified_tables)
119
+ if result.missing_tables:
120
+ result.score_pct = max(70, 100 - (len(result.missing_tables) * 15))
121
+ elif total_checked > 0:
122
+ result.score_pct = 100
123
+ else:
124
+ result.score_pct = 95
125
+
126
+ return result
127
+
128
+ @classmethod
129
+ def _clean_sql(cls, sql: str) -> str:
130
+ """Strip SQL comments and string literals to avoid false identifier matching."""
131
+ clean = re.sub(r"--[^\n]*\n?", " ", sql)
132
+ clean = re.sub(r"/\*.*?\*/", " ", clean, flags=re.DOTALL)
133
+ clean = re.sub(r"'(''|[^'])*'", "''", clean)
134
+ return clean
135
+
136
+ @classmethod
137
+ def _clean_identifier(cls, ident: str) -> str:
138
+ """Remove quotes, brackets, and schema qualifiers (e.g. `public."table"` -> `table`)."""
139
+ ident = ident.strip()
140
+ if "." in ident:
141
+ ident = ident.split(".")[-1]
142
+ ident = ident.strip('"`[] \t\n;')
143
+ return ident
144
+
145
+ @classmethod
146
+ def validate_sql_against_schema(
147
+ cls, sql: str, schema: Optional[DatabaseSchema]
148
+ ) -> Tuple[bool, List[str]]:
149
+ """
150
+ Validates that all tables, columns, primary keys, and foreign keys referenced
151
+ in an SQL statement or script exist in the physical database schema metadata.
152
+
153
+ Returns:
154
+ (is_valid, list_of_error_messages)
155
+ """
156
+ if not schema or not schema.tables:
157
+ return True, []
158
+
159
+ if not sql or not sql.strip():
160
+ return True, []
161
+
162
+ from dbagent.sql.validator import SQLValidator
163
+
164
+ statements = SQLValidator.split_statements(sql)
165
+ errors: List[str] = []
166
+
167
+ schema_table_map: Dict[str, TableModel] = {t.name.lower(): t for t in schema.tables}
168
+ sql_keywords = {
169
+ "select", "from", "where", "join", "inner", "left", "right", "full", "cross", "outer",
170
+ "on", "using", "group", "by", "order", "having", "limit", "offset", "insert", "into",
171
+ "values", "update", "set", "delete", "create", "alter", "drop", "table", "index", "view",
172
+ "and", "or", "not", "in", "is", "null", "like", "ilike", "between", "exists", "case",
173
+ "when", "then", "else", "end", "as", "distinct", "all", "any", "some", "union",
174
+ "intersect", "except", "with", "recursive", "returning", "asc", "desc", "true", "false",
175
+ "count", "sum", "avg", "min", "max", "coalesce", "cast", "extract", "now", "date",
176
+ "current_date", "current_timestamp", "upper", "lower", "trim", "concat", "row_number",
177
+ }
178
+
179
+ for stmt in statements:
180
+ clean_stmt = cls._clean_sql(stmt)
181
+ table_alias_map: Dict[str, TableModel] = {}
182
+ stmt_tables: List[TableModel] = []
183
+
184
+ # 1. Discover all tables and aliases in statement
185
+ table_ref_patterns = [
186
+ r"\bFROM\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
187
+ r"\bJOIN\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
188
+ r"\bUPDATE\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
189
+ r"\bINSERT\s+INTO\s+([a-zA-Z0-9_\"`.]+)",
190
+ r"\bDELETE\s+FROM\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
191
+ r"\bMERGE\s+INTO\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?",
192
+ ]
193
+
194
+ for pat in table_ref_patterns:
195
+ for match in re.finditer(pat, clean_stmt, re.IGNORECASE):
196
+ raw_table = cls._clean_identifier(match.group(1))
197
+ raw_alias = cls._clean_identifier(match.group(2)) if match.lastindex and match.lastindex >= 2 and match.group(2) else None
198
+ if not raw_table or raw_table.lower() in sql_keywords:
199
+ continue
200
+
201
+ tbl_lower = raw_table.lower()
202
+ if tbl_lower in schema_table_map:
203
+ t_model = schema_table_map[tbl_lower]
204
+ if t_model not in stmt_tables:
205
+ stmt_tables.append(t_model)
206
+ table_alias_map[tbl_lower] = t_model
207
+ table_alias_map[t_model.name.lower()] = t_model
208
+ if raw_alias:
209
+ alias_clean = raw_alias.lower()
210
+ if alias_clean not in sql_keywords:
211
+ table_alias_map[alias_clean] = t_model
212
+ else:
213
+ # Table not in schema - check for near matches
214
+ similar = [
215
+ s_name for s_name in schema_table_map
216
+ if raw_table.lower() in s_name or s_name in raw_table.lower()
217
+ ]
218
+ hint = f" (Did you mean: {', '.join(similar[:3])}?)" if similar else ""
219
+ errors.append(f"Table '{raw_table}' does not exist in the database schema.{hint}")
220
+
221
+ if not stmt_tables:
222
+ continue
223
+
224
+ # 2. Check UPDATE statements for invalid column assignments in SET clause
225
+ update_match = re.search(
226
+ r"\bUPDATE\s+([a-zA-Z0-9_\"`.]+)(?:\s+(?:AS\s+)?([a-zA-Z0-9_\"`]+))?\s+SET\s+(.*?)(?:\s+WHERE|\s+RETURNING|;|\Z)",
227
+ clean_stmt,
228
+ re.IGNORECASE | re.DOTALL,
229
+ )
230
+ if update_match:
231
+ upd_table_name = cls._clean_identifier(update_match.group(1)).lower()
232
+ target_table = schema_table_map.get(upd_table_name) or (stmt_tables[0] if stmt_tables else None)
233
+ if target_table:
234
+ set_clause = update_match.group(3)
235
+ assigns = re.findall(r"([a-zA-Z0-9_\"`.]+)\s*=", set_clause)
236
+ for col_expr in assigns:
237
+ col_clean = cls._clean_identifier(col_expr).lower()
238
+ col_model_map = {c.name.lower(): c for c in target_table.columns}
239
+ if col_clean not in col_model_map and col_clean not in sql_keywords:
240
+ avail = ", ".join(c.name for c in target_table.columns[:10])
241
+ pk_str = ", ".join(target_table.primary_key) if isinstance(target_table.primary_key, list) else str(target_table.primary_key or "")
242
+ pk_hint = f" Primary key is '{pk_str}'." if pk_str else ""
243
+ fk_cols = [fk.constrained_columns[0] for fk in target_table.foreign_keys if fk.constrained_columns]
244
+ fk_hint = f" Valid foreign keys: [{', '.join(fk_cols)}]." if fk_cols else ""
245
+ errors.append(
246
+ f"Column '{col_clean}' does not exist in table '{target_table.name}'. Available columns: [{avail}].{pk_hint}{fk_hint}"
247
+ )
248
+
249
+ # 3. Check INSERT statements for invalid column names in column list
250
+ insert_match = re.search(
251
+ r"\bINSERT\s+INTO\s+([a-zA-Z0-9_\"`.]+)\s*\((.*?)\)",
252
+ clean_stmt,
253
+ re.IGNORECASE | re.DOTALL,
254
+ )
255
+ if insert_match:
256
+ ins_table_name = cls._clean_identifier(insert_match.group(1)).lower()
257
+ target_table = schema_table_map.get(ins_table_name)
258
+ if target_table:
259
+ col_list_str = insert_match.group(2)
260
+ ins_cols = [cls._clean_identifier(c) for c in col_list_str.split(",") if c.strip()]
261
+ col_model_map = {c.name.lower(): c for c in target_table.columns}
262
+ for ic in ins_cols:
263
+ if ic.lower() not in col_model_map and ic.lower() not in sql_keywords:
264
+ avail = ", ".join(c.name for c in target_table.columns[:10])
265
+ errors.append(
266
+ f"Column '{ic}' does not exist in table '{target_table.name}'. Available columns: [{avail}]."
267
+ )
268
+
269
+ # 4. Check qualified column references (alias.column or table.column)
270
+ qualified_cols = re.findall(r"\b([a-zA-Z0-9_]+)\.([a-zA-Z0-9_]+)\b", clean_stmt)
271
+ for tbl_or_alias, col_name in qualified_cols:
272
+ tbl_or_alias_clean = tbl_or_alias.lower()
273
+ col_clean = col_name.lower()
274
+ if tbl_or_alias_clean in table_alias_map and col_clean != "*" and col_clean not in sql_keywords:
275
+ t_model = table_alias_map[tbl_or_alias_clean]
276
+ col_map = {c.name.lower(): c for c in t_model.columns}
277
+ if col_clean not in col_map:
278
+ avail = ", ".join(c.name for c in t_model.columns[:10])
279
+ pk_str = ", ".join(t_model.primary_key) if isinstance(t_model.primary_key, list) else str(t_model.primary_key or "")
280
+ pk_hint = f" Primary key: '{pk_str}'." if pk_str else ""
281
+ fk_cols = [fk.constrained_columns[0] for fk in t_model.foreign_keys if fk.constrained_columns]
282
+ fk_hint = f" Foreign keys: [{', '.join(fk_cols)}]." if fk_cols else ""
283
+ errors.append(
284
+ f"Column '{col_name}' does not exist in table '{t_model.name}'. Available columns: [{avail}].{pk_hint}{fk_hint}"
285
+ )
286
+
287
+ # 5. Check unqualified column references in WHERE clauses
288
+ where_match = re.search(
289
+ r"\bWHERE\s+(.*?)(?:\s+ORDER\s+BY|\s+GROUP\s+BY|\s+LIMIT|\s+RETURNING|;|\Z)",
290
+ clean_stmt,
291
+ re.IGNORECASE | re.DOTALL,
292
+ )
293
+ if where_match:
294
+ where_clause = where_match.group(1)
295
+ # Find identifiers on the LHS of comparisons: col = val, col IN (...), etc.
296
+ lhs_idents = re.findall(
297
+ r"(?:^|\s+)([a-zA-Z0-9_]+)\s*(?:=|!=|<>|<=|>=|<|>|\bLIKE\b|\bILIKE\b|\bIS\b|\bIN\b|\bBETWEEN\b)",
298
+ where_clause,
299
+ re.IGNORECASE,
300
+ )
301
+ for ident in lhs_idents:
302
+ ident_clean = ident.strip().lower()
303
+ if ident_clean in sql_keywords or ident_clean.isdigit():
304
+ continue
305
+ if "." in ident:
306
+ continue # Handled by qualified check
307
+
308
+ if len(stmt_tables) == 1:
309
+ target_t = stmt_tables[0]
310
+ col_map = {c.name.lower(): c for c in target_t.columns}
311
+ if ident_clean not in col_map and ident_clean not in table_alias_map:
312
+ avail = ", ".join(c.name for c in target_t.columns[:10])
313
+ pk_str = ", ".join(target_t.primary_key) if isinstance(target_t.primary_key, list) else str(target_t.primary_key or "")
314
+ pk_hint = f" Primary key is '{pk_str}'." if pk_str else ""
315
+ errors.append(
316
+ f"Column '{ident}' in WHERE clause does not exist in table '{target_t.name}'. Available columns: [{avail}].{pk_hint}"
317
+ )
318
+ elif len(stmt_tables) > 1:
319
+ # Multi-table query: check if column exists in ANY of the statement's tables
320
+ all_stmt_cols = set()
321
+ for t in stmt_tables:
322
+ for c in t.columns:
323
+ all_stmt_cols.add(c.name.lower())
324
+ if ident_clean not in all_stmt_cols and ident_clean not in table_alias_map:
325
+ tables_str = ", ".join(t.name for t in stmt_tables)
326
+ errors.append(
327
+ f"Column '{ident}' in WHERE clause does not exist in any of the referenced tables: [{tables_str}]."
328
+ )
329
+
330
+ # Deduplicate error messages while preserving order
331
+ deduped_errors: List[str] = []
332
+ for e in errors:
333
+ if e not in deduped_errors:
334
+ deduped_errors.append(e)
335
+
336
+ is_valid = len(deduped_errors) == 0
337
+ return is_valid, deduped_errors
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.11.2
3
+ Version: 0.11.4
4
4
  Summary: The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)
5
5
  Author: Santhosh Gupta
6
6
  License: MIT
@@ -81,6 +81,7 @@ tests/test_doctor_and_viz.py
81
81
  tests/test_formats_and_exit_codes.py
82
82
  tests/test_generator.py
83
83
  tests/test_graph.py
84
+ tests/test_ground_truth_sql_validation.py
84
85
  tests/test_knowledge.py
85
86
  tests/test_multi_turn_conversation.py
86
87
  tests/test_plugins.py
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "dbagent-cli"
7
- version = "0.11.2"
7
+ version = "0.11.4"
8
8
  description = "The Privacy-First AI Database Copilot, Schema Graph Introspector & Post-Commit Recovery Engine (CLI)"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.9"
@@ -11,7 +11,7 @@ from dbagent.llm.mock_provider import MockProvider
11
11
  def test_is_direct_sql_detection():
12
12
  # Direct SQL queries
13
13
  assert SQLValidator.is_direct_sql("SELECT * FROM users;") is True
14
- assert SQLValidator.is_direct_sql("select pk_user_id, username from ivis_user where active = 'Y'") is True
14
+ assert SQLValidator.is_direct_sql("select id, username from users where active = 'Y'") is True
15
15
  assert SQLValidator.is_direct_sql("SELECT count(*) FROM orders") is True
16
16
  assert SQLValidator.is_direct_sql("SELECT 1;") is True
17
17
  assert SQLValidator.is_direct_sql("WITH cte AS (SELECT 1) SELECT * FROM cte") is True
@@ -74,7 +74,7 @@ def test_validator_destructive_delete_without_where():
74
74
 
75
75
 
76
76
  def test_classify_intent_analysis():
77
- prompt = "analyze the database and give me the linkage architecture that how ivis_user and ivis_tenant was related and how they were mapped and how to know that which sites in ivis_site were mapped to which user"
77
+ prompt = "analyze the database and give me the linkage architecture of how users and accounts are related and how they are mapped and how to know which devices in customer_devices are mapped to which user"
78
78
  assert ScriptGenerator.classify_intent(prompt) == "analysis"
79
79
 
80
80
  assert ScriptGenerator.classify_intent("how are users and orders related?") == "analysis"
@@ -0,0 +1,150 @@
1
+ """
2
+ Unit tests for Schema Ground-Truth SQL Validation and Auto-Correction.
3
+ Verifies table, column, PK, and FK validation against physical database schema using generic domain models.
4
+ """
5
+
6
+ import pytest
7
+ from unittest.mock import MagicMock
8
+ from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
9
+ from dbagent.schema.ground_truth import SchemaGroundTruthValidator
10
+ from dbagent.agent.pipeline import QueryPipeline, PipelineResult
11
+ from dbagent.security.modes import OperationMode
12
+
13
+
14
+ @pytest.fixture
15
+ def sample_schema():
16
+ """Create a sample generic schema with realistic primary keys and foreign keys."""
17
+ customers = TableModel(
18
+ name="customers",
19
+ columns=[
20
+ ColumnModel(name="customer_id", data_type="INTEGER", is_primary_key=True),
21
+ ColumnModel(name="customer_name", data_type="VARCHAR(255)"),
22
+ ColumnModel(name="account_id", data_type="INTEGER"),
23
+ ColumnModel(name="region_id", data_type="INTEGER"),
24
+ ColumnModel(name="is_active", data_type="BOOLEAN"),
25
+ ],
26
+ foreign_keys=[
27
+ ForeignKeyModel(
28
+ constrained_columns=["account_id"],
29
+ referred_table="accounts",
30
+ referred_columns=["account_id"],
31
+ ),
32
+ ForeignKeyModel(
33
+ constrained_columns=["region_id"],
34
+ referred_table="regions",
35
+ referred_columns=["region_id"],
36
+ ),
37
+ ],
38
+ primary_key=["customer_id"],
39
+ )
40
+
41
+ accounts = TableModel(
42
+ name="accounts",
43
+ columns=[
44
+ ColumnModel(name="account_id", data_type="INTEGER", is_primary_key=True),
45
+ ColumnModel(name="account_name", data_type="VARCHAR(255)"),
46
+ ColumnModel(name="billing_code", data_type="VARCHAR(50)"),
47
+ ],
48
+ foreign_keys=[],
49
+ primary_key=["account_id"],
50
+ )
51
+
52
+ regions = TableModel(
53
+ name="regions",
54
+ columns=[
55
+ ColumnModel(name="region_id", data_type="INTEGER", is_primary_key=True),
56
+ ColumnModel(name="region_name", data_type="VARCHAR(255)"),
57
+ ],
58
+ foreign_keys=[],
59
+ primary_key=["region_id"],
60
+ )
61
+
62
+ return DatabaseSchema(
63
+ database_name="sample_db",
64
+ dialect_name="postgresql",
65
+ tables=[customers, accounts, regions],
66
+ )
67
+
68
+
69
+ def test_valid_select_query_passes_ground_truth(sample_schema):
70
+ sql = "SELECT c.customer_id, c.customer_name, a.account_name FROM customers c JOIN accounts a ON c.account_id = a.account_id WHERE c.is_active = true;"
71
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
72
+ assert is_valid is True
73
+ assert len(errors) == 0
74
+
75
+
76
+ def test_valid_update_query_passes_ground_truth(sample_schema):
77
+ sql = "UPDATE customers SET account_id = 10, region_id = 20 WHERE customer_id = 100;"
78
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
79
+ assert is_valid is True
80
+ assert len(errors) == 0
81
+
82
+
83
+ def test_hallucinated_column_in_update_set_detected(sample_schema):
84
+ sql = "UPDATE customers SET project = 5, site_group = 10 WHERE customer_id = 100;"
85
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
86
+ assert is_valid is False
87
+ assert any("Column 'project' does not exist in table 'customers'" in e for e in errors)
88
+ assert any("Column 'site_group' does not exist in table 'customers'" in e for e in errors)
89
+ assert any("Valid foreign keys" in e for e in errors)
90
+
91
+
92
+ def test_hallucinated_column_in_where_clause_detected(sample_schema):
93
+ sql = "UPDATE customers SET account_id = 5 WHERE legacy_user_code = 100;"
94
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
95
+ assert is_valid is False
96
+ assert any("Column 'legacy_user_code' in WHERE clause does not exist in table 'customers'" in e for e in errors)
97
+ assert any("Primary key is 'customer_id'" in e for e in errors)
98
+
99
+
100
+ def test_hallucinated_table_name_detected(sample_schema):
101
+ sql = "SELECT * FROM non_existent_table WHERE id = 1;"
102
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
103
+ assert is_valid is False
104
+ assert any("Table 'non_existent_table' does not exist" in e for e in errors)
105
+
106
+
107
+ def test_hallucinated_qualified_column_detected(sample_schema):
108
+ sql = "SELECT c.customer_id, c.fake_column FROM customers c;"
109
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
110
+ assert is_valid is False
111
+ assert any("Column 'fake_column' does not exist in table 'customers'" in e for e in errors)
112
+
113
+
114
+ def test_multi_statement_validation(sample_schema):
115
+ sql = """
116
+ UPDATE customers SET account_id = 2 WHERE customer_id = 1;
117
+ UPDATE customers SET bad_col = 3 WHERE customer_id = 2;
118
+ """
119
+ is_valid, errors = SchemaGroundTruthValidator.validate_sql_against_schema(sql, sample_schema)
120
+ assert is_valid is False
121
+ assert any("Column 'bad_col' does not exist in table 'customers'" in e for e in errors)
122
+
123
+
124
+ def test_pipeline_auto_corrects_hallucinated_columns(sample_schema):
125
+ mock_connector = MagicMock()
126
+ mock_connector.get_table_names.return_value = ["customers", "accounts", "regions"]
127
+ mock_connector.resolve_tables.return_value = (["customers"], [])
128
+ mock_connector.inspect_schema.return_value = sample_schema
129
+ mock_connector.execute_query.return_value = (["customer_id", "account_id"], [{"customer_id": 1, "account_id": 2}], None)
130
+
131
+ mock_llm = MagicMock()
132
+ mock_llm.generate.side_effect = [
133
+ "UPDATE customers SET account = 2 WHERE cust_id = 1;",
134
+ "UPDATE customers SET account_id = 2 WHERE customer_id = 1;",
135
+ ]
136
+
137
+ pipeline = QueryPipeline(
138
+ connector=mock_connector,
139
+ llm=mock_llm,
140
+ mode=OperationMode.CONFIRM,
141
+ auto_execute=True,
142
+ )
143
+
144
+ result = pipeline.run("Change the account for customer 1 to account 2", force=True)
145
+
146
+ assert result.success is True
147
+ assert "account_id = 2" in result.sql
148
+ assert "customer_id = 1" in result.sql
149
+ assert result.retries >= 1
150
+ assert result.ground_truth is not None
@@ -17,27 +17,27 @@ def temp_kb(tmp_path):
17
17
  def test_knowledge_base_learn_and_find(temp_kb):
18
18
  rule = temp_kb.learn(
19
19
  db_name="test_db",
20
- topic="Site to User Group Mapping",
21
- rule_description="ivis_site maps to ivis_user through ivis_site_group -> ivis_group -> ivis_user_group_map",
22
- relevant_tables=["ivis_site", "ivis_site_group", "ivis_group", "ivis_user_group_map", "ivis_user"],
20
+ topic="User to Role Permission Mapping",
21
+ rule_description="users map to permissions through user_roles -> roles -> role_permissions",
22
+ relevant_tables=["users", "user_roles", "roles", "role_permissions", "permissions"],
23
23
  source="user_correction",
24
24
  )
25
25
  assert rule.rule_id.startswith("rule_")
26
- assert rule.topic == "Site to User Group Mapping"
26
+ assert rule.topic == "User to Role Permission Mapping"
27
27
 
28
28
  # List all
29
29
  all_rules = temp_kb.list_all("test_db")
30
30
  assert len(all_rules) == 1
31
31
 
32
32
  # Find relevant by prompt
33
- matched = temp_kb.find_relevant("test_db", "how to map sites to user")
33
+ matched = temp_kb.find_relevant("test_db", "how to map users to permissions")
34
34
  assert len(matched) == 1
35
- assert matched[0].topic == "Site to User Group Mapping"
35
+ assert matched[0].topic == "User to Role Permission Mapping"
36
36
 
37
37
  # Context formatting
38
- ctx = temp_kb.format_prompt_context("test_db", "how to map sites to user")
38
+ ctx = temp_kb.format_prompt_context("test_db", "how to map users to permissions")
39
39
  assert "Persistent Database Domain Knowledge" in ctx
40
- assert "Site to User Group Mapping" in ctx
40
+ assert "User to Role Permission Mapping" in ctx
41
41
 
42
42
  # Delete
43
43
  deleted = temp_kb.delete("test_db", rule.rule_id)
@@ -26,10 +26,10 @@ def test_prompt_queue_fifo():
26
26
 
27
27
  def test_split_and_enqueue_chained():
28
28
  qm = PromptQueueManager()
29
- raw = "how ivis_user and ivis_site were linked ; give the flow diagram ; count total users"
29
+ raw = "how users and orders were linked ; give the flow diagram ; count total users"
30
30
  first = qm.split_and_enqueue_chained(raw)
31
31
 
32
- assert first == "how ivis_user and ivis_site were linked"
32
+ assert first == "how users and orders were linked"
33
33
  assert qm.count() == 2
34
34
  assert qm.list_all() == ["give the flow diagram", "count total users"]
35
35
 
@@ -11,10 +11,10 @@ def test_queue_does_not_split_multiline_sql_with_semicolons():
11
11
  queue_mgr = PromptQueueManager()
12
12
 
13
13
  sql_input = """
14
- DELETE FROM ivis_customer_user_map WHERE fk_user_id = 5201;
15
- DELETE FROM ivis_user_access WHERE fk_user_id = 5201;
14
+ DELETE FROM customer_user_map WHERE user_id = 5201;
15
+ DELETE FROM user_access WHERE user_id = 5201;
16
16
  DELETE FROM user_activity WHERE user_id = 5201;
17
- DELETE FROM ivis_user WHERE pk_user_id = 5201;
17
+ DELETE FROM users WHERE user_id = 5201;
18
18
  COMMIT; correct this
19
19
  """
20
20
 
@@ -27,12 +27,12 @@ def test_classify_query_with_cte_delete():
27
27
  """Verify that data-modifying WITH CTE statements are classified as 'write'."""
28
28
  cte_delete = """
29
29
  WITH target_user AS (
30
- SELECT pk_user_id FROM ivis_user WHERE username = 'target_user'
30
+ SELECT id FROM users WHERE username = 'target_user'
31
31
  ),
32
32
  del_access AS (
33
- DELETE FROM ivis_user_access WHERE fk_user_id IN (SELECT pk_user_id FROM target_user)
33
+ DELETE FROM user_access WHERE user_id IN (SELECT id FROM target_user)
34
34
  )
35
- DELETE FROM ivis_user WHERE pk_user_id IN (SELECT pk_user_id FROM target_user);
35
+ DELETE FROM users WHERE id IN (SELECT id FROM target_user);
36
36
  """
37
37
  assert SQLValidator.classify_query(cte_delete) == "write"
38
38
 
@@ -41,7 +41,7 @@ def test_classify_query_with_cte_read():
41
41
  """Verify that regular SELECT WITH CTE statements remain classified as 'read'."""
42
42
  cte_read = """
43
43
  WITH active_users AS (
44
- SELECT pk_user_id, username FROM ivis_user WHERE active = 'Y'
44
+ SELECT id, username FROM users WHERE active = 'Y'
45
45
  )
46
46
  SELECT * FROM active_users LIMIT 10;
47
47
  """
@@ -1,126 +0,0 @@
1
- """
2
- Schema Ground-Truth Verification Engine.
3
- Fact-checks AI-generated architectural reports and relationship mappings against the physical database schema metadata:
4
- - Table existence verification
5
- - Column existence and data type verification
6
- - Relational foreign key constraint verification
7
- - Logical / semantic mapping verification
8
- - SQL query dialect validation
9
- """
10
-
11
- import re
12
- from typing import List, Dict, Any, Optional
13
- from dataclasses import dataclass, field
14
- from dbagent.schema.models import DatabaseSchema, TableModel, ColumnModel, ForeignKeyModel
15
-
16
-
17
- @dataclass
18
- class GroundTruthValidationResult:
19
- """Result of validating AI architectural response against database schema ground truth."""
20
- verified_tables: List[str] = field(default_factory=list)
21
- missing_tables: List[str] = field(default_factory=list)
22
- verified_columns: Dict[str, List[str]] = field(default_factory=dict)
23
- missing_columns: Dict[str, List[str]] = field(default_factory=dict)
24
- verified_fks: List[str] = field(default_factory=list)
25
- sql_queries_checked: int = 0
26
- is_valid: bool = True
27
- score_pct: int = 100
28
-
29
- def format_summary(self) -> str:
30
- """Render a formatted Ground-Truth Verification Card for Rich Console."""
31
- color = "green" if self.score_pct >= 90 else ("yellow" if self.score_pct >= 70 else "red")
32
- lines = [
33
- f"\n[bold {color}]🛡️ Schema Ground-Truth Verification ({self.score_pct}% Confirmed):[/bold {color}]"
34
- ]
35
- if self.verified_tables:
36
- lines.append(f" [bold green]✓[/bold green] [bold white]Tables Verified ({len(self.verified_tables)}):[/bold white] [cyan]{', '.join(self.verified_tables)}[/cyan]")
37
- if self.verified_fks:
38
- lines.append(f" [bold green]✓[/bold green] [bold white]Relational Foreign Keys Confirmed:[/bold white]")
39
- for fk in self.verified_fks[:4]:
40
- lines.append(f" [dim]•[/dim] {fk}")
41
- if len(self.verified_fks) > 4:
42
- lines.append(f" [dim]• ... and {len(self.verified_fks) - 4} more[/dim]")
43
- if self.sql_queries_checked > 0:
44
- lines.append(f" [bold green]✓[/bold green] [bold white]SQL Queries Verified:[/bold white] {self.sql_queries_checked} dialect statement(s) validated against schema")
45
- for missing in self.missing_tables:
46
- lines.append(f" [bold yellow]⚠ Note:[/bold yellow] Query referenced [dim]{missing}[/dim]")
47
- for tbl, cols in self.missing_columns.items():
48
- lines.append(f" [bold red]⚠ Warning:[/bold red] Column(s) [yellow]{', '.join(cols)}[/yellow] not found in table [cyan]{tbl}[/cyan]")
49
- return "\n".join(lines)
50
-
51
-
52
- class SchemaGroundTruthValidator:
53
- """Validates that all tables, columns, and relationships in an AI response match real database ground-truth."""
54
-
55
- @classmethod
56
- def validate_response(cls, response_text: str, schema: Optional[DatabaseSchema]) -> Optional[GroundTruthValidationResult]:
57
- if not schema or not schema.tables:
58
- return None
59
-
60
- result = GroundTruthValidationResult()
61
- schema_table_map = {t.name.lower(): t for t in schema.tables}
62
-
63
- # 1. Identify all database tables referenced in the response
64
- resp_lower = response_text.lower()
65
- for t_name, t_model in schema_table_map.items():
66
- if re.search(r'\b' + re.escape(t_name) + r'\b', resp_lower):
67
- result.verified_tables.append(t_model.name)
68
-
69
- # Check columns for this table
70
- col_map = {c.name.lower(): c.name for c in t_model.columns}
71
- found_cols = []
72
- for c_lower, c_orig in col_map.items():
73
- if len(c_lower) > 2 and re.search(r'\b' + re.escape(c_lower) + r'\b', resp_lower):
74
- found_cols.append(c_orig)
75
- if found_cols:
76
- result.verified_columns[t_model.name] = found_cols
77
-
78
- # Check foreign keys
79
- for fk in t_model.foreign_keys:
80
- ref_tbl = fk.referred_table
81
- constrained = ", ".join(fk.constrained_columns)
82
- ref_cols = ", ".join(fk.referred_columns)
83
- if ref_tbl and ref_tbl.lower() in resp_lower:
84
- fk_desc = f"[cyan]{t_model.name}[/cyan] ([dim]{constrained}[/dim]) ➔ [cyan]{ref_tbl}[/cyan] ([dim]{ref_cols}[/dim])"
85
- if fk_desc not in result.verified_fks:
86
- result.verified_fks.append(fk_desc)
87
-
88
- # 2. Extract and validate only actual SQL queries against physical schema
89
- sql_blocks = re.findall(r'```(?:sql)\s*([\s\S]*?)\s*```', response_text, re.IGNORECASE)
90
- result.sql_queries_checked = len(sql_blocks)
91
- sql_referenced_tables = set()
92
- sql_ignore_keywords = {"via", "lateral", "on", "using", "where", "select", "table", "as", "and", "or", "values"}
93
-
94
- for block in sql_blocks:
95
- # Extract FROM and JOIN table references
96
- matches = re.findall(r'\b(?:FROM|JOIN)\s+([a-zA-Z0-9_]+)', block, re.IGNORECASE)
97
- for m in matches:
98
- m_clean = m.strip().lower()
99
- if m_clean not in sql_ignore_keywords:
100
- sql_referenced_tables.add(m_clean)
101
-
102
- for sql_tbl in sql_referenced_tables:
103
- if sql_tbl in schema_table_map:
104
- if schema_table_map[sql_tbl].name not in result.verified_tables:
105
- result.verified_tables.append(schema_table_map[sql_tbl].name)
106
- else:
107
- # Check if it's a common alias or slight naming divergence
108
- found_match = None
109
- for real_tbl in schema_table_map:
110
- if sql_tbl in real_tbl or real_tbl in sql_tbl:
111
- found_match = schema_table_map[real_tbl].name
112
- break
113
- if found_match:
114
- result.missing_tables.append(f"{sql_tbl} ➔ (Actual: {found_match})")
115
- else:
116
- result.missing_tables.append(sql_tbl)
117
-
118
- total_checked = len(result.verified_tables)
119
- if result.missing_tables:
120
- result.score_pct = max(70, 100 - (len(result.missing_tables) * 15))
121
- elif total_checked > 0:
122
- result.score_pct = 100
123
- else:
124
- result.score_pct = 95
125
-
126
- return result
File without changes
File without changes