SQuery-Lens

What is SQuery-Lens?

SQuery-Lens is a lightweight ML-powered system for SQL query understanding and schema relevance ranking, designed specifically for RAG (Retrieval-Augmented Generation) pipelines. It provides two core models:

1. Query Analyzer - Classifies SQL queries by complexity, keywords, category, and table count
2. Schema Ranker - Ranks database tables/columns by relevance to a natural language query

Key Performance: Filters 500 tables in ~7ms using efficient bi-encoder architecture with pre-computed embeddings.

Why SQuery-Lens?

The Problem:

• Text-to-SQL systems struggle with large schemas (hundreds of tables)
• LLM context windows get overwhelmed with irrelevant table definitions
• Generic embedding models don’t understand SQL semantics

The Solution:

• Query Understanding - Know what type of SQL is needed before generation
• Schema Pre-filtering - Reduce 500 tables to top 10-20 relevant ones
• Optimized for RAG - Fast inference, structured outputs, Python API

Architecture

Query Analyzer

Classifies queries using multi-head DistilBERT:

Query → DistilBERT → Shared Layer (768→512)
                           ↓
         ┌────────┬────────┼────────┬────────┐
         ↓        ↓        ↓        ↓        ↓
    Complexity Keywords Category  Subcat  TableCount

Outputs:

• Complexity: Simple, Medium, Complex, Very Complex
• Category: SELECT, JOIN, AGGREGATION, SUBQUERY, etc.
• Keywords: Predicted SQL keywords (SELECT, WHERE, GROUP BY…)
• Table Count: Expected number of tables needed

Schema Ranker (Bi-Encoder)

Efficient bi-encoder for table relevance scoring:

Query  → Encoder → Query Embedding
                        ↓
             Cosine Similarity → Scores
                        ↑
Tables → Encoder → Table Embeddings (pre-computed)

Benefits:

• Pre-compute table embeddings once
• O(1) query encoding + O(n) cosine similarity
• Scales to thousands of tables efficiently

Recommended RAG Pipeline

            User Query
                ↓
┌─────────────────────────────────────┐
│          Query Analyzer             │
│  → complexity, keywords, category   │
│  → Route: simple→template,          │
│           complex→LLM               │
└─────────────────────────────────────┘
                ↓
┌─────────────────────────────────────┐
│     Schema Ranker (optional)        │
│  → Pre-filter 500 tables → 20       │
└─────────────────────────────────────┘
                ↓
┌─────────────────────────────────────┐
│         RAG Vector Search           │
│  → Get detailed schema for top      │
│    tables                           │
└─────────────────────────────────────┘
                ↓
┌─────────────────────────────────────┐
│              LLM                    │
│  → Generate SQL with focused        │
│    context                          │
└─────────────────────────────────────┘

Quick Start

Installation

# Clone the repository
git clone https://github.com/amannirala13/SQuery-Lens.git
cd SQuery-Lens

# Install dependencies
pip install -r requirements.txt

# Download models from Hugging Face
huggingface-cli download amannirala13/squery-lens-models --local-dir ./models

Python API

from squery_lens import QueryAnalyzer, SchemaRanker

# Query Analysis
analyzer = QueryAnalyzer("./models")
result = analyzer.analyze("Find all customers who ordered more than 5 times")
print(result.complexity)   # "Medium"
print(result.category)     # "AGGREGATION"
print(result.keywords)     # ["SELECT", "FROM", "WHERE", "GROUP BY", "HAVING"]

# Schema Ranking
ranker = SchemaRanker("./models")
tables = ["customers", "orders", "products", "reviews", "inventory"]
hints = ranker.rank("Find top selling products", tables)
print(hints.top_tables)    # ["products", "orders"]
print(hints.scores)        # [0.89, 0.76, 0.32, ...]

API Server

Run the FastAPI server for REST access:

python -m squery_lens.server --port 8000

Endpoints:

• POST /analyze - Query analysis
• POST /rank - Schema ranking
• GET /health - Health check

TypeScript Client

For frontend/Node.js integration:

import { SQueryLensClient } from "squery-lens";

const client = new SQueryLensClient("http://localhost:8000");

const analysis = await client.analyze("Get user order history");
const ranking = await client.rank("Top products by revenue", tableList);

Limitations

• Query Analyzer: Trained on synthetic data, may not generalize to all SQL dialects
• Schema Ranker: Works best with descriptive table/column names
• Cold Start: First inference is slower due to model loading (~2-3 seconds)

Hardware Requirements

• Minimum: 4GB RAM, CPU inference supported
• Recommended: 8GB RAM, GPU for batch processing
• Model Size: ~250MB for both models combined