style: format code with ruff

adapters/db/base.py

@@ -1,8 +1,10 @@
 from typing import Tuple, List, Dict, Any, Protocol
 from typing import List, Tuple, Any
 
+
 class DBAdapter(Protocol):
     """Abstract database adapter for read-only queries."""
+
     name: str
     dialect: str
 
@@ -10,4 +12,4 @@ class DBAdapter(Protocol):
         """Generate a readable summary of the database schema with optional sample rows per table."""
 
     def execute(self, sql: str) -> Tuple[List[Tuple[Any, ...]], List[str]]:
-        """Execute a SELECT query and return (rows, columns)."""
+        """Execute a SELECT query and return (rows, columns)."""

adapters/db/postgres_adapter.py

@@ -2,6 +2,7 @@ import psycopg
 from typing import Any, List, Tuple
 from adapters.db.base import DBAdapter
 
+
 class PostgresAdapter(DBAdapter):
     name = "postgres"
     dialect = "postgres"
@@ -24,11 +25,14 @@ class PostgresAdapter(DBAdapter):
             tables = [t[0] for t in cur.fetchall()]
             lines = []
             for t in tables:
-                cur.execute(f"""
+                cur.execute(
+                    f"""
                     SELECT column_name, data_type
                     FROM information_schema.columns
                     WHERE table_name = %s;
-                """, (t,))
+                """,
+                    (t,),
+                )
                 cols = [f"{c[0]}:{c[1]}" for c in cur.fetchall()]
                 lines.append(f"- {t} ({', '.join(cols)})")
             return "\n".join(lines)

adapters/db/sqlite_adapter.py

@@ -2,6 +2,7 @@ import sqlite3
 from typing import List, Tuple, Any
 from adapters.db.base import DBAdapter
 
+
 class SQLiteAdapter(DBAdapter):
     name = "sqlite"
     dialect = "sqlite"

adapters/llm/base.py

@@ -2,15 +2,26 @@
 from __future__ import annotations
 from typing import Tuple, List, Dict, Any, Protocol
 
+
 class LLMProvider(Protocol):
     provider_id: str
 
-    def plan(self, *, user_query: str, schema_preview: str) -> Tuple[str, int, int, float]:
+    def plan(
+        self, *, user_query: str, schema_preview: str
+    ) -> Tuple[str, int, int, float]:
         """Return (plan_text, token_in, token_out, cost_usd)."""
 
-    def generate_sql(self, *, user_query: str, schema_preview: str, plan_text: str,
-                     clarify_answers: Dict[str, Any] | None = None) -> Tuple[str, str, int, int, float]:
+    def generate_sql(
+        self,
+        *,
+        user_query: str,
+        schema_preview: str,
+        plan_text: str,
+        clarify_answers: Dict[str, Any] | None = None,
+    ) -> Tuple[str, str, int, int, float]:
         """Return (sql, rationale, token_in, token_out, cost_usd)."""
 
-    def repair(self, *, sql: str, error_msg: str, schema_preview: str) -> Tuple[str, int, int, float]:
+    def repair(
+        self, *, sql: str, error_msg: str, schema_preview: str
+    ) -> Tuple[str, int, int, float]:
         """Return (patched_sql, token_in, token_out, cost_usd)."""

adapters/llm/openai_provider.py

@@ -11,14 +11,13 @@ from openai import OpenAI
 #  - OPENAI_MODEL_ID  (e.g., "gpt-4o-mini")
 
 
-
 class OpenAIProvider(LLMProvider):
     provider_id = "openai"
 
     def __init__(self) -> None:
         self.client = OpenAI(
             api_key=os.environ["OPENAI_API_KEY"],
-            base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+            base_url=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
         )
         self.model = os.getenv("OPENAI_MODEL_ID", "gpt-4o-mini")
 
@@ -27,16 +26,25 @@ class OpenAIProvider(LLMProvider):
             model=self.model,
             messages=[
                 {"role": "system", "content": "You create SQL query plans."},
-                {"role": "user", "content": f"Query: {user_query}\nSchema:\n{schema_preview}"}
+                {
+                    "role": "user",
+                    "content": f"Query: {user_query}\nSchema:\n{schema_preview}",
+                },
             ],
-            temperature=0
+            temperature=0,
         )
         msg = completion.choices[0].message.content
         usage = completion.usage
-        return msg, usage.prompt_tokens, usage.completion_tokens, self._estimate_cost(usage)
-
+        return (
+            msg,
+            usage.prompt_tokens,
+            usage.completion_tokens,
+            self._estimate_cost(usage),
+        )
 
