slm_text_to_sql/README.md

# SLM Text-to-SQL

`slm_text_to_sql` is a lightweight, local Text-to-SQL translation library powered entirely by a Small Language Model (SLM) running on CPU. It translates natural language questions into database-specific SQL queries. 

To achieve production-grade reliability on quantized 1.5B models, the library features **context-aware schema pruning** and an **agentic self-correction loop** that validates SQL queries against an in-memory database and corrects execution errors dynamically.

---

## Key Features

- **Local & Private**: Runs completely on CPU / RAM. Zero API keys, zero network latency, and complete data privacy.
- **Agentic Self-Correction**: Automatically creates an ephemeral in-memory SQLite database, runs the DDL schema, inserts dummy data, and tests the generated query. If execution fails (e.g. syntax, column mismatch), the agent reads the error and dynamically debugs the SQL query in a correction feedback loop.
- **Context-Aware Schema Pruning**: Prunes large schemas down to a maximum number of relevant tables (default: 8) based on keyword overlap and foreign key relationship paths to prevent model confusion.
- **Agnostic & Custom Few-Shots**: Overrides/augments few-shot examples with domain-specific templates. Abstract templates are used by default to prevent model copying bias.
- **Column Descriptions**: Supports passing a glossary dictionary of table/column semantics (e.g. enums, flags) to guide logical mapping.
- **Streaming Support**: Stream token-by-token output in real-time via a Python generator.

---

## Installation

Install locally for development:

```bash
# 1. Create a fresh virtual environment
python3 -m venv .venv
source .venv/bin/activate

# 2. Install the package in editable mode
pip install -e .
```

*Note: Requires `onnxruntime-genai`, `huggingface_hub`, and `pyyaml`.*

---

## Quick Start

```python
from slm_text_to_sql import SLMTextToSQL

# Initialize the agent
agent = SLMTextToSQL()

schema = """
CREATE TABLE users (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    role TEXT DEFAULT 'customer'
);
CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER REFERENCES users(id),
    total_amount REAL NOT NULL
);
"""

# Generates SQL query and validates/corrects it via SQLite in the background
query = agent.generate_sql(
    schema=schema,
    question="Find the total amount spent by user 'John Doe'."
)

print(query)
# Output:
# SELECT SUM(o.total_amount) FROM orders o JOIN users u ON o.customer_id = u.id WHERE u.name = 'John';
```

### Streaming Example (Bypasses Self-Correction)

```python
from slm_text_to_sql import SLMTextToSQL

agent = SLMTextToSQL()

# Stream tokens as they are generated
for token in agent.generate_sql(
    schema=schema,
    question="List all user names.",
    stream=True
):
    print(token, end="", flush=True)
print()
```

---

## Achieving Great Accuracy (Best Practices)

Small language models (1.5B parameters) are highly sensitive to prompt structure. Follow these design guidelines to achieve **95%+ execution accuracy**:

### 1. Enable Agentic Self-Correction
Keep `stream=False` (default) to run the self-correction validation loop. If the model hallucinations a column name or a bad join path, the agent captures the exact SQLite error, matches it against table definitions, and feeds it back to the model to correct itself.

### 2. Document Your Schema DDL (Inline Comments)
Quantized local models rely heavily on semantic hints inside the DDL. Adding inline comments directly to table schemas helps the model map synonyms:
```sql
CREATE TABLE users ( -- Represents users, clients, and customers
    id INTEGER PRIMARY KEY,
    status TEXT -- Either 'active' or 'inactive'
);
```

### 3. Prevent Few-Shot Name Copying Bias
If you pass custom few-shot examples via `few_shot_examples`, ensure the table and column names used inside the examples are **completely distinct** from your actual database tables. Small models tend to prioritize exact strings in examples over DDL schemas if they match.

### 4. Provide Column Glossaries
Use the `column_descriptions` parameter to map complex columns or enums to their definitions:
```python
column_descriptions = {
    "users.role": "Can be 'customer', 'manager', or 'admin'",
    "orders.status": "Can be 'pending', 'completed', or 'cancelled'"
}
```

---

## Self-Correction Loop Walkthrough

When a query fails validation, the agent automatically executes the following loop:

```mermaid
graph TD
    A[Initial Generation] --> B[Execute against SQLite DB]
    B -->|Success| C[Return SQL Query]
    B -->|Database Error| D[Extract Table Column Specs]
    D --> E[Construct History & Feedback Prompt]
    E --> F[Correction Attempt]
    F --> B
```

- **Example Database Error:** `no such column: o.product_id`
- **Feedback Injected:**
  ```text
  Failed Attempt #1 SQL:
  SELECT p.name FROM products p JOIN orders o ON p.id = o.product_id;
  
  Failed Attempt #1 Database Error:
  no such column: o.product_id
  
  Available column definitions for tables in your query:
  - Table 'products' columns: id, name, price, stock
  - Table 'orders' columns: id, customer_id, total_amount, status
  - Table 'order_items' columns: id, order_id, product_id, quantity
  ```
- **Result:** The model immediately notices `orders` lacks `product_id`, looks at `order_items` columns, and joins them correctly.

---

## Configuration API

### Constructor Configuration

```python
SLMTextToSQL(
    model_path=None,   # Explicit path to an ONNX model directory (optional)
    cache_dir=None,    # Cache directory for auto-downloads
    n_ctx=2048,        # Context window size (default: 2048)
    n_threads=4        # Number of CPU threads (default: 4)
)
```

### Query Generation Method

```python
agent.generate_sql(
    schema: str,                     # Database DDL schema (CREATE TABLE statements)
    question: str,                   # User natural language query
    temperature: float = 0.0,        # Sampling temperature (0.0 for deterministic answers)
    max_tokens: int = None,          # Maximum token limit for the response
    stream: bool = False,            # Returns generator yielding tokens if True
    column_descriptions: dict = None,# Optional dict mapping "table.col" -> "desc"
    few_shot_examples: list = None,  # Custom list of example dicts to append to system prompt
    system_prompt: str = None,       # Overrides the built-in system rules entirely
    max_iterations: int = 5,         # Maximum self-correction retry loops
    max_pruned_tables: int = 8       # Schema pruning table limit cap
)
```

---

## Environment Variables

All constructor parameters can be overridden via environment variables:

| Variable | Description | Default |
|---|---|---|
| `SLM_TEXT_TO_SQL_CONFIG` | Path to a custom `config.yaml` file | — |
| `SLM_TEXT_TO_SQL_CACHE_DIR` | Override model download/cache directory | — |
| `SLM_TEXT_TO_SQL_N_THREADS` | Number of CPU threads | `4` |
| `SLM_TEXT_TO_SQL_N_CTX` | Context window size | `2048` |
| `SLM_TEXT_TO_SQL_MAX_TOKENS` | Default max tokens per answer | `512` |

---

## Fine-Tuning with QLoRA

The library includes a complete script to fine-tune base models (e.g. `Qwen/Qwen2.5-Coder-1.5B-Instruct`) for Text-to-SQL using QLoRA.

### Running Fine-Tuning

```bash
python -m slm_text_to_sql.fine_tune
```

---

## License

Apache License 2.0.