RAG for Structured Data: Natural Language to SQL Guide

from openai import OpenAI
import asyncpg

client = OpenAI()

async def text_to_sql(question: str, schema: str, conn: asyncpg.Connection) -> dict:
    """Convert natural language to SQL with schema context."""
    sql_prompt = f"""You are a SQL expert. Generate a SQL query to answer the question.

Database Schema:
{schema}

Rules:
- Use only tables and columns from the schema
- Return SELECT queries only (no INSERT, UPDATE, DELETE)
- Use proper joins and WHERE clauses
- Include LIMIT for safety

Question: {question}

Return JSON: {{"sql": str, "reasoning": str}}
"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": sql_prompt}],
        response_format={"type": "json_object"},
        temperature=0.1
    )

    result = json.loads(response.choices[0].message.content)
    sql_query = result["sql"]

    # Step 2: Execute SQL
    try:
        rows = await conn.fetch(sql_query)

        # Step 3: Format results
        return {
            "sql": sql_query,
            "rows": [dict(r) for r in rows],
            "reasoning": result["reasoning"],
            "success": True,
        }

    except Exception as e:
        return {
            "sql": sql_query,
            "error": str(e),
            "success": False,
        }

# Usage
schema = """
Table: orders
- order_id: int (PK)
- customer_id: int (FK)
- order_date: date
- total_amount: decimal

Table: customers
- customer_id: int (PK)
- name: varchar
- region: varchar
"""

result = await text_to_sql("Show total revenue by region in 2025", schema, conn)
print(result["sql"])
# SELECT region, SUM(total_amount) as revenue
# FROM orders JOIN customers ON orders.customer_id = customers.customer_id
# WHERE EXTRACT(YEAR FROM order_date) = 2025
# GROUP BY region

from typing import List, Dict

def extract_schema_metadata(conn: asyncpg.Connection) -> List[Dict]:
    """Extract tables and columns with descriptions."""

    # Query information schema
    schema_query = """
    SELECT 
        table_name,
        column_name,
        data_type,
        col_description((table_schema||'.'||table_name)::regclass::oid, ordinal_position) as description
    FROM information_schema.columns
    WHERE table_schema = 'public'
    ORDER BY table_name, ordinal_position
    """

    rows = await conn.fetch(schema_query)

    # Group by table
    tables = {}
    for row in rows:
        table = row["table_name"]
        if table not in tables:
            tables[table] = {
                "table_name": table,
                "columns": [],
                "description": "",  # Add table-level descriptions if available
            }

        tables[table]["columns"].append({
            "column_name": row["column_name"],
            "data_type": row["data_type"],
            "description": row["description"] or "",
        })

    return list(tables.values())

def format_table_for_embedding(table: Dict) -> str:
    """Format table metadata as text for embedding."""
    lines = [f"Table: {table['table_name']}"]

    if table.get("description"):
        lines.append(f"Description: {table['description']}")

    lines.append("Columns:")
    for col in table["columns"]:
        col_desc = f"  - {col['column_name']} ({col['data_type']})"
        if col.get("description"):
            col_desc += f": {col['description']}"
        lines.append(col_desc)

    return "\n".join(lines)

# Embed schemas
tables = await extract_schema_metadata(conn)

for table in tables:
    text = format_table_for_embedding(table)
    embedding = client.embeddings.create(
        input=[text],
        model="text-embedding-3-small"
    ).data[0].embedding

    # Store in vector DB
    await store_schema_embedding(table["table_name"], text, embedding)

async def retrieve_relevant_schema(question: str, conn: asyncpg.Connection, top_k: int = 3) -> str:
    """Retrieve relevant tables for the question."""

    # Embed question
    query_embedding = client.embeddings.create(
        input=[question],
        model="text-embedding-3-small"
    ).data[0].embedding

    # Retrieve top-k relevant tables
    await register_vector(conn)
    relevant_tables = await conn.fetch(
        """
        SELECT schema_text, 1 - (embedding <=> $1) AS similarity
        FROM schema_embeddings
        ORDER BY embedding <=> $1
        LIMIT $2
        """,
        query_embedding,
        top_k
    )

    # Combine retrieved schemas
    schema_context = "\n\n".join(t["schema_text"] for t in relevant_tables)

    return schema_context

# Usage
question = "Show revenue by customer region last quarter"
relevant_schema = await retrieve_relevant_schema(question, conn)

# Now pass only relevant schema to text-to-SQL
result = await text_to_sql(question, relevant_schema, conn)

async def hybrid_rag_query(question: str, conn: asyncpg.Connection) -> dict:
    """Query both documents and structured data."""

    # Step 1: Classify query type
    classification = await classify_query_type(question)

    unstructured_results = []
    structured_results = []

    # Step 2: Retrieve from unstructured sources (documents)
    if classification["needs_documents"]:
        doc_chunks = await retrieve_document_chunks(question, conn)
        unstructured_results = [c["content"] for c in doc_chunks]

    # Step 3: Query structured data (SQL)
    if classification["needs_database"]:
        schema = await retrieve_relevant_schema(question, conn)
        sql_result = await text_to_sql(question, schema, conn)

        if sql_result["success"]:
            structured_results = sql_result["rows"]

    # Step 4: Synthesize answer from both sources
    answer = await synthesize_answer(
        question=question,
        unstructured=unstructured_results,
        structured=structured_results
    )

    return {
        "answer": answer,
        "unstructured_sources": unstructured_results,
        "structured_query": sql_result.get("sql") if classification["needs_database"] else None,
        "structured_data": structured_results,
    }

async def classify_query_type(question: str) -> dict:
    """Classify if query needs documents, database, or both."""

    classify_prompt = f"""Classify the query:

