File size: 6,718 Bytes
939c0c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """
SQL Intelligence Specialist Agent Node
=======================================
Translates natural language queries into executable SQL commands.
Enforces strict security guardrails (SELECT only; blocks DDL/DML mutations).
Formats tabular SQL results for consumption by the Synthesizer node.
"""
from __future__ import annotations
import re
import structlog
from typing import Any, Dict, List, Tuple
from sqlalchemy import inspect, text
from app.core.config import settings
from app.core.database import AsyncSessionLocal, engine
from app.services.llm_gateway import llm_gateway
from agents.state import CopilotState
logger = structlog.get_logger(__name__)
# Forbidden SQL keywords that alter database state
FORBIDDEN_KEYWORDS = {
"INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE",
"CREATE", "GRANT", "REVOKE", "EXEC", "EXECUTE", "MERGE"
}
TEXT_TO_SQL_PROMPT = """You are a Database Expert AI Agent.
Your job is to translate the user's natural language request into a valid, read-only SQL query.
Database Dialect: {dialect}
Database Schema:
{schema_description}
Rules:
1. Generate ONLY a single SELECT query. Never attempt to INSERT, UPDATE, DELETE, DROP, or ALTER.
2. Return ONLY the raw SQL statement inside ```sql ... ``` block or plain text without explanations.
3. Use proper joins and table aliases if referencing multiple tables.
4. Restrict query results with LIMIT 50 if returning many rows.
"""
def is_safe_sql(query_sql: str) -> Tuple[bool, str]:
"""
Validate that the generated SQL is strictly read-only (SELECT).
"""
cleaned = re.sub(r"--.*?\n|/\*.*?\*/", "", query_sql, flags=re.DOTALL).strip()
uppercase_words = set(re.findall(r"\b[A-Z]+\b", cleaned.upper()))
forbidden_found = uppercase_words.intersection(FORBIDDEN_KEYWORDS)
if forbidden_found:
return False, f"Security Violation: Query contains forbidden statement(s) {forbidden_found}"
if not cleaned.upper().startswith("SELECT") and not cleaned.upper().startswith("WITH"):
return False, "Security Violation: Only SELECT queries are permitted"
return True, "Safe"
async def get_database_schema() -> str:
"""Extract schema tables and columns for LLM prompt context."""
schema_info = []
try:
async with engine.connect() as conn:
def _inspect_schema(connection):
inspector = inspect(connection)
tables = inspector.get_table_names()
res = []
for table in tables[:10]: # inspect first 10 tables
columns = inspector.get_columns(table)
col_str = ", ".join([f"{c['name']} ({c['type']})" for c in columns])
res.append(f"Table '{table}': {col_str}")
return "\n".join(res)
schema_info = await conn.run_sync(_inspect_schema)
except Exception as e:
logger.warning("Could not introspect database schema", error=str(e))
schema_info = "Tables: users (id, email, full_name, role), tenants (id, name, slug), documents (id, filename, doc_type, status, chunk_count)"
return schema_info or "No tables found"
async def execute_sql_query(query_sql: str) -> Tuple[List[str], List[Dict[str, Any]]]:
"""Execute safe read-only SQL query and return (columns, rows)."""
async with AsyncSessionLocal() as session:
result = await session.execute(text(query_sql))
columns = list(result.keys())
rows = [dict(zip(columns, row)) for row in result.fetchall()]
return columns, rows
async def sql_node(state: CopilotState) -> CopilotState:
"""
SQL Specialist Node.
Generates, validates, and executes read-only SQL.
"""
query = state.get("query", "")
logger.info("SQL Agent executing", query=query)
try:
# 1. Introspect Schema
schema_str = await get_database_schema()
dialect = "SQLite" if settings.use_sqlite else "PostgreSQL"
# 2. Text-to-SQL LLM Prompt
prompt = f"User Request: {query}\n\nGenerated SQL:"
system_prompt = TEXT_TO_SQL_PROMPT.format(dialect=dialect, schema_description=schema_str)
response = await llm_gateway.generate(
prompt=prompt,
system_prompt=system_prompt,
temperature=0.0,
)
sql_text = response.content.strip()
# Extract SQL from markdown code block if present
if "```" in sql_text:
sql_text = sql_text.split("```")[1]
if sql_text.lower().startswith("sql"):
sql_text = sql_text[3:].strip()
logger.info("Generated SQL statement", sql=sql_text)
# 3. Security Validation
safe, reason = is_safe_sql(sql_text)
if not safe:
logger.warning("SQL execution blocked by security guardrail", reason=reason)
state["error"] = reason
state["retrieved_chunks"] = [{
"document_name": "SQL Guardrail",
"text": f"⚠️ Query blocked: {reason}. Only read-only SELECT queries are allowed."
}]
return state
# 4. Execute Query
columns, rows = await execute_sql_query(sql_text)
logger.info("SQL Query executed successfully", row_count=len(rows))
# Format tabular output into markdown text for Synthesizer
formatted_table = ""
if rows:
headers = " | ".join(columns)
formatted_table += f"| {headers} |\n"
formatted_table += f"| {' | '.join(['---'] * len(columns))} |\n"
for row in rows[:20]:
formatted_table += f"| {' | '.join(str(row.get(c, '')) for c in columns)} |\n"
else:
formatted_table = "Query returned 0 rows."
# Update state
state["retrieved_chunks"] = [{
"document_id": "sql_result",
"document_name": f"SQL Query ({sql_text})",
"text": f"Executed SQL: `{sql_text}`\n\nResults:\n{formatted_table}",
"score": 1.0,
"doc_type": "sql"
}]
outputs = state.get("agent_outputs", [])
outputs.append({
"agent_name": "sql",
"content": f"Executed SQL query: `{sql_text}`. Returned {len(rows)} rows.",
"metadata": {"sql": sql_text, "row_count": len(rows)}
})
state["agent_outputs"] = outputs
state["active_agent"] = "sql"
except Exception as e:
logger.error("SQL Agent execution error", error=str(e))
state["error"] = f"SQL Error: {str(e)}"
state["retrieved_chunks"] = [{
"document_name": "SQL Error",
"text": f"Failed to execute SQL query: {str(e)}"
}]
return state
|