build(mypy): fix type errors and add safety guards for None values

.github/workflows/ci.yml

@@ -49,7 +49,7 @@ jobs:
         run: python -m ruff check .
 
       - name: Type check (mypy)
-        run: python -m mypy .
+        run: python -m mypy . --ignore-missing-imports --explicit-package-bases
 
       - name: Run tests
         run: python -m pytest -q

adapters/db/postgres_adapter.py

@@ -15,34 +15,56 @@ class PostgresAdapter(DBAdapter):
         self.dsn = dsn
 
     def preview_schema(self, limit_per_table: int = 0) -> str:
+        """
+        Return a simple textual preview of tables and their columns in public schema.
+        Example line: "- users (id:integer, name:text)"
+        """
+        lines: List[str] = []
         with psycopg.connect(self.dsn) as conn:
-            cur = conn.cursor()
-            cur.execute("""
-                SELECT table_name
-                FROM information_schema.tables
-                WHERE table_schema = 'public';
-            """)
-            tables = [t[0] for t in cur.fetchall()]
-            lines = []
-            for t in tables:
+            with conn.cursor() as cur:
+                # list tables
                 cur.execute(
                     """
-                    SELECT column_name, data_type
-                    FROM information_schema.columns
-                    WHERE table_name = %s;
-                """,
-                    (t,),
+                    SELECT table_name
+                    FROM information_schema.tables
+                    WHERE table_schema = 'public'
+                    ORDER BY table_name;
+                    """
                 )
-                cols = [f"{c[0]}:{c[1]}" for c in cur.fetchall()]
-                lines.append(f"- {t} ({', '.join(cols)})")
-            return "\n".join(lines)
+                table_rows = cur.fetchall() or []
+                tables: List[str] = [t[0] for t in table_rows if t and t[0]]
+
+                for t in tables:
+                    # list columns for table t
+                    cur.execute(
+                        """
+                        SELECT column_name, data_type
+                        FROM information_schema.columns
+                        WHERE table_schema = 'public' AND table_name = %s
+                        ORDER BY ordinal_position;
+                        """,
+                        (t,),
+                    )
+                    col_rows = cur.fetchall() or []
+                    # guard against None; build "name:type"
+                    cols: List[str] = [
+                        f"{c[0]}:{c[1]}" for c in col_rows if c and len(c) >= 2
+                    ]
+                    lines.append(f"- {t} ({', '.join(cols)})")
+
+        return "\n".join(lines)
 
     def execute(self, sql: str) -> Tuple[List[Tuple[Any, ...]], List[str]]:
-        if not sql.strip().lower().startswith("select"):
+        """
+        Execute a read-only SELECT query and return (rows, columns).
+        """
+        if not sql or not sql.strip().lower().startswith("select"):
             raise ValueError("Only SELECT statements are allowed.")
+
         with psycopg.connect(self.dsn) as conn:
-            cur = conn.cursor()
-            cur.execute(sql)
-            rows = cur.fetchall()
-            cols = [desc[0] for desc in cur.description]
-            return rows, cols
+            with conn.cursor() as cur:
+                cur.execute(sql)
+                rows = cur.fetchall() or []
+                desc = cur.description or ()
+                cols: List[str] = [d[0] for d in desc if d]
+                return rows, cols

app/routers/nl2sql.py

@@ -13,11 +13,13 @@ from nl2sql.repair import Repair
 from adapters.db.sqlite_adapter import SQLiteAdapter
 from adapters.db.postgres_adapter import PostgresAdapter
 import os
+from typing import Union
 
 
 router = APIRouter(prefix="/nl2sql")
 
 
+_db: Union[PostgresAdapter, SQLiteAdapter]
 if os.getenv("DB_MODE", "sqlite") == "postgres":
     _db = PostgresAdapter(os.environ["POSTGRES_DSN"])
 else:

app/schemas.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from typing import List, Optional, Any, Dict
+from typing import List, Optional, Any, Dict, Union
 
 
 class NL2SQLRequest(BaseModel):
@@ -19,9 +19,9 @@ class TraceModel(BaseModel):
 
 class NL2SQLResponse(BaseModel):
     ambiguous: bool = False
-    sql: str
+    sql: Optional[str] = None
     rationale: Optional[str] = None
-    traces: List[TraceModel] = []
+    traces: List[Union[TraceModel, dict]] = []
 
 
 class ClarifyResponse(BaseModel):

benchmarks/run.py

@@ -5,9 +5,10 @@ import os
 import json
 import time
 from pathlib import Path
+from typing import Iterable, List, Dict, Any
 
 # ---- app imports
-from nl2sql.pipeline import Pipeline
+from nl2sql.pipeline import Pipeline, FinalResult
 from nl2sql.ambiguity_detector import AmbiguityDetector
 from nl2sql.planner import Planner
 from nl2sql.generator import Generator
@@ -26,7 +27,10 @@ class DummyLLM:
     provider_id = "dummy-llm"
 
     def plan(self, *, user_query: str, schema_preview: str):
-        text = f"- understand question: {user_query}\n- identify tables\n- join if needed\n- filter\n- order/limit"
+        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(
@@ -68,11 +72,13 @@ def build_pipeline(db_path: Path, use_openai: bool) -> Pipeline:
     # DB adapter
     db = SQLiteAdapter(str(db_path))
     executor = Executor(db)
+
     # LLM provider
     if use_openai and os.getenv("OPENAI_API_KEY"):
         llm = OpenAIProvider()
     else:
         llm = DummyLLM()
+
     # stages
     detector = AmbiguityDetector()
     planner = Planner(llm)
@@ -80,6 +86,7 @@ def build_pipeline(db_path: Path, use_openai: bool) -> Pipeline:
     safety = Safety()
     verifier = Verifier()
     repair = Repair(llm)
+
     # pipeline
     return Pipeline(
         detector=detector,
@@ -92,33 +99,49 @@ def build_pipeline(db_path: Path, use_openai: bool) -> Pipeline:
     )
 
 
-def run_benchmark(queries, schema_preview, pipeline: Pipeline, outfile: Path):
-    results = []
+def _sum_cost(traces: Iterable[Dict[str, Any]]) -> float:
+    total = 0.0
+    for tr in traces:
+        try:
+            total += float(tr.get("cost_usd", 0.0))
+        except Exception:
+            # ignore bad values
+            pass
+    return total
+
+
+def _is_safe_fail(ok: bool, details: List[str] | None) -> float:
+    """Return 1.0 when pipeline failed due to unsafe SQL (heuristic)."""
+    if ok:
+        return 0.0
+    txt = " ".join(details or []).lower()
+    return 1.0 if "unsafe" in txt else 0.0
+
+
+def run_benchmark(
+    queries: List[str], schema_preview: str, pipeline: Pipeline, outfile: Path
+) -> None:
+    results: List[Dict[str, Any]] = []
     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
-        ok = (not r.get("ambiguous")) and ("error" not in r)
-
-        traces = r.get("traces", [])
-        cost_sum = 0.0
-        for t in traces:
-            try:
-                cost_sum += float(t.get("cost_usd", 0.0))
-            except Exception:
-                pass
+        res: FinalResult = pipeline.run(user_query=q, schema_preview=schema_preview)
+        latency_ms = (time.perf_counter() - t0) * 1000.0
+
+        ok = (not res.ambiguous) and (not res.error) and bool(res.ok)
+        traces = res.traces or []
+        cost_sum = _sum_cost(traces)
 
         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,
+                "safe_fail": _is_safe_fail(ok, res.details),
                 "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",
+                "provider": getattr(
+                    getattr(pipeline.generator, "llm", None), "provider_id", "unknown"
+                ),
             }
         )
 
@@ -129,7 +152,7 @@ def run_benchmark(queries, schema_preview, pipeline: Pipeline, outfile: Path):
     print(f"[OK] wrote {len(results)} rows → {outfile}")
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--outfile", default="benchmarks/results/demo.jsonl")
     parser.add_argument("--db", default="data/bench_demo.db")
@@ -140,9 +163,9 @@ def main():
     )
     args = parser.parse_args()
 
-    ROOT = Path(__file__).resolve().parents[1]  # project root
-    outfile = (ROOT / args.outfile).resolve()
-    db_path = (ROOT / args.db).resolve()
+    root = Path(__file__).resolve().parents[1]  # project root
+    outfile = (root / args.outfile).resolve()
+    db_path = (root / args.db).resolve()
 
     ensure_demo_db(db_path)
     pipe = build_pipeline(db_path, use_openai=args.use_openai)

nl2sql/pipeline.py

@@ -14,7 +14,6 @@ from nl2sql.repair import Repair
 from nl2sql.stubs import NoOpExecutor, NoOpRepair, NoOpVerifier
 
 
-# ---- NEW: FinalResult as domain-level, type-safe result ----
 @dataclass(frozen=True)
 class FinalResult:
     ok: bool
@@ -144,7 +143,7 @@ class Pipeline:
             self.generator.run,
             user_query=user_query,
             schema_preview=schema_preview,
-            plan_text=r_plan.data.get("plan"),
+            plan_text=(r_plan.data or {}).get("plan"),
             clarify_answers=clarify_answers or {},
         )
         traces.extend(self._trace_list(r_gen))
@@ -160,11 +159,10 @@ class Pipeline:
                 verified=None,
                 traces=traces,
             )
-        sql = r_gen.data.get("sql")
-        rationale = r_gen.data.get("rationale")
+        sql = (r_gen.data or {}).get("sql")
+        rationale = (r_gen.data or {}).get("rationale")
 
         # --- 4) safety
-        # fix: align with DummySafety signature → use .run (not .check)
         r_safe = self._safe_stage(self.safety.run, sql=sql)
         traces.extend(self._trace_list(r_safe))
         if not r_safe.ok:
@@ -181,13 +179,17 @@ class Pipeline:
             )
 
         # --- 5) executor
-        r_exec = self._safe_stage(self.executor.run, sql=r_safe.data.get("sql", sql))
+        r_exec = self._safe_stage(
+            self.executor.run, sql=(r_safe.data or {}).get("sql", sql)
+        )
         traces.extend(self._trace_list(r_exec))
         if not r_exec.ok:
             details.extend(r_exec.error or [])
 
         # --- 6) verifier
-        r_ver = self._safe_stage(self.verifier.run, sql=sql, exec_result=r_exec.data)
+        r_ver = self._safe_stage(
+            self.verifier.run, sql=sql, exec_result=(r_exec.data or {})
+        )
         traces.extend(self._trace_list(r_ver))
         verified = bool(r_ver.ok)
 
@@ -203,7 +205,7 @@ class Pipeline:
                 traces.extend(self._trace_list(r_fix))
                 if not r_fix.ok:
                     break
-                sql = r_fix.data.get("sql")
+                sql = (r_fix.data or {}).get("sql")
 
                 r_safe = self._safe_stage(self.safety.run, sql=sql)
                 traces.extend(self._trace_list(r_safe))
@@ -212,7 +214,7 @@ class Pipeline:
                     continue
 
                 r_exec = self._safe_stage(
-                    self.executor.run, sql=r_safe.data.get("sql", sql)
+                    self.executor.run, sql=(r_safe.data or {}).get("sql", sql)
                 )
                 traces.extend(self._trace_list(r_exec))
                 if not r_exec.ok:
@@ -220,7 +222,7 @@ class Pipeline:
                     continue
 
                 r_ver = self._safe_stage(
-                    self.verifier.run, sql=sql, exec_result=r_exec.data
+                    self.verifier.run, sql=sql, exec_result=(r_exec.data or {})
                 )
                 traces.extend(self._trace_list(r_ver))
                 verified = bool(r_ver.ok)

nl2sql/safety.py

@@ -86,3 +86,5 @@ class Safety:
                 stage=self.name, duration_ms=(time.perf_counter() - t0) * 1000
             ),
         )
+
+    run = check