Natural Language → SQL Analytics Pipeline

User Question
    │
    ▼
┌─────────────────────────────────────────┐
│   Embedding (sentence-transformers)    │
│              ↓                           │
│   "Show me top 5 trades by value"      │
│   → [0.12, 0.45, -0.23, ...]          │
└──────────────┬──────────────────────────┘
               │ 384-dim vector
               ▼
    ┌──────────────────────────┐
    │   FAISS Retriever        │
    │   4 Category Indexes:    │
    │   - SCHEMA               │
    │   - GLOSSARY             │
    │   - RELATIONSHIP         │
    │   - POLICY               │
    └──────────────┬───────────┘
                   │ context docs
                   ▼
    ┌──────────────────────────┐
    │  CRISP Prompt Template   │
    │  + Retrieved Context     │
    │  + User Question         │
    └──────────────┬───────────┘
                   │
                   ▼
    ┌──────────────────────────┐
    │   Ollama (Mistral 7B)    │
    │   SQL Generation         │
    └──────────────┬───────────┘
                   │ generated SQL
                   ▼
    ┌──────────────────────────┐      ┌─────────────┐
    │   SQL Validator          │─────▶│  Retry Loop?│
    │   12 Guardrails          │◀─────│  Feedback   │
    └──────────────┬───────────┘      └─────────────┘
                   │ valid SQL
                   ▼
    ┌──────────────────────────┐
    │   PostgreSQL Executor    │
    │   Run Query              │
    └──────────────┬───────────┘
                   │ results
                   ▼
    ┌──────────────────────────┐
    │   Response Formatter     │
    │   (Table/CSV/JSON/HTML)  │
    └──────────────────────────┘

# Every query retrieval uses this strategy:
SCHEMA:       Top 2 matches       (always included)
RELATIONSHIP: Top 1 match         (always included)
GLOSSARY:     Top 2 matches       (only if similarity > 0.40)
POLICY:       All documents       (always included)

KNOWLEDGE_BASE = [
    {
        "text": "Table: trades\nColumns:\n  - trade_id (INTEGER, PRIMARY KEY)\n  - client_id (INTEGER, FOREIGN KEY)\n  - symbol (TEXT)\n  - side (TEXT: BUY or SELL)\n  - quantity (INTEGER)\n  - price (DECIMAL)\n  - trade_date (DATE)",
        "category": "SCHEMA",
        "metadata": {"table": "trades"},
    },
    {
        "text": "Business term: Net Exposure\nDefinition: Sum of absolute notional value per counterparty\nSQL pattern: SUM(CASE WHEN side='BUY' THEN quantity*price ELSE -quantity*price END)",
        "category": "GLOSSARY",
        "metadata": {"term": "net_exposure"},
    },
    # ... more documents ...
]

<ROLE>
You are a precise SQL translator. Your ONLY job is to convert the user's
question into a single PostgreSQL-compatible SELECT query.

You are NOT a helpful assistant. Do NOT add anything the user didn't ask for.
Do NOT explain your reasoning. Output ONLY the SQL query.
</ROLE>

<SCHEMA>
Table: trades
Columns: trade_id, client_id, symbol, side, quantity, price, trade_date
Table: clients
Columns: client_id, client_name, risk_tier
</SCHEMA>

<RELATIONSHIPS>
trades JOIN clients ON trades.client_id = clients.id
</RELATIONSHIPS>

<GLOSSARY>
"Net Exposure" = SUM(CASE WHEN side='BUY' THEN quantity*price ...)
"Active Client" = clients with trades in last 30 days
</GLOSSARY>

<RULES>
1. Output ONLY a single SELECT statement. No explanation.
2. Use ONLY tables and columns listed in SCHEMA.
3. Do NOT add WHERE clauses unless the user explicitly asks.
4. Do NOT use subqueries unless necessary for the question.
5. Always include a LIMIT clause (max 1000 rows).
6. Use table aliases (SELECT t.trade_id FROM trades t).
</RULES>

<QUESTION>
{user_question}
</QUESTION>

# User asks: "What's the average trade value?"
# LLM generates:
SELECT AVG(t.trade_value) FROM trades t

# Column 'trade_value' doesn't exist. The correct columns are:
# - trade_id, client_id, symbol, side, quantity, price, trade_date

# Validator detects:
# - FROM clause: "trades t"
# - alias_map = {"t": "trades"}
# - SELECT clause: "t.trade_value"
# - Resolve: t.trade_value → trades.trade_value
# - Check schema: trades.trade_value EXISTS? NO
# - Block and return error

def retry_with_feedback(question, context, failed_sql, error_message):
    """
    Retry prompt: Ask LLM to fix ONLY the specific error.
    """
    retry_prompt = f"""
    Your previous SQL was rejected:
    {failed_sql}

    Error: {error_message}

    Fix ONLY this error. Use ONLY columns in the schema.
    Output ONLY the corrected SQL query.
    """

    # Send to LLM with temperature=0.1
    corrected_sql = mistral.generate(retry_prompt, temperature=0.1)
    return corrected_sql

# docker-compose.yml for production
version: '3.8'
services:
  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: nl2sql
      POSTGRES_PASSWORD: nl2sql_prod
      POSTGRES_DB: financial_db
    ports:
      - "5432:5432"
    volumes:
      - postgres_data:/var/lib/postgresql/data

  ollama:
    image: ollama/ollama:latest
    ports:
      - "11434:11434"
    volumes:
      - ollama_data:/root/.ollama

volumes:
  postgres_data:
  ollama_data:

class DatabaseConnection:
    """Works with SQLite or PostgreSQL — pipeline doesn't care which."""

    def __init__(self, db_type="postgresql", config=None):
        if db_type == "postgresql":
            self.conn = psycopg2.connect(
                host=config["host"],
                user=config["user"],
                password=config["password"],
                database=config["database"]
            )
        elif db_type == "sqlite":
            self.conn = sqlite3.connect(config["path"])

    def execute(self, sql):
        """Same interface for both. Pipeline doesn't change."""
        cursor = self.conn.cursor()
        cursor.execute(sql)
        return cursor.fetchall()

faiss-cpu==1.7.4               # Vector similarity search
sentence-transformers==2.2.2   # Embeddings (all-MiniLM)
psycopg2==2.9.9                # PostgreSQL driver
requests==2.31.0               # Ollama HTTP client
sqlparse==0.4.4                # SQL parsing
pytest==7.4.3                  # Testing