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

nl2sql/pipeline_factory.py

@@ -1,4 +1,3 @@
-# nl2sql/pipeline_factory.py
 from __future__ import annotations
 
 import os
@@ -15,7 +14,15 @@ from nl2sql.registry import (
     VERIFIERS,
     REPAIRS,
 )
-from nl2sql.types import StageResult
+from nl2sql.types import StageResult, StageTrace
+
+from nl2sql.ambiguity_detector import AmbiguityDetector
+from nl2sql.planner import Planner
+from nl2sql.generator import Generator
+from nl2sql.executor import Executor
+from nl2sql.verifier import Verifier
+from nl2sql.repair import Repair
+
 from adapters.db.base import DBAdapter
 from adapters.db.sqlite_adapter import SQLiteAdapter
 from adapters.db.postgres_adapter import PostgresAdapter
@@ -35,15 +42,12 @@ def _build_adapter(adapter_cfg: Dict[str, Any]) -> DBAdapter:
         dsn = _require_str(adapter_cfg.get("dsn"), name="adapter.dsn")
         return SQLiteAdapter(dsn)
     if kind == "postgres":
-        # Pass through any kwargs your adapter expects (dsn, host, user, ...)
         return PostgresAdapter(**adapter_cfg)
     raise ValueError(f"Unknown adapter kind: {kind}")
 
 
 def _build_llm(llm_cfg: Optional[Dict[str, Any]] = None) -> Any:
-    """
-    Build the LLM provider. Under pytest we return None so stubs are used.
-    """
+    """Under pytest return None (stubs handle logic); otherwise real OpenAI provider."""
     if os.getenv("PYTEST_CURRENT_TEST"):
         return None
     _ = llm_cfg or {}
@@ -54,6 +58,25 @@ def _is_pytest() -> bool:
     return bool(os.getenv("PYTEST_CURRENT_TEST"))
 
 
