refactor(core): DI-ready Pipeline; add registry + YAML factory + typed trace/result

app/routers/nl2sql.py

@@ -7,10 +7,10 @@ import os
 from pathlib import Path
 import time
 import uuid
-from typing import Any, Dict, Optional, TypedDict, Union, Protocol, cast
+from typing import Any, Dict, Optional, TypedDict, Union, Protocol, cast, List
 
 # --- Third-party ---
-from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Depends
+from fastapi import APIRouter, HTTPException, Request, UploadFile, File
 
 # --- Local ---
 from app.schemas import NL2SQLRequest, NL2SQLResponse, ClarifyResponse
@@ -25,6 +25,10 @@ from nl2sql.repair import Repair
 from adapters.llm.openai_provider import OpenAIProvider
 from adapters.db.sqlite_adapter import SQLiteAdapter
 from adapters.db.postgres_adapter import PostgresAdapter
+from nl2sql.pipeline_factory import (
+    pipeline_from_config,
+    pipeline_from_config_with_adapter,
+)
 
 
 # Stable public re-exports
@@ -53,6 +57,10 @@ _DB_MAP_PATH.parent.mkdir(parents=True, exist_ok=True)
 UPLOAD_DIR = Path("data/uploads")
 UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
 
+CONFIG_PATH = os.getenv("PIPELINE_CONFIG", "configs/sqlite_pipeline.yaml")
+# Build a default pipeline once from config; adapter inside the config will be used.
+_PIPELINE = pipeline_from_config(CONFIG_PATH)
+
 
 class DBEntry(TypedDict):
     path: str
@@ -110,42 +118,55 @@ _load_db_map()
 # -------------------------------
 # Adapter selection (lazy)
 # -------------------------------
+# ---------- SELECT ADAPTER ----------
 def _select_adapter(db_id: Optional[str]) -> Union[PostgresAdapter, SQLiteAdapter]:
     """
-    Resolve a DB adapter:
-      - postgres: requires POSTGRES_DSN
-      - sqlite with db_id: uploaded file or fallback locations
-      - sqlite default: DEFAULT_SQLITE_PATH must exist
+    Resolve a DB adapter based on module-level DB_MODE and an optional db_id.
+
+    - postgres mode:
+        requires POSTGRES_DSN in env
+    - sqlite mode:
+        if db_id provided, resolve file by:
+            1) absolute path (if user supplied a full path)
+            2) uploads/{db_id}.sqlite
+            3) uploads/{db_id}.db
+            4) data/{db_id}.sqlite
+            5) data/{db_id}.db
+        else fallback to DEFAULT_SQLITE_PATH
     """
-    mode = os.getenv("DB_MODE", "sqlite").lower()
-    if mode == "postgres":
+    if DB_MODE == "postgres":
         dsn = os.environ.get("POSTGRES_DSN")
         if not dsn:
             raise HTTPException(status_code=500, detail="POSTGRES_DSN env is missing")
         return PostgresAdapter(dsn)
 
     # sqlite mode
-    _cleanup_db_map()
     if db_id:
-        # Check runtime map
-        entry = _DB_MAP.get(db_id)
-        candidates = []
-        if entry and os.path.exists(entry["path"]):
-            candidates.append(entry["path"])
-        # Fallback locations based on convention
-        candidates.append(os.path.join(_DB_UPLOAD_DIR, f"{db_id}.sqlite"))
-        candidates.append(str(UPLOAD_DIR / f"{db_id}.sqlite"))
-
-        for p in candidates:
-            if p and os.path.exists(p):
-                return SQLiteAdapter(p)
+        # 1) absolute path
+        p = Path(db_id)
+        candidates: List[Path] = []
+        if p.is_absolute():
+            candidates.append(p)
+
+        # 2) uploads/
+        candidates.append(UPLOAD_DIR / f"{db_id}.sqlite")
+        candidates.append(UPLOAD_DIR / f"{db_id}.db")
+
+        # 3) data/
+        candidates.append(Path("data") / f"{db_id}.sqlite")
+        candidates.append(Path("data") / f"{db_id}.db")
+
+        for c in candidates:
+            if c.exists() and c.is_file():
+                return SQLiteAdapter(str(c))
 
         raise HTTPException(status_code=400, detail="invalid db_id (file not found)")
 