-    def generate_sql(self, *, user_query, schema_preview, plan_text, clarify_answers=None):
+    def generate_sql(
+        self, *, user_query, schema_preview, plan_text, clarify_answers=None
+    ):
         prompt = f"""
         You are a precise SQL generator. 
         Return ONLY valid JSON with two keys: "sql" and "rationale".
@@ -60,9 +68,9 @@ class OpenAIProvider(LLMProvider):
             model=self.model,
             messages=[
                 {"role": "system", "content": "You convert natural language to SQL."},
-                {"role": "user", "content": prompt}
+                {"role": "user", "content": prompt},
             ],
-            temperature=0
+            temperature=0,
         )
         content = completion.choices[0].message.content.strip()
         usage = completion.usage  # ← لازم داریم
@@ -78,7 +86,7 @@ class OpenAIProvider(LLMProvider):
             end = content.rfind("}")
             if start != -1 and end != -1:
                 try:
-                    parsed = json.loads(content[start:end + 1])
+                    parsed = json.loads(content[start : end + 1])
                 except Exception:
                     raise ValueError(f"Invalid LLM JSON output: {content[:200]}")
             else:
@@ -93,19 +101,29 @@ class OpenAIProvider(LLMProvider):
         # IMPORTANT: return the expected 5-tuple
         return sql, rationale, t_in, t_out, cost
 
-
     def repair(self, *, sql, error_msg, schema_preview):
         completion = self.client.chat.completions.create(
             model=self.model,
             messages=[
-                {"role": "system", "content": "You fix SQL queries keeping them SELECT-only."},
-                {"role": "user", "content": f"SQL:\n{sql}\nError:\n{error_msg}\nSchema:\n{schema_preview}"}
+                {
+                    "role": "system",
+                    "content": "You fix SQL queries keeping them SELECT-only.",
+                },
+                {
+                    "role": "user",
+                    "content": f"SQL:\n{sql}\nError:\n{error_msg}\nSchema:\n{schema_preview}",
+                },
             ],
-            temperature=0
+            temperature=0,
         )
         msg = completion.choices[0].message.content
         usage = completion.usage
-        return msg, usage.prompt_tokens, usage.completion_tokens, self._estimate_cost(usage)
+        return (
+            msg,
+            usage.prompt_tokens,
+            usage.completion_tokens,
+            self._estimate_cost(usage),
+        )
 
     def _estimate_cost(self, usage):
         # Rough estimation example — can be refined with official token pricing

app/main.py

@@ -1,29 +1,29 @@
 from dotenv import load_dotenv
+
 load_dotenv()
 
 from fastapi import FastAPI
 from app.routers import nl2sql
+
 app = FastAPI(
     title="NL2SQL Copilot Prototype",
     version="0.1.0",
-    description="Natural Language -> SQL Copilot API"
+    description="Natural Language -> SQL Copilot API",
 )
 
 app.include_router(nl2sql.router, prefix="/api/v1")
 
+
 @app.get("/healthz")
 def health_check():
     return {"status": "ok"}
 
+
 @app.get("/")
 def root():
     return {"status": "ok", "message": "NL2SQL Copilot API is running"}
 
+
 @app.get("/health")
 def health():
-    return {
-        "status": "ok",
-        "db": "connected",
-        "llm": "reachable",
-        "uptime_sec": 123.4
-    }
+    return {"status": "ok", "db": "connected", "llm": "reachable", "uptime_sec": 123.4}

app/routers/nl2sql.py

@@ -19,7 +19,6 @@ import os
 router = APIRouter(prefix="/nl2sql")
 
 
-
 if os.getenv("DB_MODE", "sqlite") == "postgres":
     _db = PostgresAdapter(os.environ["POSTGRES_DSN"])
 else:
@@ -40,7 +39,7 @@ _pipeline = Pipeline(
     safety=Safety(),
     executor=_executor,
     verifier=_verifier,
-    repair=_repair
+    repair=_repair,
 )
 
 
@@ -48,6 +47,7 @@ def _to_dict(obj):
     """Helper: safely convert dataclass → dict."""
     return asdict(obj) if is_dataclass(obj) else obj
 
+
 def _round_trace(t: dict) -> dict:
     if t.get("cost_usd") is not None:
         t["cost_usd"] = round(t["cost_usd"], 6)
@@ -55,9 +55,12 @@ def _round_trace(t: dict) -> dict:
         t["duration_ms"] = round(t["duration_ms"], 2)
     return t
 
+
 @router.post("", name="nl2sql_handler")
 def nl2sql_handler(request: NL2SQLRequest):
-    result = _pipeline.run(user_query=request.query, schema_preview=request.schema_preview)
+    result = _pipeline.run(
+        user_query=request.query, schema_preview=request.schema_preview
+    )
 
     # --- Ensure result type ---
     if not isinstance(result, StageResult):

app/schemas.py

@@ -1,11 +1,13 @@
 from pydantic import BaseModel
 from typing import List, Optional, Any, Dict
 
+
 class NL2SQLRequest(BaseModel):
     query: str
     schema_preview: str
     db_name: Optional[str] = "default"
 
+
 class TraceModel(BaseModel):
     stage: str
     duration_ms: float
@@ -14,16 +16,19 @@ class TraceModel(BaseModel):
     cost_usd: float | None = 0
     notes: Dict[str, Any] | None = None
 
+
 class NL2SQLResponse(BaseModel):
     ambiguous: bool = False
     sql: str
     rationale: Optional[str] = None
     traces: List[TraceModel] = []
 
+
 class ClarifyResponse(BaseModel):
     ambiguous: bool = True
     questions: List[str]
 
+
 class ErrorResponse(BaseModel):
     error: str
     details: List[str] | None = None

benchmarks/evaluate_spider.py

@@ -13,16 +13,19 @@ from sqlglot.errors import ParseError
 LOG_DIR = Path("logs/spider_eval")
 LOG_DIR.mkdir(parents=True, exist_ok=True)
 
+
 def normalize_sql(sql: str) -> str:
     # نسخه ساده؛ می‌تونی قوی‌ترش کنی با پارس + بازسازی
     return " ".join(sql.lower().strip().split())
 
+
 def compare_results(pred_rows, gold_rows):
     if pred_rows is None or gold_rows is None:
         return False
     # اگر ترتیب مهم نیست
     return set(pred_rows) == set(gold_rows)
 
+
 def try_execute_sql(sql_db, sql, timeout: float = None):
     start = time.time()
     try:
@@ -31,6 +34,7 @@ def try_execute_sql(sql_db, sql, timeout: float = None):
     except Exception as e:
         return None, time.time() - start, str(e)
 
+
 def exact_match_structural(sql_pred: str, sql_gold: str) -> bool:
     try:
         ast_pred = parse_one(sql_pred)
@@ -54,13 +58,19 @@ def exact_match_structural(sql_pred: str, sql_gold: str) -> bool:
     norm_gold = normalize_ast(ast_gold)
     return norm_prd == norm_gold
 
+
 def get_git_commit_hash() -> str:
     try:
-        out = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("ascii")
+        out = (
+            subprocess.check_output(["git", "rev-parse", "HEAD"])
+            .strip()
+            .decode("ascii")
+        )
         return out
     except Exception:
         return "UNKNOWN"
 
+
 FORBIDDEN_NODES = (
     exp.Insert,
     exp.Delete,
@@ -72,6 +82,7 @@ FORBIDDEN_NODES = (
     exp.Create,
 )
 
+
 def is_safe_sql(sql: str, dialect: str | None = None) -> bool:
     try:
         ast = parse_one(sql, read=dialect)
@@ -84,6 +95,7 @@ def is_safe_sql(sql: str, dialect: str | None = None) -> bool:
             return False
     return True
 
+
 def run_eval(split="dev", limit=100, resume=True, sleep_time: float = 0.01):
     data = load_spider_sqlite(split)
     if len(data) < limit:
@@ -94,8 +106,8 @@ def run_eval(split="dev", limit=100, resume=True, sleep_time: float = 0.01):
     commit_hash = get_git_commit_hash()
     start_ts = int(time.time())
 
-    pred_txt   = LOG_DIR / f"{split}_pred_{start_ts}.txt"
-    gold_txt   = LOG_DIR / f"{split}_gold_{start_ts}.txt"
+    pred_txt = LOG_DIR / f"{split}_pred_{start_ts}.txt"
+    gold_txt = LOG_DIR / f"{split}_gold_{start_ts}.txt"
     results_fn = LOG_DIR / f"{split}_results_{start_ts}.jsonl"
     metrics_fn = LOG_DIR / f"{split}_metrics_{start_ts}.json"
 
@@ -112,10 +124,11 @@ def run_eval(split="dev", limit=100, resume=True, sleep_time: float = 0.01):
                     pass
 
     write_header = not results_fn.exists()
-    with results_fn.open("a", encoding="utf-8") as fout, \
-         pred_txt.open("a", encoding="utf-8") as fpred, \
-         gold_txt.open("a", encoding="utf-8") as fgold:
-
+    with (
+        results_fn.open("a", encoding="utf-8") as fout,
+        pred_txt.open("a", encoding="utf-8") as fpred,
+        gold_txt.open("a", encoding="utf-8") as fgold,
+    ):
         if write_header:
             header = {
                 "commit_hash": commit_hash,
@@ -228,21 +241,28 @@ def run_eval(split="dev", limit=100, resume=True, sleep_time: float = 0.01):
             if sleep_time > 0:
                 time.sleep(sleep_time)
 
-
-    valid = [r for r in agg if (not r.get("safe_check_failed", False)) and r.get("gold_error") is None]
+    valid = [
+        r
+        for r in agg
+        if (not r.get("safe_check_failed", False)) and r.get("gold_error") is None
+    ]
     total_valid = len(valid)
     total_all = len(agg)
     if total_valid == 0:
         print("No valid examples to compute metrics")
         return
 
-    em_count        = sum(1 for r in valid if r["exact_match"])
+    em_count = sum(1 for r in valid if r["exact_match"])
     em_struct_count = sum(1 for r in valid if r["exact_match_structural"])
-    exec_acc_count  = sum(1 for r in valid if r["execution_accuracy"])
-    error_count     = sum(1 for r in agg if (r.get("error") is not None) and (not r.get("safe_check_failed", False)))
+    exec_acc_count = sum(1 for r in valid if r["execution_accuracy"])
+    error_count = sum(
+        1
+        for r in agg
+        if (r.get("error") is not None) and (not r.get("safe_check_failed", False))
+    )
     safe_fail_count = sum(1 for r in agg if r.get("safe_check_failed", False))
-    avg_gen_time    = sum(r["gen_time"] for r in valid) / total_valid
-    avg_exec_time   = sum(r["exec_time"] for r in valid) / total_valid
+    avg_gen_time = sum(r["gen_time"] for r in valid) / total_valid
+    avg_exec_time = sum(r["exec_time"] for r in valid) / total_valid
 
     metrics = {
         "commit_hash": commit_hash,

benchmarks/run.py

@@ -20,6 +20,7 @@ from nl2sql.repair import Repair
 from adapters.db.sqlite_adapter import SQLiteAdapter
 from adapters.llm.openai_provider import OpenAIProvider
 
+
 # ---- fallbacks: Dummy LLM (so it runs without API keys)
 class DummyLLM:
     provider_id = "dummy-llm"
@@ -28,7 +29,14 @@ class DummyLLM:
         text = f"- understand question: {user_query}\n- identify tables\n- join if needed\n- filter\n- order/limit"
         return text, 0, 0, 0.0
 
-    def generate_sql(self, *, user_query: str, schema_preview: str, plan_text: str, clarify_answers=None):
+    def generate_sql(
+        self,
+        *,
+        user_query: str,
+        schema_preview: str,
+        plan_text: str,
+        clarify_answers=None,
+    ):
         # naive demo SQL (so pipeline flows end-to-end)
         sql = "SELECT 1 AS one;"
         rationale = "Demo SQL from DummyLLM"
@@ -43,12 +51,15 @@ def ensure_demo_db(path: Path) -> None:
     if path.exists():
         return
     import sqlite3
+
     path.parent.mkdir(parents=True, exist_ok=True)
     con = sqlite3.connect(path)
     cur = con.cursor()
     cur.execute("CREATE TABLE users(id INTEGER PRIMARY KEY, name TEXT, spend REAL);")
-    cur.executemany("INSERT INTO users(id,name,spend) VALUES(?,?,?)",
-                    [(1,"Alice",120.5),(2,"Bob",80.0),(3,"Carol",155.0)])
+    cur.executemany(
+        "INSERT INTO users(id,name,spend) VALUES(?,?,?)",
+        [(1, "Alice", 120.5), (2, "Bob", 80.0), (3, "Carol", 155.0)],
+    )
     con.commit()
     con.close()
 
@@ -86,7 +97,7 @@ def run_benchmark(queries, schema_preview, pipeline: Pipeline, outfile: Path):
     for q in queries:
         t0 = time.perf_counter()
         r = pipeline.run(user_query=q, schema_preview=schema_preview)
-        latency_ms = (time.perf_counter()-t0)*1000
+        latency_ms = (time.perf_counter() - t0) * 1000
         ok = (not r.get("ambiguous")) and ("error" not in r)
 
         traces = r.get("traces", [])
@@ -97,15 +108,19 @@ def run_benchmark(queries, schema_preview, pipeline: Pipeline, outfile: Path):
             except Exception:
                 pass
 
-        results.append({
-            "query": q,
-            "exec_acc": 1.0 if ok else 0.0,
-            "safe_fail": 0.0 if ok else 1.0 if "unsafe" in str(r).lower() else 0.0,
-            "latency_ms": latency_ms,
-            "cost_usd": cost_sum,
-            "repair_attempts": sum(1 for t in traces if t.get("stage") == "repair"),
-            "provider": pipeline.generator.llm.provider_id if hasattr(pipeline.generator, "llm") else "unknown",
-        })
+        results.append(
+            {
+                "query": q,
+                "exec_acc": 1.0 if ok else 0.0,
+                "safe_fail": 0.0 if ok else 1.0 if "unsafe" in str(r).lower() else 0.0,
+                "latency_ms": latency_ms,
+                "cost_usd": cost_sum,
+                "repair_attempts": sum(1 for t in traces if t.get("stage") == "repair"),
+                "provider": pipeline.generator.llm.provider_id
+                if hasattr(pipeline.generator, "llm")
+                else "unknown",
+            }
+        )
 
     outfile.parent.mkdir(parents=True, exist_ok=True)
     with open(outfile, "w") as f:
@@ -118,10 +133,14 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--outfile", default="benchmarks/results/demo.jsonl")
     parser.add_argument("--db", default="data/bench_demo.db")
-    parser.add_argument("--use-openai", action="store_true", help="Use OpenAI provider if API key present")
+    parser.add_argument(
+        "--use-openai",
+        action="store_true",
+        help="Use OpenAI provider if API key present",
+    )
     args = parser.parse_args()
 
-    ROOT = Path(__file__).resolve().parents[1]   # project root
+    ROOT = Path(__file__).resolve().parents[1]  # project root
     outfile = (ROOT / args.outfile).resolve()
     db_path = (ROOT / args.db).resolve()
 

benchmarks/spider_loader.py

@@ -4,9 +4,8 @@ from dataclasses import dataclass
 from typing import List, Optional
 import os
 
-SPIDER_ROOT = pathlib.Path(
-    os.getenv("SPIDER_ROOT", "data/spider")
-)
+SPIDER_ROOT = pathlib.Path(os.getenv("SPIDER_ROOT", "data/spider"))
+
 
 @dataclass
 class SpiderItem:
@@ -15,7 +14,10 @@ class SpiderItem:
     gold_sql: str
     db_path: pathlib.Path
 
-def load_spider_sqlite(split: str = "dev", limit: Optional[int] = None) -> List[SpiderItem]:
+
+def load_spider_sqlite(
+    split: str = "dev", limit: Optional[int] = None
+) -> List[SpiderItem]:
     fn = {"dev": "dev.json", "train": "train_spider.json"}[split]
     json_path = SPIDER_ROOT / fn
     try:
@@ -23,7 +25,6 @@ def load_spider_sqlite(split: str = "dev", limit: Optional[int] = None) -> List[
     except Exception as e:
         raise RuntimeError(f"Failed to read Spider split file: {json_path} ({e})")
 
-
     out: list[SpiderItem] = []
     for ex in items[: (limit or len(items))]:
         db_id = ex["db_id"]
@@ -35,14 +36,16 @@ def load_spider_sqlite(split: str = "dev", limit: Optional[int] = None) -> List[
                 db_id=db_id,
                 question=ex["question"],
                 gold_sql=ex["query"],
-                db_path=db_path
+                db_path=db_path,
             )
         )
     return out
 
 
-def open_readonly_connection(db_path: pathlib.Path, timeout: float=5.0) -> sqlite3.Connection:
+def open_readonly_connection(
+    db_path: pathlib.Path, timeout: float = 5.0
+) -> sqlite3.Connection:
     uri = f"file:{db_path}?mode=ro&uri=true"
     conn = sqlite3.connect(uri, uri=True, timeout=timeout)
     conn.row_factory = sqlite3.Row
-    return conn
+    return conn

config.py

@@ -4,12 +4,15 @@ from dotenv import load_dotenv
 load_dotenv()
 
 
-def get_env_var(name: str, required: bool = True, default: str | None = None) -> str | None:
+def get_env_var(
+    name: str, required: bool = True, default: str | None = None
+) -> str | None:
     val = os.getenv(name, default)
     if required and not val:
         raise ValueError(f"Missing required environment variable: {name}")
     return val
 
+
 proxy_key = os.getenv("PROXY_API_KEY")
 proxy_base = os.getenv("PROXY_BASE_URL")
 openai_key = os.getenv("OPENAI_API_KEY")
@@ -17,7 +20,9 @@ openai_base = os.getenv("OPENAI_BASE_URL")
 
 api_key = proxy_key or openai_key
 if not api_key:
-    raise ValueError("Missing API key: set PROXY_API_KEY or OPENAI_API_KEY in environment/secrets.")
+    raise ValueError(
+        "Missing API key: set PROXY_API_KEY or OPENAI_API_KEY in environment/secrets."
+    )
 
 base_url = proxy_base or openai_base or "https://api.openai.com/v1"
 
@@ -33,11 +38,24 @@ LLM_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")  # or gpt-4o, gpt-4o-mini
 LLM_TEMPERATURE = float(os.getenv("LLM_TEMPERATURE", "0"))
 
 FORBIDDEN_KEYWORDS = {
-    "ATTACH", "PRAGMA",
-    "CREATE", "DROP", "ALTER", "VACUUM", "REINDEX", "TRIGGER",
-    "INSERT", "UPDATE", "DELETE", "REPLACE",
-    "GRANT", "REVOKE",
-    "BEGIN", "END", "COMMIT", "ROLLBACK",
+    "ATTACH",
+    "PRAGMA",
+    "CREATE",
+    "DROP",
+    "ALTER",
+    "VACUUM",
+    "REINDEX",
+    "TRIGGER",
+    "INSERT",
+    "UPDATE",
+    "DELETE",
+    "REPLACE",
+    "GRANT",
+    "REVOKE",
+    "BEGIN",
+    "END",
+    "COMMIT",
+    "ROLLBACK",
     "DETACH",
 }
 FORBIDDEN_TABLES = {"sqlite_master", "sqlite_temp_master"}

nl2sql/ambiguity_detector.py

@@ -1,16 +1,17 @@
 import re
 from typing import List
 
+
 class AmbiguityDetector:
     """Lightweight AmbiSQL-style ambiguity detection."""
 
     AMBIGUOUS_TERMS = ["recent", "top", "name", "rank", "latest"]
 
-    def detect(self, query:str, schema_preview: str) -> list[str]:
+    def detect(self, query: str, schema_preview: str) -> list[str]:
         hits = []
         q_lower = query.lower()
         for term in self.AMBIGUOUS_TERMS:
             if re.search(rf"\b{term}\b", q_lower):
                 hits.append(f"The term '{term}' is ambiguous in this query.'")
 
-        return hits
+        return hits

nl2sql/executor.py

@@ -2,6 +2,7 @@ import time
 from nl2sql.types import StageResult, StageTrace
 from adapters.db.base import DBAdapter
 
+
 class Executor:
     name = "executor"
 
@@ -12,10 +13,18 @@ class Executor:
         t0 = time.perf_counter()
         try:
             rows, cols = self.db.execute(sql)
-            trace = StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000,
-                               notes={"row_count": len(rows), "col_count": len(cols)})
-            return StageResult(ok=True, data={"rows": rows, "columns": cols}, trace=trace)
+            trace = StageTrace(
+                stage=self.name,
+                duration_ms=(time.perf_counter() - t0) * 1000,
+                notes={"row_count": len(rows), "col_count": len(cols)},
+            )
+            return StageResult(
+                ok=True, data={"rows": rows, "columns": cols}, trace=trace
+            )
         except Exception as e:
-            trace = StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000,
-                               notes={"error": str(e)})
+            trace = StageTrace(
+                stage=self.name,
+                duration_ms=(time.perf_counter() - t0) * 1000,
+                notes={"error": str(e)},
+            )
             return StageResult(ok=False, data=None, trace=trace, error=[str(e)])

nl2sql/generator.py

@@ -4,34 +4,48 @@ from typing import Optional, Dict, Any
 from nl2sql.types import StageResult, StageTrace
 from adapters.llm.base import LLMProvider
 
+
 class Generator:
     name = "generator"
 
     def __init__(self, llm: LLMProvider) -> None:
         self.llm = llm
 
-    def run(self, *, user_query: str, schema_preview: str, plan_text: str,
-            clarify_answers: Optional[Dict[str, Any]] = None) -> StageResult:
+    def run(
+        self,
+        *,
+        user_query: str,
+        schema_preview: str,
+        plan_text: str,
+        clarify_answers: Optional[Dict[str, Any]] = None,
+    ) -> StageResult:
         t0 = time.perf_counter()
         try:
             res = self.llm.generate_sql(
                 user_query=user_query,
                 schema_preview=schema_preview,
                 plan_text=plan_text,
-                clarify_answers=clarify_answers or {}
+                clarify_answers=clarify_answers or {},
             )
         except Exception as e:
             return StageResult(ok=False, error=[f"Generator failed: {e}"])
 
         # Expect a 5-tuple
         if not isinstance(res, tuple) or len(res) != 5:
-            return StageResult(ok=False, error=["Generator contract violation: expected 5-tuple (sql, rationale, t_in, t_out, cost)"])
+            return StageResult(
+                ok=False,
+                error=[
+                    "Generator contract violation: expected 5-tuple (sql, rationale, t_in, t_out, cost)"
+                ],
+            )
 
         sql, rationale, t_in, t_out, cost = res
 
         # Type/shape checks
         if not isinstance(sql, str) or not sql.strip():
-            return StageResult(ok=False, error=["Generator produced empty or non-string SQL"])
+            return StageResult(
+                ok=False, error=["Generator produced empty or non-string SQL"]
+            )
         if not sql.lower().lstrip().startswith("select"):
             return StageResult(ok=False, error=[f"Generated non-SELECT SQL: {sql}"])
 
@@ -45,5 +59,6 @@ class Generator:
             notes={"rationale_len": len(rationale)},
         )
 
-        return StageResult(ok=True, data={"sql": sql, "rationale": rationale}, trace=trace)
-
+        return StageResult(
+            ok=True, data={"sql": sql, "rationale": rationale}, trace=trace
+        )

nl2sql/pipeline.py

@@ -17,14 +17,17 @@ class Pipeline:
     All stages return structured traces and errors but final result is JSON-safe dict.
     """
 
