refactor: unify pipeline output via FinalResult model

app/routers/nl2sql.py

@@ -1,13 +1,12 @@
 from dataclasses import asdict, is_dataclass
 from fastapi import APIRouter, HTTPException
 from app.schemas import NL2SQLRequest, NL2SQLResponse, ClarifyResponse
-from nl2sql.pipeline import Pipeline
+from nl2sql.pipeline import Pipeline, FinalResult
 from nl2sql.ambiguity_detector import AmbiguityDetector
 from nl2sql.safety import Safety
 from nl2sql.planner import Planner
 from nl2sql.generator import Generator
 from adapters.llm.openai_provider import OpenAIProvider
-from nl2sql.types import StageResult
 from nl2sql.executor import Executor
 from nl2sql.verifier import Verifier
 from nl2sql.repair import Repair
@@ -59,28 +58,28 @@ def _round_trace(t: dict) -> dict:
 @router.post("", name="nl2sql_handler")
 def nl2sql_handler(request: NL2SQLRequest):
     result = _pipeline.run(
-        user_query=request.query, schema_preview=request.schema_preview
+        user_query=request.query,
+        schema_preview=request.schema_preview,
     )
 
     # --- Ensure result type ---
-    if not isinstance(result, StageResult):
+    if not isinstance(result, FinalResult):
         raise HTTPException(status_code=500, detail="Pipeline returned unexpected type")
 
-    data = result.data or {}
-
     # --- Handle ambiguity ---
-    if isinstance(data, dict) and data.get("ambiguous") and data.get("questions"):
-        return ClarifyResponse(ambiguous=True, questions=data["questions"])
+    if result.ambiguous and result.questions:
+        return ClarifyResponse(ambiguous=True, questions=result.questions)
 
     # --- Handle error ---
-    if not result.ok:
-        detail = "; ".join(result.error) if result.error else "Unknown error"
+    if not result.ok or result.error:
+        detail = "; ".join(result.details or ["Unknown error"])
         raise HTTPException(status_code=400, detail=detail)
 
     # --- Success case ---
+    traces = [ _round_trace(t) for t in (result.traces or []) ]
     return NL2SQLResponse(
         ambiguous=False,
-        sql=data.get("sql"),
-        rationale=data.get("rationale"),
-        traces=[_to_dict(t) for t in data.get("traces", [])],
+        sql=result.sql,
+        rationale=result.rationale,
+        traces=traces,
     )

nl2sql/pipeline.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 import traceback
+from dataclasses import dataclass, asdict
 from typing import Dict, Any, Optional, List
+
 from nl2sql.types import StageResult
 from nl2sql.ambiguity_detector import AmbiguityDetector
 from nl2sql.planner import Planner
@@ -11,10 +13,26 @@ from nl2sql.verifier import Verifier
 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
+    ambiguous: bool
+    error: bool
+    details: Optional[List[str]]
+    sql: Optional[str]
+    rationale: Optional[str]
+    verified: Optional[bool]
+    questions: Optional[List[str]]
+    traces: List[dict]
+
+
 class Pipeline:
     """
-    NL2SQL Copilot pipeline with guaranteed dict output.
-    All stages return structured traces and errors but final result is JSON-safe dict.
+    NL2SQL Copilot pipeline.
+    Stages return StageResult; final result is a type-safe FinalResult.
+    Adapters (e.g. FastAPI) can serialize with dataclasses.asdict().
     """
 
     def __init__(
@@ -25,7 +43,7 @@ class Pipeline:
         generator: Generator,
         safety: Safety,
         executor: Optional[Executor] = None,
-        verifier: Optional[Verifier] = None ,
+        verifier: Optional[Verifier] = None,
         repair: Optional[Repair] = None,
     ):
         self.detector = detector
@@ -55,7 +73,7 @@ class Pipeline:
             if isinstance(r, StageResult):
                 return r
             else:
-                # not ideal, but wrap it
+                # Normalize non-StageResult returns
                 return StageResult(ok=True, data=r, trace=None)
         except Exception as e:
             tb = traceback.format_exc()
@@ -68,41 +86,40 @@ class Pipeline:
         user_query: str,
         schema_preview: str,
         clarify_answers: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
-        """
-        Always returns:
-        {
-            "ambiguous": bool,
-            "error": bool,
-            "details": list[str] | None,
-            "sql": str | None,
-            "rationale": str | None,
-            "verified": bool | None,
-            "traces": list[dict]
-        }
-        """
+    ) -> FinalResult:
         traces: List[dict] = []
         details: List[str] = []