-    # default sqlite
-    if not Path(DEFAULT_SQLITE_PATH).exists():
-        raise HTTPException(status_code=500, detail="default DB not found")
-    return SQLiteAdapter(DEFAULT_SQLITE_PATH)
+    # default sqlite fallback
+    default_path = Path(DEFAULT_SQLITE_PATH)
+    if not default_path.exists():
+        raise HTTPException(status_code=500, detail="default SQLite DB not found")
+    return SQLiteAdapter(str(default_path))
 
 
 # -------------------------------
@@ -289,57 +310,52 @@ async def upload_db(file: UploadFile = File(...)):
 # Main NL2SQL endpoint
 # -------------------------------
 @router.post("", name="nl2sql_handler")
-def nl2sql_handler(
-    request: NL2SQLRequest,
-    run: Runner = Depends(get_runner),
-):
+def nl2sql_handler(request: NL2SQLRequest):
     """
-    Handles NL→SQL conversion requests.
-    Uses dependency-injected pipeline runner (get_runner).
-    If db_id provided → builds a temporary per-request pipeline.
+    NL→SQL handler using YAML-driven DI. If 'db_id' is provided, we override only the adapter
+    while keeping all other stages from the YAML config intact.
     """
     db_id = getattr(request, "db_id", None)
-    provided_preview: Optional[str] = cast(
-        Optional[str], getattr(request, "schema_preview", None)
+    provided_preview = (
+        cast(Optional[str], getattr(request, "schema_preview", None)) or ""
     )
 
-    # Select pipeline (DI default vs per-request)
+    # Choose runner: default pipeline from YAML OR per-request override with a specific adapter
     if db_id:
         adapter = _select_adapter(db_id)
-        pipeline = _build_pipeline(adapter)
-        derived_preview = _derive_schema_preview(adapter)
-        runner: Runner = pipeline.run
-        final_preview = provided_preview or derived_preview or ""
+        # Build a temporary pipeline from YAML but bind the per-request adapter
+        pipeline = pipeline_from_config_with_adapter(CONFIG_PATH, adapter=adapter)
+        runner = pipeline.run
+        final_preview = provided_preview  # keep simple; derive only if you have a SQLite schema helper
     else:
-        runner = run
-        final_preview = provided_preview or ""
+        runner = _PIPELINE.run
+        final_preview = provided_preview
 
-    # Execute safely
+    # Execute pipeline
     try:
         result = runner(user_query=request.query, schema_preview=final_preview)
     except Exception as exc:
         raise HTTPException(status_code=500, detail=f"Pipeline crash: {exc!s}")
 
+    # Type sanity
     if not isinstance(result, FinalResult):
         raise HTTPException(status_code=500, detail="Pipeline returned unexpected type")
 
-    # Ambiguous → 200
+    # Ambiguity path → 200 with questions
     if result.ambiguous and (result.questions is not None):
         return ClarifyResponse(ambiguous=True, questions=result.questions)
 
-    # Error → 400 + dump
+    # Error path → 400 with joined details
     if (not result.ok) or result.error:
         print("❌ Pipeline failure dump:")
         print("  ok:", result.ok)
         print("  error:", result.error)
         print("  details:", result.details)
         print("  traces:", result.traces)
-        raise HTTPException(
-            status_code=400,
-            detail="; ".join(result.details or []) or (result.error or "Unknown error"),
-        )
+        message = "; ".join(result.details or []) or "Unknown error"
+        raise HTTPException(status_code=400, detail=message)
 
-    # Success → 200
+    # Success path → 200
     traces = [_round_trace(t) for t in (result.traces or [])]
     return NL2SQLResponse(
         ambiguous=False,

config/sqlite_pipeline.yaml

@@ -0,0 +1,11 @@
+detector: default
+planner: default
+generator: rules    # or "llm" when available
+safety: default
+executor: default
+verifier: basic
+repair: default
+
+adapter:
+  kind: sqlite
+  dsn: data/chinook.db

nl2sql/pipeline.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 import traceback
 from dataclasses import dataclass
 from typing import Dict, Any, Optional, List
+import time
 
 from nl2sql.types import StageResult
 from nl2sql.ambiguity_detector import AmbiguityDetector
@@ -31,6 +32,7 @@ class Pipeline:
     """
     NL2SQL Copilot pipeline.
     Stages return StageResult; final result is a type-safe FinalResult.
+    DI-ready: all dependencies are injected via __init__.
     """
 
     def __init__(
@@ -53,19 +55,26 @@ class Pipeline:
         self.repair = repair or NoOpRepair()
 
     # ------------------------------------------------------------
-    def _trace_list(self, *stages: StageResult) -> List[dict]:
-        traces = []
+    @staticmethod
+    def _trace_list(*stages: Optional[StageResult]) -> List[dict]:
+        """Collect .trace objects (as dict) from StageResult items if present."""
+        traces: List[dict] = []
         for s in stages:
             if not s:
                 continue
             t = getattr(s, "trace", None)
-            if t:
-                traces.append(t.__dict__)
+            if t is not None:
+                # t is likely a dataclass – expose as plain dict for JSON safety
+                traces.append(getattr(t, "__dict__", t))
         return traces
 
     # ------------------------------------------------------------
-    def _safe_stage(self, fn, **kwargs) -> StageResult:
-        """Run a stage safely; if it throws, catch and convert to StageResult."""
+    @staticmethod
+    def _safe_stage(fn, **kwargs) -> StageResult:
+        """
+        Run a stage safely; if it throws, return a StageResult(ok=False, error=[...]).
+        If fn returns a non-StageResult (e.g., dict), coerce to StageResult(ok=True, data=...).
+        """
         try:
             r = fn(**kwargs)
             if isinstance(r, StageResult):
@@ -75,6 +84,18 @@ class Pipeline:
             tb = traceback.format_exc()
             return StageResult(ok=False, data=None, trace=None, error=[f"{e}", tb])
 
+    # ------------------------------------------------------------
+    @staticmethod
+    def _mk_trace(
+        stage: str, duration_ms: float, notes: Optional[Dict[str, Any]] = None
+    ) -> dict:
+        """Create a normalized trace dict."""
+        return {
+            "stage": stage,
+            "duration_ms": float(duration_ms),
+            "notes": notes or {},
+        }
+
     # ------------------------------------------------------------
     def run(
         self,
@@ -88,11 +109,27 @@ class Pipeline:
         sql: Optional[str] = None
         rationale: Optional[str] = None
         verified: Optional[bool] = None
+
+        # Normalize inputs
         schema_preview = schema_preview or ""
+        clarify_answers = clarify_answers or {}
 
-        # --- 1) ambiguity detection ---
+        # --- 1) ambiguity detection (with explicit timing & trace) ---
         try:
+            t0 = time.perf_counter()
             questions = self.detector.detect(user_query, schema_preview)
+            t1 = time.perf_counter()
+            traces.append(
+                self._mk_trace(
+                    "detector",
+                    (t1 - t0) * 1000.0,
+                    {
+                        "ambiguous": bool(questions),
+                        "questions_len": len(questions or []),
+                    },
+                )
+            )
+
             if questions:
                 return FinalResult(
                     ok=True,
@@ -103,9 +140,11 @@ class Pipeline:
                     sql=None,
                     rationale=None,
                     verified=None,
-                    traces=[],
+                    traces=traces,
                 )
         except Exception as e:
+            # detector crash – mark as error but keep trace so far
+            traces.append(self._mk_trace("detector", 0.0, {"error": str(e)}))
             return FinalResult(
                 ok=False,
                 ambiguous=True,
@@ -115,7 +154,7 @@ class Pipeline:
                 sql=None,
                 rationale=None,
                 verified=None,
-                traces=[],
+                traces=traces,
             )
 
         # --- 2) planner ---
@@ -142,7 +181,7 @@ class Pipeline:
             user_query=user_query,
             schema_preview=schema_preview,
             plan_text=(r_plan.data or {}).get("plan"),
-            clarify_answers=clarify_answers or {},
+            clarify_answers=clarify_answers,
         )
         traces.extend(self._trace_list(r_gen))
         if not r_gen.ok:
@@ -183,7 +222,9 @@ class Pipeline:
         )
         traces.extend(self._trace_list(r_exec))
         if not r_exec.ok:
-            details.extend(r_exec.error or [])
+            # executor failure does not hard-fail the pipeline; accumulate details
+            if r_exec.error:
+                details.extend(r_exec.error)
 
         # --- 6) verifier ---
         r_ver = self._safe_stage(
@@ -203,13 +244,17 @@ class Pipeline:
                 )
                 traces.extend(self._trace_list(r_fix))
                 if not r_fix.ok:
+                    # repair failed – stop trying further
                     break
 
-                sql = (r_fix.data or {}).get("sql")
+                # re-run safety → executor → verifier on the fixed SQL
+                sql = (r_fix.data or {}).get("sql", sql)
+
                 r_safe = self._safe_stage(self.safety.run, sql=sql)
                 traces.extend(self._trace_list(r_safe))
                 if not r_safe.ok:
-                    details.extend(r_safe.error or [])
+                    if r_safe.error:
+                        details.extend(r_safe.error)
                     continue
 
                 r_exec = self._safe_stage(
@@ -217,7 +262,8 @@ class Pipeline:
                 )
                 traces.extend(self._trace_list(r_exec))
                 if not r_exec.ok:
-                    details.extend(r_exec.error or [])
+                    if r_exec.error:
+                        details.extend(r_exec.error)
                     continue
 
                 r_ver = self._safe_stage(
@@ -230,19 +276,19 @@ class Pipeline:
 
         # --- 8) fallback: verifier silent but executor succeeded ---
         if (verified is None or not verified) and not details:
-            any_exec = any(
-                t.get("stage") == "executor" and t.get("notes", {}).get("row_count")
+            any_exec_ok = any(
+                t.get("stage") == "executor" and (t.get("notes") or {}).get("row_count")
                 for t in traces
             )
-            if any_exec:
+            if any_exec_ok:
                 traces.append(
-                    {
-                        "stage": "pipeline",
-                        "notes": {
+                    self._mk_trace(
+                        "pipeline",
+                        0.0,
+                        {
                             "auto_fix": "verified=True (executor succeeded, verifier silent)"
                         },
-                        "duration_ms": 0.0,
-                    }
+                    )
                 )
                 verified = True
 
@@ -252,11 +298,11 @@ class Pipeline:
         err = has_errors and not bool(verified)
 
         traces.append(
-            {
-                "stage": "pipeline",
-                "notes": {"final_verified": verified, "details_len": len(details)},
-                "duration_ms": 0.0,
-            }
+            self._mk_trace(
+                "pipeline",
+                0.0,
+                {"final_verified": bool(verified), "details_len": len(details)},
+            )
         )
 
         return FinalResult(

nl2sql/pipeline_factory.py

@@ -0,0 +1,84 @@
+import yaml
+from typing import Any, Dict
+from nl2sql.pipeline import Pipeline
+from nl2sql.registry import (
+    DETECTORS,
+    PLANNERS,
+    GENERATORS,
+    SAFETIES,
+    EXECUTORS,
+    VERIFIERS,
+    REPAIRS,
+)
+from adapters.db.sqlite_adapter import SQLiteAdapter
+from adapters.db.postgres_adapter import PostgresAdapter
+from adapters.db.base import DBAdapter
+
+
+def _build_adapter(adapter_cfg: Dict[str, Any]) -> DBAdapter:
+    kind = adapter_cfg.get("kind", "sqlite")
+    if kind == "sqlite":
+        return SQLiteAdapter(adapter_cfg.get("dsn"))
+    if kind == "postgres":
+        return PostgresAdapter(**adapter_cfg)
+    raise ValueError(f"Unknown adapter kind: {kind}")
+
+
+def pipeline_from_config(path: str) -> Pipeline:
+    with open(path, "r", encoding="utf-8") as fh:
+        cfg: Dict[str, Any] = yaml.safe_load(fh)
+
+    detector = DETECTORS[cfg.get("detector", "default")]()
+    planner = PLANNERS[cfg.get("planner", "default")]()
+    generator = GENERATORS[cfg.get("generator", "rules")]()
+    safety = SAFETIES[cfg.get("safety", "default")]()
+    executor = EXECUTORS[cfg.get("executor", "default")]()
+    verifier = VERIFIERS[cfg.get("verifier", "basic")]()
+    repair = REPAIRS[cfg.get("repair", "default")]()
+
+    # If your Executor needs an adapter inside, set it there (common pattern):
+    adapter_cfg = cfg.get("adapter", {"kind": "sqlite", "dsn": "data/chinook.db"})
+    adapter = _build_adapter(adapter_cfg)
+    if hasattr(executor, "bind_adapter"):
+        executor.bind_adapter(adapter)
+    elif hasattr(executor, "adapter"):
+        executor.adapter = adapter  # fallback
+
+    return Pipeline(
+        detector=detector,
+        planner=planner,
+        generator=generator,
+        safety=safety,
+        executor=executor,
+        verifier=verifier,
+        repair=repair,
+    )
+
+
+def pipeline_from_config_with_adapter(path: str, *, adapter: DBAdapter) -> Pipeline:
+    """Same as pipeline_from_config, but force a specific adapter (per-request override)."""
+    with open(path, "r", encoding="utf-8") as fh:
+        cfg: Dict[str, Any] = yaml.safe_load(fh)
+
+    detector = DETECTORS[cfg.get("detector", "default")]()
+    planner = PLANNERS[cfg.get("planner", "default")]()
+    generator = GENERATORS[cfg.get("generator", "rules")]()
+    safety = SAFETIES[cfg.get("safety", "default")]()
+    executor = EXECUTORS[cfg.get("executor", "default")]()
+    verifier = VERIFIERS[cfg.get("verifier", "basic")]()
+    repair = REPAIRS[cfg.get("repair", "default")]()
+
+    if hasattr(executor, "bind_adapter"):
+        executor.bind_adapter(adapter)
+    elif hasattr(executor, "adapter"):
+        executor.adapter = adapter
+
+    return Pipeline(
+        detector=detector,
+        planner=planner,
+        generator=generator,
+        safety=safety,
+        executor=executor,
+        verifier=verifier,
+        repair=repair,
+    )

nl2sql/registry.py

@@ -0,0 +1,22 @@
+"""
+Registry mapping simple string keys to concrete component classes.
+Used by pipeline_factory to perform lightweight dependency injection.
+"""
+
+from typing import Dict, Type
+from nl2sql.ambiguity_detector import AmbiguityDetector
+from nl2sql.planner import Planner
+from nl2sql.generator import Generator
+from nl2sql.safety import Safety
+from nl2sql.executor import Executor
+from nl2sql.verifier import Verifier
+from nl2sql.repair import Repair
+
+# later you can add llm-aware generator variants, etc.
+PLANNERS: Dict[str, Type[Planner]] = {"default": Planner}
+DETECTORS: Dict[str, Type[AmbiguityDetector]] = {"default": AmbiguityDetector}
+GENERATORS: Dict[str, Type[Generator]] = {"rules": Generator}
+SAFETIES: Dict[str, Type[Safety]] = {"default": Safety}
+EXECUTORS: Dict[str, Type[Executor]] = {"default": Executor}
+VERIFIERS: Dict[str, Type[Verifier]] = {"basic": Verifier}
+REPAIRS: Dict[str, Type[Repair]] = {"default": Repair}

tests/test_nl2sql_router.py

@@ -252,3 +252,12 @@ def test_traces_are_rounded_to_ints():
         assert isinstance(traces[0]["duration_ms"], int)
     finally:
         app.dependency_overrides.pop(nl2sql.get_runner, None)
+
+
+def test_nl2sql_handler_returns_sql(monkeypatch):
+    payload = {"query": "Top 5 albums by sales"}
+    r = client.post("/nl2sql", json=payload)
+    assert r.status_code == 200
+    data = r.json()
+    assert "sql" in data
+    assert "traces" in data

tests/test_pipeline_factory.py

@@ -0,0 +1,32 @@
+from nl2sql.pipeline_factory import (
+    pipeline_from_config,
+    pipeline_from_config_with_adapter,
+)
+from adapters.db.sqlite_adapter import SQLiteAdapter
+
+
+def test_pipeline_from_config_builds_and_runs(tmp_path):
+    p = pipeline_from_config("configs/sqlite_pipeline.yaml")
+    result = p.run(user_query="Top 3 albums by sales")
+    assert result.sql is not None
+    assert isinstance(result.traces, list)
+
+
+def test_pipeline_from_config_with_adapter_override(tmp_path):
+    adapter = SQLiteAdapter("data/chinook.db")
+    p = pipeline_from_config_with_adapter(
+        "configs/sqlite_pipeline.yaml", adapter=adapter
+    )
+    result = p.run(user_query="Count customers")
+    assert "SELECT" in result.sql.upper()
+    assert isinstance(result.traces, list)
+
+
+def test_full_pipeline_from_yaml(monkeypatch):
+    from nl2sql.pipeline_factory import pipeline_from_config
+
+    p = pipeline_from_config("configs/sqlite_pipeline.yaml")
+    res = p.run(user_query="List all artists")
+    assert res.ok
+    assert isinstance(res.sql, str)
+    assert any(t["stage"] == "executor" for t in res.traces)