+def _tr(
+    stage: str,
+    *,
+    duration_ms: int = 0,
+    notes: Optional[Dict[str, Any]] = None,
+    token_in: Optional[int] = None,
+    token_out: Optional[int] = None,
+    cost_usd: Optional[float] = None,
+) -> StageTrace:
+    return StageTrace(
+        stage=stage,
+        duration_ms=duration_ms,
+        notes=notes,
+        token_in=token_in,
+        token_out=token_out,
+        cost_usd=cost_usd,
+    )
+
+
 # ------------------------------ factory ------------------------------ #
 def pipeline_from_config(path: str) -> Pipeline:
     """
@@ -68,7 +91,6 @@ def pipeline_from_config(path: str) -> Pipeline:
     # --- Adapter ---
     adapter_cfg = cast(Dict[str, Any], cfg.get("adapter", {}))
     if is_pytest:
-        # Avoid filesystem errors during tests
         adapter_cfg = {"kind": "sqlite", "dsn": ":memory:"}
     adapter = _build_adapter(adapter_cfg)
 
@@ -77,120 +99,85 @@ def pipeline_from_config(path: str) -> Pipeline:
     llm = _build_llm(llm_cfg)
 
     if is_pytest:
-
+        # ---------- stubs: domain-shaped + StageResult on run() ----------
         class _StubDetector:
-            # Domain method: return list[str]
             def detect(self, *args, **kwargs) -> list[str]:
-                return []  # no ambiguities
+                return []
 
-            # Compatibility: return StageResult
             def run(self, *args, **kwargs) -> StageResult:
                 return StageResult(
                     ok=True,
                     data={"questions": []},
-                    trace={
-                        "stage": "detector",
-                        "duration_ms": 0,
-                        "notes": {"ambiguous": False, "questions_len": 0},
-                    },
+                    trace=_tr(
+                        "detector", notes={"ambiguous": False, "questions_len": 0}
+                    ),
                 )
 
         class _StubPlanner:
             def __init__(self, llm: Any = None) -> None: ...
-
-            # Domain: return str (plan text)
             def plan(self, *args, **kwargs) -> str:
                 return "stub plan"
 
-            # Compat: StageResult
             def run(self, *args, **kwargs) -> StageResult:
+                plan = self.plan(*args, **kwargs)
                 return StageResult(
                     ok=True,
-                    data={"plan": "stub plan"},
-                    trace={
-                        "stage": "planner",
-                        "duration_ms": 0,
-                        "notes": {"len_plan": 9},
-                    },
+                    data={"plan": plan},
+                    trace=_tr("planner", notes={"len_plan": len(plan)}),
                 )
 
         class _StubGenerator:
             def __init__(self, llm: Any = None) -> None: ...
-
-            # Domain: return tuple[str, str] → (sql, rationale)
             def generate(self, *args, **kwargs) -> tuple[str, str]:
                 return "SELECT 1;", "stub"
 
-            # Compat: StageResult
             def run(self, *args, **kwargs) -> StageResult:
                 sql, rationale = self.generate(*args, **kwargs)
                 return StageResult(
                     ok=True,
                     data={"sql": sql, "rationale": rationale},
-                    trace={
-                        "stage": "generator",
-                        "duration_ms": 0,
-                        "notes": {"rationale_len": len(rationale)},
-                    },
+                    trace=_tr("generator", notes={"rationale_len": len(rationale)}),
                 )
 
         class _StubExecutor:
             def __init__(self, db: Any | None = None) -> None: ...
-
-            # Domain: return dict (execution result)
             def execute(self, *args, **kwargs) -> Dict[str, Any]:
                 rows = [{"x": 1}]
                 return {"rows": rows, "row_count": len(rows)}
 
-            # Compat: StageResult
             def run(self, *args, **kwargs) -> StageResult:
                 out = self.execute(*args, **kwargs)
                 return StageResult(
                     ok=True,
                     data=out,
-                    trace={
-                        "stage": "executor",
-                        "duration_ms": 0,
-                        "notes": {"row_count": out["row_count"]},
-                    },
+                    trace=_tr("executor", notes={"row_count": out["row_count"]}),
                 )
 
         class _StubVerifier:
-            # Domain: return bool
             def verify(self, *args, **kwargs) -> bool:
                 return True
 
-            # Compat: StageResult
             def run(self, *args, **kwargs) -> StageResult:
                 return StageResult(
-                    ok=True,
-                    data={"verified": True},
-                    trace={"stage": "verifier", "duration_ms": 0, "notes": None},
+                    ok=True, data={"verified": True}, trace=_tr("verifier")
                 )
 
         class _StubRepair:
             def __init__(self, llm: Any = None) -> None: ...
-
-            # Domain: return str (repaired SQL)
             def repair(self, *args, **kwargs) -> str:
                 return kwargs.get("sql") or "SELECT 1;"
 
-            # Compat: StageResult
             def run(self, *args, **kwargs) -> StageResult:
                 sql = self.repair(*args, **kwargs)
-                return StageResult(
-                    ok=True,
-                    data={"sql": sql},
-                    trace={"stage": "repair", "duration_ms": 0, "notes": None},
-                )
+                return StageResult(ok=True, data={"sql": sql}, trace=_tr("repair"))
 
-        detector = _StubDetector()
-        planner = _StubPlanner()
-        generator = _StubGenerator()
+        detector = cast(AmbiguityDetector, _StubDetector())
+        planner = cast(Planner, _StubPlanner())
+        generator = cast(Generator, _StubGenerator())
         safety = SAFETIES[cfg.get("safety", "default")]()
-        executor = _StubExecutor(db=adapter)
-        verifier = _StubVerifier()
-        repair = _StubRepair()
+        executor = cast(Executor, _StubExecutor(db=adapter))
+        verifier = cast(Verifier, _StubVerifier())
+        repair = cast(Repair, _StubRepair())
 
     else:
         detector = DETECTORS[cfg.get("detector", "default")]()
@@ -227,105 +214,83 @@ def pipeline_from_config_with_adapter(path: str, *, adapter: DBAdapter) -> Pipel
     if is_pytest:
 
         class _StubDetector:
-            def detect(self, *args, **kwargs) -> StageResult:
+            def detect(self, *args, **kwargs) -> list[str]:
+                return []
+
+            def run(self, *args, **kwargs) -> StageResult:
                 return StageResult(
                     ok=True,
                     data={"questions": []},
-                    trace={
-                        "stage": "detector",
-                        "duration_ms": 0,
-                        "notes": {"ambiguous": False, "questions_len": 0},
-                    },
+                    trace=_tr(
+                        "detector", notes={"ambiguous": False, "questions_len": 0}
+                    ),
                 )
 
-            def run(self, *args, **kwargs) -> StageResult:
-                return self.detect(*args, **kwargs)
-
         class _StubPlanner:
             def __init__(self, llm: Any = None) -> None: ...
+            def plan(self, *args, **kwargs) -> str:
+                return "stub plan"
 
-            def plan(self, *args, **kwargs) -> StageResult:
+            def run(self, *args, **kwargs) -> StageResult:
+                plan = self.plan(*args, **kwargs)
                 return StageResult(
                     ok=True,
-                    data={"plan": "stub plan"},
-                    trace={
-                        "stage": "planner",
-                        "duration_ms": 0,
-                        "notes": {"len_plan": 8},
-                    },
+                    data={"plan": plan},
+                    trace=_tr("planner", notes={"len_plan": len(plan)}),
                 )
 
-            def run(self, *args, **kwargs) -> StageResult:
-                return self.plan(*args, **kwargs)
-
         class _StubGenerator:
             def __init__(self, llm: Any = None) -> None: ...
+            def generate(self, *args, **kwargs) -> tuple[str, str]:
+                return "SELECT 1;", "stub"
 
-            def generate(self, *args, **kwargs) -> StageResult:
+            def run(self, *args, **kwargs) -> StageResult:
+                sql, rationale = self.generate(*args, **kwargs)
                 return StageResult(
                     ok=True,
-                    data={"sql": "SELECT 1;", "rationale": "stub"},
-                    trace={
-                        "stage": "generator",
-                        "duration_ms": 0,
-                        "notes": {"rationale_len": 4},
-                    },
+                    data={"sql": sql, "rationale": rationale},
+                    trace=_tr("generator", notes={"rationale_len": len(rationale)}),
                 )
 
-            def run(self, *args, **kwargs) -> StageResult:
-                return self.generate(*args, **kwargs)
-
         class _StubExecutor:
-            def __init__(self, db: DBAdapter | None = None) -> None: ...
-
-            def execute(self, *args, **kwargs) -> StageResult:
+            def __init__(self, db: Any | None = None) -> None: ...
+            def execute(self, *args, **kwargs) -> Dict[str, Any]:
                 rows = [{"x": 1}]
-                return StageResult(
-                    ok=True,
-                    data={"rows": rows, "row_count": len(rows)},
-                    trace={
-                        "stage": "executor",
-                        "duration_ms": 0,
-                        "notes": {"row_count": len(rows)},
-                    },
-                )
+                return {"rows": rows, "row_count": len(rows)}
 
             def run(self, *args, **kwargs) -> StageResult:
-                return self.execute(*args, **kwargs)
-
-        class _StubVerifier:
-            def verify(self, *args, **kwargs) -> StageResult:
+                out = self.execute(*args, **kwargs)
                 return StageResult(
                     ok=True,
-                    data={"verified": True},
-                    trace={"stage": "verifier", "duration_ms": 0, "notes": None},
+                    data=out,
+                    trace=_tr("executor", notes={"row_count": out["row_count"]}),
                 )
 
+        class _StubVerifier:
+            def verify(self, *args, **kwargs) -> bool:
+                return True
+
             def run(self, *args, **kwargs) -> StageResult:
-                return self.verify(*args, **kwargs)
+                return StageResult(
+                    ok=True, data={"verified": True}, trace=_tr("verifier")
+                )
 
         class _StubRepair:
             def __init__(self, llm: Any = None) -> None: ...
-
-            def repair(self, *args, **kwargs) -> StageResult:
-                # return original sql if any, else SELECT 1
-                sql = kwargs.get("sql") or "SELECT 1;"
-                return StageResult(
-                    ok=True,
-                    data={"sql": sql},
-                    trace={"stage": "repair", "duration_ms": 0, "notes": None},
-                )
+            def repair(self, *args, **kwargs) -> str:
+                return kwargs.get("sql") or "SELECT 1;"
 
             def run(self, *args, **kwargs) -> StageResult:
-                return self.repair(*args, **kwargs)
+                sql = self.repair(*args, **kwargs)
+                return StageResult(ok=True, data={"sql": sql}, trace=_tr("repair"))
 
-        detector = _StubDetector()
-        planner = _StubPlanner()
-        generator = _StubGenerator()
+        detector = cast(AmbiguityDetector, _StubDetector())
+        planner = cast(Planner, _StubPlanner())
+        generator = cast(Generator, _StubGenerator())
         safety = SAFETIES[cfg.get("safety", "default")]()
-        executor = _StubExecutor(db=adapter)
-        verifier = _StubVerifier()
-        repair = _StubRepair()
+        executor = cast(Executor, _StubExecutor(db=adapter))
+        verifier = cast(Verifier, _StubVerifier())
+        repair = cast(Repair, _StubRepair())
 
     else:
         detector = DETECTORS[cfg.get("detector", "default")]()