-    def __init__(self, *,
-                 detector: AmbiguityDetector,
-                 planner: Planner,
-                 generator: Generator,
-                 safety: Safety,
-                 executor: Executor,
-                 verifier: Verifier,
-                 repair: Repair):
+    def __init__(
+        self,
+        *,
+        detector: AmbiguityDetector,
+        planner: Planner,
+        generator: Generator,
+        safety: Safety,
+        executor: Executor,
+        verifier: Verifier,
+        repair: Repair,
+    ):
         self.detector = detector
         self.planner = planner
         self.generator = generator
@@ -59,8 +62,13 @@ class Pipeline:
             return StageResult(ok=False, data=None, trace=None, errors=[f"{e}", tb])
 
     # ------------------------------------------------------------
-    def run(self, *, user_query: str, schema_preview: str,
-            clarify_answers: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    def run(
+        self,
+        *,
+        user_query: str,
+        schema_preview: str,
+        clarify_answers: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
         """
         Always returns:
         {
@@ -86,26 +94,45 @@ class Pipeline:
                     "error": False,
                     "details": [f"Ambiguities found: {len(questions)}"],
                     "questions": questions,
-                    "traces": []
+                    "traces": [],
                 }
         except Exception as e:
-            return {"ambiguous": True, "error": True, "details": [f"Detector failed: {e}"], "traces": []}
+            return {
+                "ambiguous": True,
+                "error": True,
+                "details": [f"Detector failed: {e}"],
+                "traces": [],
+            }
 
         # --- 2) planner
-        r_plan = self._safe_stage(self.planner.run, user_query=user_query, schema_preview=schema_preview)
+        r_plan = self._safe_stage(
+            self.planner.run, user_query=user_query, schema_preview=schema_preview
+        )
         traces.extend(self._trace_list(r_plan))
         if not r_plan.ok:
-            return {"ambiguous": False, "error": True, "details": r_plan.errors, "traces": traces}
+            return {
+                "ambiguous": False,
+                "error": True,
+                "details": r_plan.errors,
+                "traces": traces,
+            }
 
         # --- 3) generator
-        r_gen = self._safe_stage(self.generator.run,
-                                 user_query=user_query,
-                                 schema_preview=schema_preview,
-                                 plan_text=r_plan.data.get("plan"),
-                                 clarify_answers=clarify_answers or {})
+        r_gen = self._safe_stage(
+            self.generator.run,
+            user_query=user_query,
+            schema_preview=schema_preview,
+            plan_text=r_plan.data.get("plan"),
+            clarify_answers=clarify_answers or {},
+        )
         traces.extend(self._trace_list(r_gen))
         if not r_gen.ok:
-            return {"ambiguous": False, "error": True, "details": r_gen.errors, "traces": traces}
+            return {
+                "ambiguous": False,
+                "error": True,
+                "details": r_gen.errors,
+                "traces": traces,
+            }
         sql = r_gen.data.get("sql")
         rationale = r_gen.data.get("rationale")
 
@@ -113,7 +140,12 @@ class Pipeline:
         r_safe = self._safe_stage(self.safety.check, sql=sql)
         traces.extend(self._trace_list(r_safe))
         if not r_safe.ok:
-            return {"ambiguous": False, "error": True, "details": r_safe.errors, "traces": traces}
+            return {
+                "ambiguous": False,
+                "error": True,
+                "details": r_safe.errors,
+                "traces": traces,
+            }
 
         # --- 5) executor
         r_exec = self._safe_stage(self.executor.run, sql=r_safe.data["sql"])
@@ -129,10 +161,12 @@ class Pipeline:
         # --- 7) repair loop if verification failed
         if not verified:
             for attempt in range(2):
-                r_fix = self._safe_stage(self.repair.run,
-                                         sql=sql,
-                                         error_msg="; ".join(details or ["unknown"]),
-                                         schema_preview=schema_preview)
+                r_fix = self._safe_stage(
+                    self.repair.run,
+                    sql=sql,
+                    error_msg="; ".join(details or ["unknown"]),
+                    schema_preview=schema_preview,
+                )
                 traces.extend(self._trace_list(r_fix))
                 if not r_fix.ok:
                     break

nl2sql/planner.py

@@ -3,14 +3,24 @@ import time
 from nl2sql.types import StageResult, StageTrace
 from adapters.llm.base import LLMProvider
 
+
 class Planner:
     name = "planner"
+
     def __init__(self, llm: LLMProvider) -> None:
         self.llm = llm
 
     def run(self, *, user_query: str, schema_preview: str) -> StageResult:
         t0 = time.perf_counter()
-        plan_text, t_in, t_out, cost = self.llm.plan(user_query=user_query, schema_preview=schema_preview)
-        trace = StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000,
-                           token_in=t_in, token_out=t_out, cost_usd=cost, notes={"len_plan": len(plan_text)})
+        plan_text, t_in, t_out, cost = self.llm.plan(
+            user_query=user_query, schema_preview=schema_preview
+        )
+        trace = StageTrace(
+            stage=self.name,
+            duration_ms=(time.perf_counter() - t0) * 1000,
+            token_in=t_in,
+            token_out=t_out,
+            cost_usd=cost,
+            notes={"len_plan": len(plan_text)},
+        )
         return StageResult(ok=True, data={"plan": plan_text}, trace=trace)

nl2sql/repair.py

@@ -14,16 +14,26 @@ When repairing:
 Return only the corrected SQL.
 """
 
+
 class Repair:
     name = "repair"
+
     def __init__(self, llm: LLMProvider):
         self.llm = llm
 
-    def run(self, sql:str, error_msg: str, schema_preview: str) -> StageResult:
+    def run(self, sql: str, error_msg: str, schema_preview: str) -> StageResult:
         t0 = time.perf_counter()
-        fixed_sql, t_in, t_out, cost = self.llm.repair(sql=sql, error_msg=f"{GUIDELINES}\n\n{error_msg}",
-                                                      schema_preview=schema_preview)
-        trace = StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000,
-                           token_in=t_in, token_out=t_out, cost_usd=cost,
-                           notes={"old_sql_len": len(sql), "new_sql_len": len(fixed_sql)})
+        fixed_sql, t_in, t_out, cost = self.llm.repair(
+            sql=sql,
+            error_msg=f"{GUIDELINES}\n\n{error_msg}",
+            schema_preview=schema_preview,
+        )
+        trace = StageTrace(
+            stage=self.name,
+            duration_ms=(time.perf_counter() - t0) * 1000,
+            token_in=t_in,
+            token_out=t_out,
+            cost_usd=cost,
+            notes={"old_sql_len": len(sql), "new_sql_len": len(fixed_sql)},
+        )
         return StageResult(ok=True, data={"sql": fixed_sql}, trace=trace)

nl2sql/safety.py

@@ -4,7 +4,7 @@ from nl2sql.types import StageResult, StageTrace
 
 # --- Regex utils ---
 _COMMENT_BLOCK = re.compile(r"/\*.*?\*/", re.DOTALL)
-_COMMENT_LINE  = re.compile(r"--.*?$", re.MULTILINE)
+_COMMENT_LINE = re.compile(r"--.*?$", re.MULTILINE)
 # string literals (single & double quotes), allow escaped quotes
 _STRING_SINGLE = re.compile(r"'([^'\\]|\\.)*'", re.DOTALL)
 _STRING_DOUBLE = re.compile(r'"([^"\\]|\\.)*"', re.DOTALL)
@@ -18,20 +18,24 @@ _FORBIDDEN = re.compile(
 # allow: SELECT ...   or   WITH <cte...> SELECT ...
 _ALLOW_SELECT = re.compile(r"^(?:WITH\b.*?\)\s*)?SELECT\b", re.IGNORECASE | re.DOTALL)
 
+
 def _strip_comments(s: str) -> str:
     s = _COMMENT_BLOCK.sub(" ", s)
     s = _COMMENT_LINE.sub(" ", s)
     return s
 
+
 def _mask_strings(s: str) -> str:
     s = _STRING_SINGLE.sub("'X'", s)
     s = _STRING_DOUBLE.sub('"X"', s)
     return s
 
+
 def _split_statements(s: str) -> list[str]:
     parts = [p.strip() for p in s.split(";")]
     return [p for p in parts if p]
 
+
 class Safety:
     name = "safety"
 
@@ -46,7 +50,9 @@ class Safety:
             return StageResult(
                 ok=False,
                 error=["Multiple statements detected"],
-                trace=StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000),
+                trace=StageTrace(
+                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
+                ),
             )
 
         body = stmts[0]
@@ -55,14 +61,18 @@ class Safety:
             return StageResult(
                 ok=False,
                 error=["Forbidden keyword detected"],
-                trace=StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000),
+                trace=StageTrace(
+                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
+                ),
             )
 
         if not _ALLOW_SELECT.match(body):
             return StageResult(
                 ok=False,
                 error=["Non-SELECT statement"],
-                trace=StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000),
+                trace=StageTrace(
+                    stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
+                ),
             )
 
         return StageResult(
@@ -71,5 +81,7 @@ class Safety:
                 "sql": sql.strip(),
                 "rationale": "Statement validated as SELECT-only (strings/comments ignored).",
             },