-        sql, rationale, verified = None, None, None
+        sql: Optional[str] = None
+        rationale: Optional[str] = None
+        verified: Optional[bool] = None
 
         # --- 1) ambiguity detection
         try:
             questions = self.detector.detect(user_query, schema_preview)
             if questions:
-                return {
-                    "ambiguous": True,
-                    "error": False,
-                    "details": [f"Ambiguities found: {len(questions)}"],
-                    "questions": questions,
-                    "traces": [],
-                }
+                return FinalResult(
+                    ok=True,
+                    ambiguous=True,
+                    error=False,
+                    details=[f"Ambiguities found: {len(questions)}"],
+                    questions=questions,
+                    sql=None,
+                    rationale=None,
+                    verified=None,
+                    traces=[],
+                )
         except Exception as e:
-            return {
-                "ambiguous": True,
-                "error": True,
-                "details": [f"Detector failed: {e}"],
-                "traces": [],
-            }
+            return FinalResult(
+                ok=False,
+                ambiguous=True,
+                error=True,
+                details=[f"Detector failed: {e}"],
+                questions=None,
+                sql=None,
+                rationale=None,
+                verified=None,
+                traces=[],
+            )
 
         # --- 2) planner
         r_plan = self._safe_stage(
@@ -110,12 +127,17 @@ class Pipeline:
         )
         traces.extend(self._trace_list(r_plan))
         if not r_plan.ok:
-            return {
-                "ambiguous": False,
-                "error": True,
-                "details": r_plan.error,
-                "traces": traces,
-            }
+            return FinalResult(
+                ok=False,
+                ambiguous=False,
+                error=True,
+                details=r_plan.error,
+                questions=None,
+                sql=None,
+                rationale=None,
+                verified=None,
+                traces=traces,
+            )
 
         # --- 3) generator
         r_gen = self._safe_stage(
@@ -127,40 +149,51 @@ class Pipeline:
         )
         traces.extend(self._trace_list(r_gen))
         if not r_gen.ok:
-            return {
-                "ambiguous": False,
-                "error": True,
-                "details": r_gen.errors,
-                "traces": traces,
-            }
+            return FinalResult(
+                ok=False,
+                ambiguous=False,
+                error=True,
+                details=r_gen.error,
+                questions=None,
+                sql=None,
+                rationale=None,
+                verified=None,
+                traces=traces,
+            )
         sql = r_gen.data.get("sql")
         rationale = r_gen.data.get("rationale")
 
         # --- 4) safety
-        r_safe = self._safe_stage(self.safety.check, sql=sql)
+        # 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:
-            return {
-                "ambiguous": False,
-                "error": True,
-                "details": r_safe.error,
-                "traces": traces,
-            }
+            return FinalResult(
+                ok=False,
+                ambiguous=False,
+                error=True,
+                details=r_safe.error,
+                questions=None,
+                sql=sql,
+                rationale=rationale,
+                verified=None,
+                traces=traces,
+            )
 
         # --- 5) executor
-        r_exec = self._safe_stage(self.executor.run, sql=r_safe.data["sql"])
+        r_exec = self._safe_stage(self.executor.run, sql=r_safe.data.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)
+        r_ver = self._safe_stage(self.verifier.run, sql=sql, exec_result=r_exec.data)
         traces.extend(self._trace_list(r_ver))
         verified = bool(r_ver.ok)
 
         # --- 7) repair loop if verification failed
         if not verified:
