Database QA Agent

Privacy-first local text-to-SQL system powered by defog/sqlcoder-7b-2 and ChromaDB vector retrieval, operating completely offline with zero third-party API dependencies.

---

Overview

Transmitting internal operational databases or sensitive financial schemas to external commercial LLM APIs introduces compliance, privacy, and cost concerns.

Database QA Agent is an offline text-to-SQL architecture designed for EV battery leasing operations. It pairs a locally loaded 7-billion parameter language model (defog/sqlcoder-7b-2) in FP16 with ChromaDB vector stores to retrieve relevant DDL schemas, domain-specific business rules, and few-shot query examples. The resulting SQL is sanitized through a regex-based security validator and executed locally inside an in-memory DuckDB engine.

---

Architecture

The system coordinates retrieval, prompt injection, local weights inference, and execution:

flowchart TD
    UserQuery([User Natural Language Query]) --> Embedder[ChromaDB Vector Retrieval]
    
    subgraph Offline Context Stores
        DDLStore[(DDL Schema Embeddings: Garage & Transaction DBs)] --> Embedder
        RuleStore[(Domain Business Rules & Column Docs)] --> Embedder
        FewShotStore[(77+ Few-Shot SQL Examples)] --> Embedder
    end

    Embedder --> Prompt[Injected Context Prompt]
    Prompt --> LocalLLM[Local LLM: defog/sqlcoder-7b-2 in FP16]
    LocalLLM --> RawSQL[Generated SQL Query]

    RawSQL --> SecurityValidator{has_dangerous_keywords_v2}
    SecurityValidator -->|Destructive: DROP, ALTER, UPDATE, INSERT| SecurityError([Execution Blocked])
    SecurityValidator -->|Read-Only: Safe SELECT| DuckDBExec[DuckDB In-Memory Execution]

    DuckDBExec --> ResultDataFrame[Pandas DataFrame & Tabular UI]

Loading

---

Key Features

• 100% Air-Gapped Local Inference: Runs locally via Hugging Face Transformers (defog/sqlcoder-7b-2). Zero network requests, zero token costs, and zero data leakage.
• RAG for Schema & Context Injection: Overrides Vanna AI's core vector storage to index DDL definitions, column-level documentation, and 77+ curated few-shot question-SQL pairs across operational and financial databases.
• Phase 1 to Phase 2A Evolution:
  • Phase 1 (Baseline): Reached ~50% accuracy. A naive substring validator blocked queries on valid timestamp columns (created_at, updated_at) because it detected "create" and "update" substrings.
  • Phase 2A (Refined): Implemented regex word boundaries (\b(DROP|CREATE|UPDATE)\b) with string literal stripping, added complex multi-table JOIN and date interval examples, boosting internal regression benchmark accuracy from 50% to 100%.
• Isolated In-Memory Querying: Executes against DuckDB instances containing local schemas (GARAGE_DB with 11 tables, TRANSACTION_DB with 14 tables).

---

Technical Stack

• Model: defog/sqlcoder-7b-2 (7B parameter quantized/FP16 weights)
• Vector Retrieval: ChromaDB
• Framework: Custom integration extending Vanna AI (LocalLLM_Vanna)
• Execution: DuckDB (in-memory relational database)
• Deep Learning Runtime: PyTorch, Hugging Face Transformers, Accelerate
• Language: Python 3.10+

---

Evaluation & Benchmark Results

The system was evaluated against a 10-query regression benchmark testing diverse SQL constructs (aggregations, multi-table JOINs, time-window logic, and business formulas):

#	Benchmark Query Description	Target SQL Construct	Result
1	Simple count (garage count)	COUNT(*) on single table	PASS
2	Date filter (garages created last month)	Date truncation & interval comparison	PASS
3	Year filter (contracts from this year)	Timestamp extraction & year filtering	PASS
4	2-table JOIN (contracts with garage names)	Inner JOIN across foreign keys	PASS
5	Business logic (overdue garages)	Status flags & operational condition filters	PASS
6	Aggregation (total revenue)	SUM() across transaction records	PASS
7	Date range (payments last month)	Bounded date window evaluation	PASS
8	Complex JOIN + GROUP BY (top 5 garages by revenue)	Multi-table JOIN, grouping, ordering, LIMIT	PASS
9	Business formula (collection rate)	Derived metric calculation with division	PASS
10	3-table JOIN (contracts with garage and org details)	Triple table relationship traversal	PASS

Result: 100% pass rate (10/10) on the internal regression benchmark suite.

---

Hardware Requirements

• GPU: NVIDIA GPU with at least 14 GB VRAM (tested on NVIDIA Tesla T4 / RTX 3090 / A100) for FP16 inference.
• System RAM: 16 GB minimum.
• Disk Space: ~15 GB for model weight caching (defog/sqlcoder-7b-2).

---

Project Structure

database-qa-agent/
├── README.md
├── LICENSE
├── .gitignore
├── requirements.txt
├── src/
│   └── chatbot/
│       ├── __init__.py
│       ├── llm.py                  # LocalLLM_Vanna subclass for offline inference
│       ├── vector_store.py         # VectorStoreManager for ChromaDB indices
│       ├── security.py             # has_dangerous_keywords_v2 regex validator
│       ├── database.py             # DatabaseManager for DuckDB connections
│       └── schemas/
│           ├── garage.sql          # 11-table PostgreSQL DDL
│           └── transaction.sql     # 14-table PostgreSQL DDL
├── data/
│   ├── training/
│   │   ├── few_shots.yaml          # 77+ curated question-SQL training pairs
│   │   ├── business_rules.yaml     # Domain business terms to SQL mappings
│   │   └── column_docs.yaml        # Schema column descriptions
│   └── mock/
│       └── generate_fixtures.py    # Synthetic test fixture generator
├── notebooks/
│   └── database_chatbot_experiments.ipynb
└── tests/
    ├── test_security.py            # Unit tests for SQL validator (5 valid + 6 malicious)
    └── test_sql_generation.py      # Regression test harness

---

Installation & Setup

# Clone the repository
git clone https://github.com/Ayush-Sharma99/database-qa-agent.git
cd database-qa-agent

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

---

Usage

Run Security Unit Tests

pytest tests/test_security.py

Run Regression Benchmark

pytest tests/test_sql_generation.py

---

Known Limitations

• VRAM Requirements: Running FP16 weights requires ~14GB VRAM. Implementing 4-bit NF4 quantization via bitsandbytes is planned to enable execution on consumer GPUs (8GB VRAM).
• In-Memory Indices: ChromaDB collections are currently seeded in-memory during setup; persisting vector stores across container lifecycles is recommended for deployment.
• Domain Scope: The 100% benchmark score applies specifically to the 10-query internal regression suite. Complex arbitrary PostgreSQL window functions or nested CTEs outside the training distribution may require additional few-shot tuning.

---

License

MIT License. See LICENSE for details.

Author

Developed by Ayush Sharma.