Sync from GitHub main

adapters/llm/openai_provider.py

@@ -43,6 +43,10 @@ class OpenAIProvider(LLMProvider):
         """Return metadata of the last LLM call (tokens, cost, sql_length, kind)."""
         return dict(self._last_usage)
 
+    def _create_chat_completion(self, **kwargs):
+        """OpenAI SDK seam for stable unit testing."""
+        return self.client.chat.completions.create(**kwargs)
+
     def __init__(self) -> None:
         """Initialize OpenAI client with config from environment."""
         api_key, base_url, model = _resolve_api_config()
@@ -84,7 +88,7 @@ Database Schema:
 
 Create a step-by-step plan to answer this question with SQL."""
 
-        completion = self.client.chat.completions.create(
+        completion = self._create_chat_completion(
             model=self.model,
             messages=[
                 {"role": "system", "content": system_prompt},
@@ -181,7 +185,7 @@ Now generate the SQL for the given question:"""
         if clarify_answers:
             user_prompt += f"\n\nAdditional context: {clarify_answers}"
 
-        completion = self.client.chat.completions.create(
+        completion = self._create_chat_completion(
             model=self.model,
             messages=[
                 {"role": "system", "content": system_prompt},
@@ -316,7 +320,7 @@ Database Schema:
 
 Return the corrected SQL (keep it simple):"""
 
-        completion = self.client.chat.completions.create(
+        completion = self._create_chat_completion(
             model=self.model,
             messages=[
                 {"role": "system", "content": system_prompt},
@@ -419,7 +423,7 @@ Database Schema:
 Please answer these clarification questions:
 {chr(10).join(f"{i + 1}. {q}" for i, q in enumerate(questions))}"""
 
-        completion = self.client.chat.completions.create(
+        completion = self._create_chat_completion(
             model=self.model,
             messages=[
                 {"role": "system", "content": system_prompt},

nl2sql/generator.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
+
 import time
 from typing import Optional, Dict, Any
-from nl2sql.types import StageResult, StageTrace
+
 from adapters.llm.base import LLMProvider
+from nl2sql.errors.codes import ErrorCode
+from nl2sql.types import StageResult, StageTrace
 
 
 class Generator:
@@ -20,6 +23,7 @@ class Generator:
         clarify_answers: Optional[Dict[str, Any]] = None,
     ) -> StageResult:
         t0 = time.perf_counter()
+
         try:
             res = self.llm.generate_sql(
                 user_query=user_query,
@@ -28,15 +32,23 @@ class Generator:
                 clarify_answers=clarify_answers or {},
             )
         except Exception as e:
-            return StageResult(ok=False, error=[f"Generator failed: {e}"])
+            # Provider/transport errors or unexpected runtime issues.
+            return StageResult(
+                ok=False,
+                error=[f"Generator failed: {e}"],
+                error_code=ErrorCode.LLM_BAD_OUTPUT,
+                trace=None,
+            )
 
-        # Expect a 5-tuple
+        # Contract: expect a 5-tuple (sql, rationale, token_in, token_out, cost_usd)
         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)"
                 ],
+                error_code=ErrorCode.LLM_BAD_OUTPUT,
+                trace=None,
             )
 
         sql, rationale, t_in, t_out, cost = res
@@ -44,12 +56,23 @@ class Generator:
         # Type/shape checks
         if not isinstance(sql, str) or not sql.strip():
             return StageResult(
-                ok=False, error=["Generator produced empty or non-string SQL"]
+                ok=False,
+                error=["Generator produced empty or non-string SQL"],
+                error_code=ErrorCode.LLM_BAD_OUTPUT,
+                trace=None,
             )
+
+        # Enforce SELECT-only at the boundary (fast fail before hitting later stages).
         if not sql.lower().lstrip().startswith("select"):
-            return StageResult(ok=False, error=[f"Generated non-SELECT SQL: {sql}"])
+            return StageResult(
+                ok=False,
+                error=[f"Generated non-SELECT SQL: {sql}"],
+                error_code=ErrorCode.SAFETY_NON_SELECT,
+                trace=None,
+            )
 
-        rationale = rationale or ""  # safe length
+        # Normalize rationale to a string
+        rationale = rationale or ""
         trace = StageTrace(
             stage=self.name,
             duration_ms=(time.perf_counter() - t0) * 1000.0,
@@ -60,5 +83,9 @@ class Generator:
         )
 
         return StageResult(
-            ok=True, data={"sql": sql, "rationale": rationale}, trace=trace
+            ok=True,
+            data={"sql": sql, "rationale": rationale},
+            trace=trace,
+            error_code=None,
+            retryable=None,
         )

nl2sql/pipeline.py

@@ -3,6 +3,7 @@ import traceback
 from dataclasses import dataclass
 from typing import Dict, Any, Optional, List
 import time
+import inspect
 
 from nl2sql.types import StageResult
 from nl2sql.ambiguity_detector import AmbiguityDetector
@@ -239,6 +240,25 @@ class Pipeline:
             "schema_preview": kwargs.get("schema_preview", ""),
         }
 
+    def _call_verifier(self, *, sql: str, exec_result: Dict[str, Any]) -> StageResult:
+        """
+        Call verifier with a backward-compatible signature.
+        Some verifiers accept `adapter=...`, some don't.
+        """
+        kwargs: Dict[str, Any] = {"sql": sql, "exec_result": exec_result}
+
+        adapter = getattr(self.executor, "adapter", None)
+        if adapter is not None:
+            try:
+                params = inspect.signature(self.verifier.run).parameters
+                if "adapter" in params:
+                    kwargs["adapter"] = adapter
+            except (TypeError, ValueError):
+                # If signature introspection fails, fall back to the minimal call.
+                pass
+
+        return self.verifier.run(**kwargs)
+
     def run(
         self,
         *,
@@ -315,6 +335,7 @@ class Pipeline:
                     ambiguous=False,
                     error=True,
                     details=r_plan.error,
+                    error_code=ErrorCode.PIPELINE_CRASH,
                     questions=None,
                     sql=None,
                     rationale=None,
@@ -347,6 +368,7 @@ class Pipeline:
                     ambiguous=False,
                     error=True,
                     details=r_gen.error,
+                    error_code=ErrorCode.LLM_BAD_OUTPUT,
                     questions=None,
                     sql=None,
                     rationale=None,
@@ -368,6 +390,7 @@ class Pipeline:
                     ambiguous=False,
                     error=True,
                     details=["empty_sql"],
+                    error_code=ErrorCode.LLM_BAD_OUTPUT,
                     questions=None,
                     sql=None,
                     rationale=rationale,
@@ -397,6 +420,7 @@ class Pipeline:
                     ambiguous=False,
                     error=True,
                     details=r_safe.error,
+                    error_code=r_safe.error_code,
                     questions=None,
                     sql=sql,
                     rationale=rationale,
@@ -431,12 +455,11 @@ class Pipeline:
             t0 = time.perf_counter()
             r_ver = self._run_with_repair(
                 "verifier",
-                self.verifier.run,
+                self._call_verifier,
                 repair_input_builder=self._sql_repair_input_builder,
                 max_attempts=1,
                 sql=sql,
                 exec_result=(r_exec.data or {}),
-                adapter=getattr(self.executor, "adapter", None),
                 traces=traces,
             )
             dt = (time.perf_counter() - t0) * 1000.0
@@ -522,10 +545,9 @@ class Pipeline:
                     # verifier again
                     t0 = time.perf_counter()
                     r_ver2 = self._safe_stage(
-                        self.verifier.run,
+                        self._call_verifier,
                         sql=sql,
                         exec_result=(r_exec2.data or {}),
-                        adapter=getattr(self.executor, "adapter", None),
                     )
                     dt2 = (time.perf_counter() - t0) * 1000.0
                     stage_duration_ms.labels("verifier").observe(dt2)