-            trace=StageTrace(stage=self.name, duration_ms=(time.perf_counter()-t0)*1000),
+            trace=StageTrace(
+                stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
+            ),
         )

nl2sql/stubs.py

@@ -1,31 +1,37 @@
 from nl2sql.types import StageResult, StageTrace
 
+
 class NoOpExecutor:
     name = "executor"
+
     def run(self, sql: str) -> StageResult:
         # pretend success, return empty result set
         return StageResult(
             ok=True,
             data={"rows": [], "columns": []},
-            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True})
+            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True}),
         )
 
+
 class NoOpVerifier:
     name = "verifier"
+
     def run(self, sql: str, exec_result: StageResult) -> StageResult:
         # always verified for legacy tests
         return StageResult(
             ok=True,
             data={"verified": True},
-            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True})
+            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True}),
         )
 
+
 class NoOpRepair:
     name = "repair"
+
     def run(self, sql: str, error_msg: str, schema_preview: str) -> StageResult:
         # return original SQL unchanged
         return StageResult(
             ok=True,
             data={"sql": sql},
-            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True})
+            trace=StageTrace(stage=self.name, duration_ms=0.0, notes={"noop": True}),
         )

nl2sql/types.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, List
 
+
 @dataclass(frozen=True)
 class StageTrace:
     stage: str