Question: {question}

Return JSON:
{{
  "needs_documents": bool,
  "needs_database": bool,
  "reasoning": str
}}
"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": classify_prompt}],
        response_format={"type": "json_object"}
    )

    return json.loads(response.choices[0].message.content)

async def synthesize_answer(question: str, unstructured: List[str], structured: List[Dict]) -> str:
    """Synthesize answer from both document and database results."""

    context_parts = []

    if unstructured:
        context_parts.append("Document Context:\n" + "\n\n---\n\n".join(unstructured))

    if structured:
        # Format structured data as markdown table
        if structured:
            headers = list(structured[0].keys())
            table = "| " + " | ".join(headers) + " |\n"
            table += "| " + " | ".join(["---"] * len(headers)) + " |\n"
            for row in structured:
                table += "| " + " | ".join(str(row[h]) for h in headers) + " |\n"

            context_parts.append(f"Database Results:\n{table}")

    combined_context = "\n\n".join(context_parts)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Synthesize an answer using document and database results."},
            {"role": "user", "content": f"{combined_context}\n\nQuestion: {question}"}
        ],
        temperature=0.1
    )

    return response.choices[0].message.content

# Example queries
# "What caused the Q4 revenue decline?" → Documents + SQL for revenue data
# "Show my last 5 orders" → SQL only
# "Explain the refund policy" → Documents only

from dataclasses import dataclass
from typing import Optional

@dataclass
class SQLConfig:
    max_retries: int = 2
    enable_query_validation: bool = True
    enable_row_level_security: bool = True
    max_rows: int = 1000

class ProductionSQLRAG:
    def __init__(self, conn: asyncpg.Connection, config: SQLConfig):
        self.conn = conn
        self.config = config

    async def query(self, question: str, tenant_id: str, user_id: str) -> dict:
        """Production SQL RAG with error recovery."""

        attempt = 0

        while attempt < self.config.max_retries:
            try:
                # Step 1: Retrieve relevant schema
                schema = await retrieve_relevant_schema(question, self.conn)

                # Step 2: Generate SQL
                sql_result = await self._generate_sql(question, schema)

                # Step 3: Validate SQL
                if self.config.enable_query_validation:
                    validation = await self._validate_sql(sql_result["sql"], tenant_id)
                    if not validation["safe"]:
                        return {
                            "error": f"Unsafe query: {validation['reason']}",
                            "success": False,
                        }

                # Step 4: Execute with RLS
                if self.config.enable_row_level_security:
                    await self.conn.execute("SET app.current_tenant_id = $1", tenant_id)

                rows = await self.conn.fetch(sql_result["sql"])

                # Step 5: Format answer
                answer = await self._format_results(question, rows)

                return {
                    "answer": answer,
                    "sql": sql_result["sql"],
                    "rows": [dict(r) for r in rows[:self.config.max_rows]],
                    "success": True,
                }

            except asyncpg.UndefinedColumnError as e:
                # Column doesn't exist, retry with error feedback
                print(f"Attempt {attempt+1}: Column error - {e}")

                if attempt < self.config.max_retries - 1:
                    # Retry with error context
                    question += f"\n\nPrevious error: {e}. Please use correct column names."
                    attempt += 1
                    continue
                else:
                    return {"error": str(e), "success": False}

            except Exception as e:
                return {"error": str(e), "success": False}

            finally:
                if self.config.enable_row_level_security:
                    await self.conn.execute("RESET app.current_tenant_id")

        return {"error": "Max retries exceeded", "success": False}

    async def _generate_sql(self, question: str, schema: str) -> dict:
        """Generate SQL with schema context."""
        return await text_to_sql(question, schema, self.conn)

    async def _validate_sql(self, sql: str, tenant_id: str) -> dict:
        """Validate SQL for safety."""

        # Check for destructive operations
        destructive_keywords = ["DROP", "DELETE", "TRUNCATE", "ALTER", "UPDATE", "INSERT"]

        sql_upper = sql.upper()
        for keyword in destructive_keywords:
            if keyword in sql_upper:
                return {
                    "safe": False,
                    "reason": f"Destructive operation not allowed: {keyword}",
                }

        # Check for proper tenant filtering (if multi-tenant)
        if "tenant_id" not in sql.lower():
            return {
                "safe": False,
                "reason": "Query must filter by tenant_id",
            }

        return {"safe": True}

    async def _format_results(self, question: str, rows: List) -> str:
        """Format SQL results as natural language."""

        if not rows:
            return "No results found."

        # Convert to markdown table
        headers = list(rows[0].keys())
        table = "| " + " | ".join(headers) + " |\n"
        table += "| " + " | ".join(["---"] * len(headers)) + " |\n"

        for row in rows[:20]:  # Limit table size
            table += "| " + " | ".join(str(row[h]) for h in headers) + " |\n"

        if len(rows) > 20:
            table += f"\n... and {len(rows) - 20} more rows"

        # Generate natural language summary
        summary_prompt = f"""Summarize these SQL results in natural language.

Question: {question}

Results (first 20 rows):
{table}

Total rows: {len(rows)}
"""

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": summary_prompt}],
            temperature=0.1
        )

        return response.choices[0].message.content + f"\n\n{table}"

