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.
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/PKG-INFO +1 -1
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/__init__.py +1 -1
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/generator.py +2 -2
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/pipeline.py +44 -3
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/cli.py +3 -3
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/base.py +2 -2
- dbagent_cli-0.11.4/dbagent/schema/ground_truth.py +337 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/PKG-INFO +1 -1
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/SOURCES.txt +1 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/pyproject.toml +1 -1
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_direct_sql.py +1 -1
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_generator.py +1 -1
- dbagent_cli-0.11.4/tests/test_ground_truth_sql_validation.py +150 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_knowledge.py +8 -8
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_prompt_queue.py +2 -2
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_prompt_queue_and_sql.py +7 -7
- dbagent_cli-0.11.2/dbagent/schema/ground_truth.py +0 -126
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/README.md +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/confidence.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/conversation.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/doctor.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/error_classifier.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/recommender.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/agent/validator.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/benchmark.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/config.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/factory.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/mongo.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/connectors/relational.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/errors.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/__init__.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/cache.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/changeset.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/cost.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/prompt_queue.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/recovery.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/sandbox.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/execution/streaming.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/anthropic_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/auto_setup.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/base.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/deepseek_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/factory.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/gemini_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/groq_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/mistral_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/mock_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/ollama_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/openai_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/llm/openrouter_provider.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/plugins.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sample_db.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/cache.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/formatter.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/graph.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/models.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/retriever.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/schema/selector.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/__init__.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/audit.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/environment.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/masking.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/modes.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/security/rbac.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/__init__.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/dictionary.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/knowledge.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/semantic/templates.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/__init__.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/optimizer.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/sql/validator.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/chat_screen.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/console.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/grid_window.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/live_streamer.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/viewer.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent/ui/visualizer.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/dependency_links.txt +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/entry_points.txt +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/requires.txt +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/dbagent_cli.egg-info/top_level.txt +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/setup.cfg +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_benchmarks.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_cache.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_cli.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_connectors.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_demo_and_sample_db.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_doctor_and_viz.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_formats_and_exit_codes.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_graph.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_multi_turn_conversation.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_plugins.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_providers.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_recovery.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_safety.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_sandbox.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_schema.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_security.py +0 -0
- {dbagent_cli-0.11.2 → dbagent_cli-0.11.4}/tests/test_semantic.py +0 -0
|
@@ -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
|
|
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
|
|
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
|
-
|
|
138
|
-
|
|
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'^(?:
|
|
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
|
|
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
|
|
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
|
|
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. "
|
|
91
|
-
unprefixed = re.sub(r'^(?:
|
|
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
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "dbagent-cli"
|
|
7
|
-
version = "0.11.
|
|
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
|
|
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
|
|
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="
|
|
21
|
-
rule_description="
|
|
22
|
-
relevant_tables=["
|
|
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 == "
|
|
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
|
|
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 == "
|
|
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
|
|
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 "
|
|
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
|
|
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
|
|
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
|
|
15
|
-
DELETE FROM
|
|
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
|
|
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
|
|
30
|
+
SELECT id FROM users WHERE username = 'target_user'
|
|
31
31
|
),
|
|
32
32
|
del_access AS (
|
|
33
|
-
DELETE FROM
|
|
33
|
+
DELETE FROM user_access WHERE user_id IN (SELECT id FROM target_user)
|
|
34
34
|
)
|
|
35
|
-
DELETE FROM
|
|
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
|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|