@@ -10,6 +11,7 @@ class StageTrace:
     token_out: Optional[int] = None
     cost_usd: Optional[float] = None
 
+
 @dataclass(frozen=True)
 class StageResult:
     ok: bool

nl2sql/verifier.py

@@ -2,15 +2,20 @@ import sqlglot
 from sqlglot import expressions as exp
 from nl2sql.types import StageResult, StageTrace
 
+
 class Verifier:
     name = "verifier"
 
     def run(self, sql: str, exec_result: StageResult) -> StageResult:
         if not exec_result.ok:
-            return StageResult(ok=False, data=None,
-                               trace=StageTrace(stage=self.name, duration_ms=0,
-                               notes={"reason": "execution_error"}),
-                               error=exec_result.errors)
+            return StageResult(
+                ok=False,
+                data=None,
+                trace=StageTrace(
+                    stage=self.name, duration_ms=0, notes={"reason": "execution_error"}
+                ),
+                error=exec_result.errors,
+            )
 
         # Rule 1: check SELECT / GROUP consistency
         issues = []
@@ -25,9 +30,16 @@ class Verifier:
             issues.append(f"Parse error during verification: {e}")
 
         if issues:
-            return StageResult(ok=False, data=None,
-                               trace=StageTrace(stage=self.name, duration_ms=0,
-                               notes={"issues": issues}),
-                               error=issues)
-        return StageResult(ok=True, data={"verified": True},
-                           trace=StageTrace(stage=self.name, duration_ms=0))
+            return StageResult(
+                ok=False,
+                data=None,
+                trace=StageTrace(
+                    stage=self.name, duration_ms=0, notes={"issues": issues}
+                ),
+                error=issues,
+            )
+        return StageResult(
+            ok=True,
+            data={"verified": True},
+            trace=StageTrace(stage=self.name, duration_ms=0),
+        )