-            for attempt in range(2):
+            for _attempt in range(2):
                 r_fix = self._safe_stage(
                     self.repair.run,
                     sql=sql,
@@ -171,29 +204,33 @@ class Pipeline:
                 if not r_fix.ok:
                     break
                 sql = r_fix.data.get("sql")
-                r_safe = self._safe_stage(self.safety.check, 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 [])
                     continue
-                r_exec = self._safe_stage(self.executor.run, sql=r_safe.data["sql"])
+
+                r_exec = self._safe_stage(self.executor.run, sql=r_safe.data.get("sql", sql))
                 traces.extend(self._trace_list(r_exec))
                 if not r_exec.ok:
                     details.extend(r_exec.error or [])
                     continue
-                r_ver = self._safe_stage(self.verifier.run, sql=sql, exec_result=r_exec)
+
+                r_ver = self._safe_stage(self.verifier.run, sql=sql, exec_result=r_exec.data)
                 traces.extend(self._trace_list(r_ver))
                 verified = bool(r_ver.ok)
                 if verified:
                     break
 
-        # --- Final result dict
-        return {
-            "ambiguous": False,
-            "error": len(details) > 0 and not verified,
-            "details": details or None,
-            "sql": sql,
-            "rationale": rationale,
-            "verified": verified,
-            "traces": traces,
-        }
+        return FinalResult(
+            ok=bool(verified) and not details,
+            ambiguous=False,
+            error=bool(details) and not bool(verified),
+            details=details or None,
+            sql=sql,
+            rationale=rationale,
+            verified=verified,
+            questions=None,
+            traces=traces,
+        )

nl2sql/types.py

@@ -19,3 +19,20 @@ class StageResult:
     trace: Optional[StageTrace] = None
     error: Optional[List[str]] = None
     notes: Optional[Dict[str, Any]] = None
+
+
+@dataclass(frozen=True)
+class FinalResult:
+    """
+    Final domain result of the whole pipeline.
+    Adapters (HTTP/CLI/UI) should serialize this to dict/JSON at the boundary.
+    """
+    ok: bool                   # end-to-end success
+    ambiguous: bool
+    error: bool
+    sql: Optional[str]
+    rationale: Optional[str]
+    verified: Optional[bool]
+    details: Optional[List[str]]
+    questions: Optional[List[str]]
+    traces: List[Dict[str, Any]]

tests/test_nl2sql_router.py

@@ -1,13 +1,12 @@
 from fastapi.testclient import TestClient
 from app.main import app
-from nl2sql.types import StageResult, StageTrace
+from nl2sql.pipeline import FinalResult
 
 client = TestClient(app)
 
-
-def fake_trace(stage: str):
-    return StageTrace(stage=stage, duration_ms=10.0)
-
+def fake_trace(stage: str) -> dict:
+    # FinalResult.traces is a list of dicts (StageTrace.__dict__)
+    return {"stage": stage, "duration_ms": 10.0, "cost_usd": None, "notes": None}
 
 path = app.url_path_for("nl2sql_handler")
 
@@ -16,15 +15,18 @@ path = app.url_path_for("nl2sql_handler")
 def test_ambiguity_route(monkeypatch):
     from app.routers import nl2sql
 
-    # mock pipeline to return StageResult with ambiguous=True
+    # mock pipeline to return FinalResult with ambiguous=True
     def fake_run(*args, **kwargs):
-        return StageResult(
+        return FinalResult(
             ok=True,
-            data={
-                "ambiguous": True,
-                "questions": ["Which table do you mean?"],
-                "traces": [fake_trace("detector")],
-            },
+            ambiguous=True,
+            error=False,
+            details=["Ambiguities found: 1"],
+            questions=["Which table do you mean?"],
+            sql=None,
+            rationale=None,
+            verified=None,
+            traces=[fake_trace("detector")],
         )
 
     monkeypatch.setattr(nl2sql._pipeline, "run", fake_run)
@@ -36,11 +38,11 @@ def test_ambiguity_route(monkeypatch):
             "schema_preview": "CREATE TABLE ...",
         },
     )
-
     assert resp.status_code == 200
     data = resp.json()
     assert data["ambiguous"] is True
     assert "questions" in data
+    assert isinstance(data["questions"], list)
 
 
 # --- 2) Error / failure case -------------------------------------------------
@@ -48,8 +50,16 @@ 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 FinalResult(
+            ok=False,
+            ambiguous=False,
+            error=True,
+            details=["Bad SQL"],
+            questions=None,
+            sql=None,
+            rationale=None,
+            verified=None,
+            traces=[fake_trace("safety")],
         )
 
     monkeypatch.setattr(nl2sql._pipeline, "run", fake_run)
@@ -61,7 +71,6 @@ def test_error_route(monkeypatch):
             "schema_preview": "CREATE TABLE users(id int);",
         },
     )
-
     assert resp.status_code == 400
     assert "Bad SQL" in resp.json()["detail"]
 
@@ -71,14 +80,16 @@ def test_success_route(monkeypatch):
     from app.routers import nl2sql
 
     def fake_run(*args, **kwargs):