def comprehensive_sql_validation(sql: str, user_permissions: dict) -> dict:
    """Comprehensive SQL safety validation."""

    issues = []

    # Check 1: Read-only operations
    if not is_select_only(sql):
        issues.append("Only SELECT queries allowed")

    # Check 2: No dangerous functions
    dangerous_funcs = ["pg_sleep", "pg_read_file", "pg_write_file", "COPY"]
    for func in dangerous_funcs:
        if func.lower() in sql.lower():
            issues.append(f"Dangerous function not allowed: {func}")

    # Check 3: Row limit enforced
    if "LIMIT" not in sql.upper():
        issues.append("Query must include LIMIT clause")

    # Check 4: Permissions check (table access)
    tables = extract_tables_from_sql(sql)
    for table in tables:
        if table not in user_permissions.get("allowed_tables", []):
            issues.append(f"No permission to access table: {table}")

    # Check 5: No SQL injection patterns
    injection_patterns = [";--", "' OR '1'='1", "UNION SELECT"]
    for pattern in injection_patterns:
        if pattern.lower() in sql.lower():
            issues.append("Potential SQL injection detected")

    return {
        "safe": len(issues) == 0,
        "issues": issues,
    }

def is_select_only(sql: str) -> bool:
    """Check if SQL is read-only."""
    sql_upper = sql.upper().strip()
    return sql_upper.startswith("SELECT") and "INTO" not in sql_upper

def extract_tables_from_sql(sql: str) -> List[str]:
    """Extract table names from SQL (simplified)."""
    import sqlparse

    parsed = sqlparse.parse(sql)[0]
    tables = []

    for token in parsed.tokens:
        if isinstance(token, sqlparse.sql.Identifier):
            tables.append(token.get_real_name())

    return tables

from functools import lru_cache
import hashlib

class SQLQueryCache:
    def __init__(self):
        self.template_cache = {}

    async def get_or_generate_sql(self, question: str, schema: str) -> str:
        """Cache SQL templates for repeated query patterns."""

        # Normalize question (remove specific values)
        normalized = self._normalize_question(question)
        cache_key = hashlib.md5(normalized.encode()).hexdigest()

        if cache_key in self.template_cache:
            print(f"Cache hit: {normalized}")
            return self._fill_template(self.template_cache[cache_key], question)

        # Cache miss, generate SQL
        sql = await text_to_sql(question, schema, self.conn)
        self.template_cache[cache_key] = self._extract_template(sql)

        return sql

    def _normalize_question(self, question: str) -> str:
        """Remove specific values, keep structure."""
        # "Show revenue for Q4 2025" → "Show revenue for Q# ####"
        import re
        normalized = re.sub(r'\d{4}', '####', question)  # Years
        normalized = re.sub(r'Q\d', 'Q#', normalized)     # Quarters
        return normalized

    def _extract_template(self, sql: str) -> str:
        """Extract parameterized template from SQL."""
        # "WHERE year = 2025" → "WHERE year = ?"
        import re
        template = re.sub(r"= \d+", "= ?", sql)
        template = re.sub(r"'[^']*'", "?", template)
        return template

    def _fill_template(self, template: str, question: str) -> str:
        """Fill template with values from question."""
        # Extract values from question and fill into template
        # Simplified implementation
        return template

async def safe_sql_execution(sql_template: str, params: dict, conn: asyncpg.Connection) -> List:
    """Execute SQL with parameterized queries."""

    # Convert template to parameterized query
    # "SELECT * FROM orders WHERE region = ?" → "SELECT * FROM orders WHERE region = $1"

    safe_sql = sql_template
    param_values = []

    for key, value in params.items():
        placeholder = f"${len(param_values) + 1}"
        safe_sql = safe_sql.replace(f":{key}", placeholder)
        param_values.append(value)

    # Execute safely
    return await conn.fetch(safe_sql, *param_values)

async def execute_with_rls(sql: str, tenant_id: str, conn: asyncpg.Connection) -> List:
    """Execute SQL with automatic tenant filtering."""

    # Set tenant context
    await conn.execute("SET app.current_tenant_id = $1", tenant_id)

    try:
        # RLS policy automatically filters results
        rows = await conn.fetch(sql)
        return rows

    finally:
        await conn.execute("RESET app.current_tenant_id")