tests/conftest.py

@@ -4,4 +4,4 @@ from dotenv import load_dotenv
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 ENV_PATH = os.path.join(ROOT_DIR, ".env")
 
-load_dotenv(dotenv_path=ENV_PATH)
+load_dotenv(dotenv_path=ENV_PATH)

tests/test_ambiguity.py

@@ -2,18 +2,23 @@ from nl2sql.ambiguity_detector import AmbiguityDetector
 from nl2sql.types import StageResult
 from app.routers import nl2sql
 
+
 def test_detects_ambiguous_terms():
     det = AmbiguityDetector()
     res = det.detect("Show me recent top singers", "table: singer(id,name,age)")
     assert len(res) >= 1
     assert "recent" in res[0].lower()
 
+
 def test_not_false_positive():
     det = AmbiguityDetector()
     res = det.detect("List all singers older than 30", "table: singer(id, name, age)")
     assert res == []
 
+
 def test_ambiguity_response():
-    fake_result = StageResult(ok=True, data={"ambiguous": True, "questions": ["Clarify column?"]})
+    fake_result = StageResult(
+        ok=True, data={"ambiguous": True, "questions": ["Clarify column?"]}
+    )
     response = nl2sql._to_dict(fake_result.data)
-    assert response["ambiguous"] is True
+    assert response["ambiguous"] is True

tests/test_executor.py

@@ -1,9 +1,11 @@
 from nl2sql.executor import Executor
 from adapters.db.sqlite_adapter import SQLiteAdapter
 
+
 def test_executor_runs_select(tmp_path):
     db_path = tmp_path / "test.db"
     import sqlite3
+
     conn = sqlite3.connect(db_path)
     conn.execute("CREATE TABLE users(id INT, name TEXT);")
     conn.execute("INSERT INTO users VALUES (1, 'Alice');")

tests/test_generator.py

@@ -5,6 +5,7 @@ from nl2sql.types import StageResult
 
 # --- Dummy LLMs (respect the 5-tuple contract) --------------------------------
 
+
 class LLM_OK:
     def generate_sql(self, **kwargs):
         # contract: (sql, rationale, t_in, t_out, cost)
@@ -37,11 +38,12 @@ class LLM_CONTRACT_SHORT:
 
 # --- Parametrized negative cases ----------------------------------------------
 
+
 @pytest.mark.parametrize(
     "llm, err_keyword",
     [
-        (LLM_EMPTY_SQL(), "empty"),          # empty or non-string sql
-        (LLM_NON_SELECT(), "non-select"),    # generated non-SELECT
+        (LLM_EMPTY_SQL(), "empty"),  # empty or non-string sql
+        (LLM_NON_SELECT(), "non-select"),  # generated non-SELECT
         (LLM_CONTRACT_NONE(), "contract violation"),
         (LLM_CONTRACT_SHORT(), "contract violation"),
     ],
@@ -52,7 +54,7 @@ def test_generator_errors_do_not_create_trace(llm, err_keyword):
         user_query="show all singers",
         schema_preview="CREATE TABLE singer(id int, name text);",
         plan_text="-- plan --",
-        clarify_answers={}
+        clarify_answers={},
     )
     assert isinstance(r, StageResult)
     assert r.ok is False
@@ -65,13 +67,14 @@ def test_generator_errors_do_not_create_trace(llm, err_keyword):
 
 # --- Positive case (success) ---------------------------------------------------
 
+
 def test_generator_success_has_valid_trace_and_data():
     gen = Generator(llm=LLM_OK())
     r = gen.run(
         user_query="show all singers",
         schema_preview="CREATE TABLE singer(id int, name text);",
         plan_text="-- plan --",
-        clarify_answers={}
+        clarify_answers={},
     )
 
     # Basic success checks

tests/test_nl2sql_router.py

@@ -9,8 +9,10 @@ client = TestClient(app)
 def fake_trace(stage: str):
     return StageTrace(stage=stage, duration_ms=10.0)
 