-        return StageResult(
+        return FinalResult(
             ok=True,
-            data={
-                "ambiguous": False,
-                "sql": "SELECT * FROM users;",
-                "rationale": "Simple listing",
-                "traces": [fake_trace("planner"), fake_trace("generator")],
-            },
+            ambiguous=False,
+            error=False,
+            details=None,
+            questions=None,
+            sql="SELECT * FROM users;",
+            rationale="Simple listing",
+            verified=True,
+            traces=[fake_trace("planner"), fake_trace("generator")],
         )
 
     monkeypatch.setattr(nl2sql._pipeline, "run", fake_run)
@@ -96,3 +107,4 @@ def test_success_route(monkeypatch):
     assert data["sql"].lower().startswith("select")
     assert isinstance(data["traces"], list)
     assert any(t["stage"] == "planner" for t in data["traces"])
+    assert any(t["stage"] == "generator" for t in data["traces"])

tests/test_pipeline_integration.py

@@ -1,14 +1,13 @@
-from nl2sql.pipeline import Pipeline
+from nl2sql.pipeline import Pipeline, FinalResult
 from nl2sql.types import StageResult, StageTrace
 
 
 # --- Dummy stages to isolate pipeline -----------------------------------------
 
-
 class DummyDetector:
     """Simulates ambiguity detector stage."""
 
-    def __init__(self, ambiguous=False):
+    def __init__(self, ambiguous: bool = False):
         self.ambiguous = ambiguous
 
     def detect(self, user_query, schema_preview):
@@ -43,10 +42,12 @@ class DummyGenerator:
 class DummySafety:
     """Simulates safety stage."""
 
-    def check(self, sql):
+    # NOTE: pipeline now calls safety.run(sql=...)
+    def run(self, *, sql):
         trace = StageTrace(stage="safety", duration_ms=1.0)
         if "DROP" in sql.upper():
             return StageResult(ok=False, error=["Unsafe SQL"], trace=trace)
+        # echo back sql in data to feed executor
         return StageResult(ok=True, data={"sql": sql, "rationale": "safe"}, trace=trace)
 
 
@@ -64,13 +65,13 @@ def test_pipeline_success():
         schema_preview="CREATE TABLE singer(id int, name text);",
     )
 
-    assert isinstance(r, StageResult)
+    assert isinstance(r, FinalResult)
     assert r.ok is True
-    data = r.data or {}
-    assert data["sql"].lower().startswith("select")
-    assert any(t.stage == "planner" for t in data["traces"])
-    assert any(t.stage == "generator" for t in data["traces"])
-    assert any(t.stage == "safety" for t in data["traces"])
+    assert r.sql is not None and r.sql.lower().startswith("select")
+    # traces is a list of dicts (StageTrace.__dict__)
+    assert any(t.get("stage") == "planner" for t in r.traces)
+    assert any(t.get("stage") == "generator" for t in r.traces)
+    assert any(t.get("stage") == "safety" for t in r.traces)
 
 
 # --- 2) Ambiguity case --------------------------------------------------------
@@ -84,10 +85,10 @@ def test_pipeline_ambiguity():
 
     r = pipeline.run(user_query="show data", schema_preview="CREATE TABLE x(id int);")
 
-    assert isinstance(r, StageResult)
+    assert isinstance(r, FinalResult)
     assert r.ok is True
-    assert r.data["ambiguous"] is True
-    assert isinstance(r.data["questions"], list)
+    assert r.ambiguous is True
+    assert isinstance(r.questions, list) and len(r.questions) > 0
 
 
 # --- 3) Planner failure -------------------------------------------------------
@@ -101,9 +102,10 @@ def test_pipeline_plan_fail():
     r = pipeline.run(
         user_query="fail_plan", schema_preview="CREATE TABLE singer(id int);"
     )
-    assert isinstance(r, StageResult)
+    assert isinstance(r, FinalResult)
     assert r.ok is False
-    assert "Planner failed" in " ".join(r.error or [])
+    assert r.details is not None
+    assert "Planner failed" in " ".join(r.details)
 
 
 # --- 4) Generator failure -----------------------------------------------------
@@ -117,8 +119,10 @@ def test_pipeline_gen_fail():
     r = pipeline.run(
         user_query="fail_gen", schema_preview="CREATE TABLE singer(id int);"
     )
+    assert isinstance(r, FinalResult)
     assert r.ok is False
-    assert "Generator failed" in " ".join(r.error or [])
+    assert r.details is not None
+    assert "Generator failed" in " ".join(r.details)
 
 
 # --- 5) Safety failure --------------------------------------------------------
@@ -140,5 +144,7 @@ def test_pipeline_safety_fail():
     r = pipeline.run(
         user_query="drop something", schema_preview="CREATE TABLE x(id int);"
     )
+    assert isinstance(r, FinalResult)
     assert r.ok is False
-    assert "unsafe" in " ".join(r.error or []).lower()
+    assert r.details is not None
+    assert "unsafe" in " ".join(r.details).lower()