+
 path = app.url_path_for("nl2sql_handler")
 
+
 # --- 1) Clarify / ambiguity case ---------------------------------------------
 def test_ambiguity_route(monkeypatch):
     from app.routers import nl2sql
@@ -47,7 +49,9 @@ def test_error_route(monkeypatch):
     from app.routers import nl2sql
 
     def fake_run(*args, **kwargs):
-        return StageResult(ok=False, error=["Bad SQL"], data={"traces": [fake_trace("safety")]})
+        return StageResult(
+            ok=False, error=["Bad SQL"], data={"traces": [fake_trace("safety")]}
+        )
 
     monkeypatch.setattr(nl2sql._pipeline, "run", fake_run)
 

tests/test_openai_provider.py

@@ -6,21 +6,23 @@ from adapters.llm.openai_provider import OpenAIProvider
 # Helper class to fake the completion object returned by OpenAI SDK
 class FakeCompletion:
     def __init__(self, content: str, prompt_tokens=5, completion_tokens=7):
-        self.choices = [type("Choice", (), {"message": type("Msg", (), {"content": content})})]
-        self.usage = type("Usage", (), {
-            "prompt_tokens": prompt_tokens,
-            "completion_tokens": completion_tokens
-        })
+        self.choices = [
+            type("Choice", (), {"message": type("Msg", (), {"content": content})})
+        ]
+        self.usage = type(
+            "Usage",
+            (),
+            {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens},
+        )
 
 
 # --- Case 1: clean valid JSON --------------------------------------------------
 def test_generate_sql_valid_json(monkeypatch):
     provider = OpenAIProvider()
 
-    fake_content = json.dumps({
-        "sql": "SELECT * FROM singer;",
-        "rationale": "List all singers."
-    })
+    fake_content = json.dumps(
+        {"sql": "SELECT * FROM singer;", "rationale": "List all singers."}
+    )
     fake_completion = FakeCompletion(fake_content)
 
     # Monkeypatch client.chat.completions.create
@@ -33,7 +35,7 @@ def test_generate_sql_valid_json(monkeypatch):
         user_query="show all singers",
         schema_preview="CREATE TABLE singer(id int, name text);",
         plan_text="-- plan --",
-        clarify_answers={}
+        clarify_answers={},
     )
 
     assert sql.strip().lower().startswith("select")
@@ -48,7 +50,7 @@ def test_generate_sql_recover_from_partial_json(monkeypatch):
     provider = OpenAIProvider()
 
     # invalid JSON with text around it
-    fake_content = "Here is the result:\n{ \"sql\": \"SELECT * FROM users;\", \"rationale\": \"list users\" }\nThanks!"
+    fake_content = 'Here is the result:\n{ "sql": "SELECT * FROM users;", "rationale": "list users" }\nThanks!'
     fake_completion = FakeCompletion(fake_content)
 
     def fake_create(*args, **kwargs):
@@ -59,7 +61,7 @@ def test_generate_sql_recover_from_partial_json(monkeypatch):
     sql, rationale, *_ = provider.generate_sql(
         user_query="show all users",
         schema_preview="CREATE TABLE users(id int, name text);",
-        plan_text="-- plan --"
+        plan_text="-- plan --",
     )
 
     assert sql.lower().startswith("select")
@@ -83,5 +85,5 @@ def test_generate_sql_invalid_json(monkeypatch):
         provider.generate_sql(
             user_query="show X",
             schema_preview="CREATE TABLE t(id int);",
-            plan_text="-- plan --"
+            plan_text="-- plan --",
         )

tests/test_pipeline_integration.py

@@ -5,8 +5,10 @@ from nl2sql.types import StageResult, StageTrace
 
 # --- Dummy stages to isolate pipeline -----------------------------------------
 
+
 class DummyDetector:
     """Simulates ambiguity detector stage."""
+
     def __init__(self, ambiguous=False):
         self.ambiguous = ambiguous
 
@@ -17,6 +19,7 @@ class DummyDetector:
 
 class DummyPlanner:
     """Simulates planner stage."""
+
     def run(self, *, user_query, schema_preview):
         trace = StageTrace(stage="planner", duration_ms=1.0)
         if "fail_plan" in user_query:
@@ -26,17 +29,21 @@ class DummyPlanner:
 
 class DummyGenerator:
     """Simulates generator stage."""
+
     def run(self, *, user_query, schema_preview, plan_text, clarify_answers):
         trace = StageTrace(stage="generator", duration_ms=1.0)
         if "fail_gen" in user_query:
             return StageResult(ok=False, error=["Generator failed"], trace=trace)
         sql = "SELECT * FROM singer;"
         rationale = "List all singers."
-        return StageResult(ok=True, data={"sql": sql, "rationale": rationale}, trace=trace)
+        return StageResult(
+            ok=True, data={"sql": sql, "rationale": rationale}, trace=trace
+        )
 
 
 class DummySafety:
     """Simulates safety stage."""
+
     def check(self, sql):
         trace = StageTrace(stage="safety", duration_ms=1.0)
         if "DROP" in sql.upper():
@@ -50,12 +57,12 @@ def test_pipeline_success():
         detector=DummyDetector(ambiguous=False),
         planner=DummyPlanner(),
         generator=DummyGenerator(),
-        safety=DummySafety()
+        safety=DummySafety(),
     )
 
     r = pipeline.run(
         user_query="show all singers",
-        schema_preview="CREATE TABLE singer(id int, name text);"
+        schema_preview="CREATE TABLE singer(id int, name text);",
     )
 
     assert isinstance(r, StageResult)
@@ -73,13 +80,10 @@ def test_pipeline_ambiguity():
         detector=DummyDetector(ambiguous=True),
         planner=DummyPlanner(),
         generator=DummyGenerator(),
-        safety=DummySafety()
+        safety=DummySafety(),
     )
 
-    r = pipeline.run(
-        user_query="show data",
-        schema_preview="CREATE TABLE x(id int);"
-    )
+    r = pipeline.run(user_query="show data", schema_preview="CREATE TABLE x(id int);")
 
     assert isinstance(r, StageResult)
     assert r.ok is True
@@ -93,11 +97,10 @@ def test_pipeline_plan_fail():
         detector=DummyDetector(),
         planner=DummyPlanner(),
         generator=DummyGenerator(),
-        safety=DummySafety()
+        safety=DummySafety(),
     )
     r = pipeline.run(
-        user_query="fail_plan",
-        schema_preview="CREATE TABLE singer(id int);"
+        user_query="fail_plan", schema_preview="CREATE TABLE singer(id int);"
     )
     assert isinstance(r, StageResult)
     assert r.ok is False
@@ -110,11 +113,10 @@ def test_pipeline_gen_fail():
         detector=DummyDetector(),
         planner=DummyPlanner(),
         generator=DummyGenerator(),
-        safety=DummySafety()
+        safety=DummySafety(),
     )
     r = pipeline.run(
-        user_query="fail_gen",
-        schema_preview="CREATE TABLE singer(id int);"
+        user_query="fail_gen", schema_preview="CREATE TABLE singer(id int);"
     )
     assert r.ok is False
     assert "Generator failed" in " ".join(r.error or [])
@@ -126,17 +128,18 @@ def test_pipeline_safety_fail():
         def run(self, **kw):
             trace = StageTrace(stage="generator", duration_ms=1.0)
             # Generate a DROP TABLE → unsafe
-            return StageResult(ok=True, data={"sql": "DROP TABLE x;", "rationale": "oops"}, trace=trace)
+            return StageResult(
+                ok=True, data={"sql": "DROP TABLE x;", "rationale": "oops"}, trace=trace
+            )
 
     pipeline = Pipeline(
         detector=DummyDetector(),
         planner=DummyPlanner(),
         generator=UnsafeGen(),
-        safety=DummySafety()
+        safety=DummySafety(),
     )
     r = pipeline.run(
-        user_query="drop something",
-        schema_preview="CREATE TABLE x(id int);"
+        user_query="drop something", schema_preview="CREATE TABLE x(id int);"
     )
     assert r.ok is False
     assert "unsafe" in " ".join(r.error or []).lower()

tests/test_safety.py

@@ -2,7 +2,6 @@ from nl2sql.safety import Safety
 import pytest
 
 
-
 def test_safety_allows_select():
     s = Safety()
     result = s.check("SELECT * FROM users;")
@@ -10,6 +9,7 @@ def test_safety_allows_select():
     assert "sql" in result.data
     assert result.trace.stage == "safety"
 
+
 def test_safety_allows_with_select_cte():
     s = Safety()
     sql = """
@@ -21,12 +21,14 @@ def test_safety_allows_with_select_cte():
     r = s.check(sql)
     assert r.ok
 
+
 def test_safety_allows_select_with_comments_and_newlines():
     s = Safety()
     sql = "/* head */ \n -- inline\n SELECT 1; -- tail"
     r = s.check(sql)
     assert r.ok
 
+
 def test_safety_allows_keywords_inside_string_literals():
     s = Safety()
     sql = "SELECT 'DROP TABLE x' as note, 'delete from y' as text;"
@@ -40,32 +42,39 @@ def test_safety_blocks_delete():
     assert not result.ok
     assert any("Forbidden" in e or "Non-SELECT" in e for e in (result.error or []))
 
-@pytest.mark.parametrize("sql", [
-    "UPDATE users SET name='X' WHERE id=1;",
-    "INSERT INTO users(id) VALUES (1);",
-    "DROP TABLE users;",
-    "CREATE TABLE x(id INT);",
-    "ALTER TABLE users ADD COLUMN x INT;",
-    "ATTACH DATABASE 'hack.db' AS h;",
-    "PRAGMA journal_mode=WAL;",
-])
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "UPDATE users SET name='X' WHERE id=1;",
+        "INSERT INTO users(id) VALUES (1);",
+        "DROP TABLE users;",
+        "CREATE TABLE x(id INT);",
+        "ALTER TABLE users ADD COLUMN x INT;",
+        "ATTACH DATABASE 'hack.db' AS h;",
+        "PRAGMA journal_mode=WAL;",
+    ],
+)
 def test_safety_blocks_forbidden_statements(sql):
     s = Safety()
     res = s.check(sql)
     assert not res.ok
 
+
 def test_safety_blocks_stacked_delete_after_select():
     s = Safety()
     sql = "SELECT * FROM users; DELETE FROM users;"
     r = s.check(sql)
     assert not r.ok
 
+
 def test_safety_blocks_stacked_delete_with_spaces():
     s = Safety()
     sql = "SELECT * FROM users ;   \n  DELETE users;"
     r = s.check(sql)
     assert not r.ok
 
+
 def test_safety_blocks_delete_inside_cte():
     s = Safety()
     sql = """
@@ -75,26 +84,35 @@ def test_safety_blocks_delete_inside_cte():
     r = s.check(sql)
     assert not r.ok
 
-@pytest.mark.parametrize("sql", [
-    "/*D*/ROP TABLE users;",
-    "PR/*x*/AGMA journal_mode=WAL;",
-    "AL/* comment */TER TABLE x ADD COLUMN y INT;",
-])
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "/*D*/ROP TABLE users;",
+        "PR/*x*/AGMA journal_mode=WAL;",
+        "AL/* comment */TER TABLE x ADD COLUMN y INT;",
+    ],
+)
 def test_safety_blocks_comment_obfuscation(sql):
     s = Safety()
     r = s.check(sql)
     assert not r.ok
 
-@pytest.mark.parametrize("sql", [
-    "pragma journal_mode=WAL;",  # lower-case
-    "  PRAGMA  user_version = 5 ; ",
-    "\nATTACH DATABASE 'hack.db' AS h;",
-])
+
+@pytest.mark.parametrize(
+    "sql",
+    [
+        "pragma journal_mode=WAL;",  # lower-case
+        "  PRAGMA  user_version = 5 ; ",
+        "\nATTACH DATABASE 'hack.db' AS h;",
+    ],
+)
 def test_safety_blocks_forbidden_case_and_spacing(sql):
     s = Safety()
     r = s.check(sql)
     assert not r.ok
 
+
 def test_safety_blocks_multiple_nonempty_statements_even_if_second_is_comment():
     s = Safety()
     sql = "SELECT 1;  -- now do something bad\n"

tests/test_stage_types.py

@@ -1,16 +1,19 @@
 from nl2sql.types import StageResult, StageTrace
 
+
 def test_error_response():
     r = StageResult(ok=False, error=["Syntax error"])
     assert not r.ok
     assert r.error == ["Syntax error"]
 
+
 def test_trace_dataclass_structure():
     t = StageTrace(stage="planner", duration_ms=12.5, token_in=10, token_out=20)
     assert t.stage == "planner"
     assert isinstance(t.duration_ms, float)
     assert t.token_out == 20
 
+
 def test_stage_result_defaults():
     r = StageResult(ok=True)
     assert r.ok

ui/benchmark_app.py

@@ -22,8 +22,8 @@ df = pd.DataFrame(rows)
 st.subheader("Aggregate Metrics")
 col1, col2, col3, col4 = st.columns(4)
 col1.metric("Total Queries", len(df))
-col2.metric("Execution Accuracy", f"{df['exec_acc'].mean()*100:.1f}%")
-col3.metric("Safety Violations", f"{df['safe_fail'].mean()*100:.1f}%")
+col2.metric("Execution Accuracy", f"{df['exec_acc'].mean() * 100:.1f}%")
+col3.metric("Safety Violations", f"{df['safe_fail'].mean() * 100:.1f}%")
 col4.metric("Average Latency (ms)", f"{df['latency_ms'].mean():.0f}")
 
 # 3. Latency Distribution
@@ -33,13 +33,23 @@ st.plotly_chart(fig1, use_container_width=True)
 
 # 4. Cost vs Accuracy
 st.subheader("Cost vs Execution Accuracy")
-fig2 = px.scatter(df, x="cost_usd", y="exec_acc", color="provider",
-                  title="Trade-off: Cost vs Accuracy", hover_data=["query"])
+fig2 = px.scatter(
+    df,
+    x="cost_usd",
+    y="exec_acc",
+    color="provider",
+    title="Trade-off: Cost vs Accuracy",
+    hover_data=["query"],
+)
 st.plotly_chart(fig2, use_container_width=True)
 
 # 5. Repair Stats
 if "repair_attempts" in df.columns:
     st.subheader("Repair Attempts")
-    fig3 = px.bar(df.groupby("repair_attempts").size().reset_index(name="count"),
-                  x="repair_attempts", y="count", title="Number of Repair Attempts per Query")
+    fig3 = px.bar(
+        df.groupby("repair_attempts").size().reset_index(name="count"),
+        x="repair_attempts",
+        y="count",
+        title="Number of Repair Attempts per Query",
+    )
     st.plotly_chart(fig3, use_container_width=True)