nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/nodes/validator/schemas.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+from pydantic import BaseModel, Field
+
+from nl2sql.common.errors import PipelineError
+
+
+class LogicalValidatorResponse(BaseModel):
+    errors: List[PipelineError] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)

nl2sql/pipeline/pipeline_runner.py

@@ -0,0 +1,72 @@
+import time
+import traceback
+from typing import Dict, Any, Optional, List
+from dataclasses import dataclass, field
+
+from nl2sql.auth import UserContext
+from nl2sql.pipeline.runtime import run_with_graph
+from nl2sql.common.settings import settings
+from nl2sql.context import NL2SQLContext
+
+
+@dataclass
+class PipelineResult:
+    """Standardized result object from the pipeline execution."""
+
+    success: bool
+    final_state: Dict[str, Any] = field(default_factory=dict)
+    duration: float = 0.0
+    error: Optional[str] = None
+    traceback: Optional[str] = None
+
+
+class PipelineRunner:
+    """
+    Orchestrates the execution of the NL2SQL pipeline.
+    Decoupled from CLI presentation logic.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.ctx = ctx
+
+    def run(
+        self,
+        query: str,
+        role: str = "admin",
+        datasource_id: Optional[str] = None,
+        execute: bool = True,
+        callbacks: List[Any] = None,
+    ) -> PipelineResult:
+        """
+        Executes the pipeline graph.
+        """
+        start_time = time.perf_counter()
+
+        user_context = UserContext(
+            roles=[role],
+            tenant_id=settings.tenant_id,
+        )
+
+        try:
+            final_state = run_with_graph(
+                self.ctx,
+                user_query=query,
+                datasource_id=datasource_id,
+                execute=execute,
+                callbacks=callbacks or [],
+                user_context=user_context,
+            )
+
+            return PipelineResult(
+                success=True,
+                final_state=final_state,
+                duration=time.perf_counter() - start_time,
+            )
+
+        except Exception as e:
+            return PipelineResult(
+                success=False,
+                error=str(e),
+                traceback=traceback.format_exc(),
+                duration=time.perf_counter() - start_time,
+            )

nl2sql/pipeline/routes.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from langgraph.graph import END
+from langgraph.types import Send
+
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.context import NL2SQLContext
+from nl2sql.pipeline.graph_utils import (
+    StateAccessor,
+    build_scan_payload,
+    next_scan_layer_ids,
+    resolve_subgraph,
+)
+from nl2sql.pipeline.state import GraphState
+from nl2sql.common.logger import get_logger
+
+logger = get_logger("router")
+
+
+def resolver_route(state: GraphState) -> str:
+    accessor = StateAccessor(state)
+    resolver_response = accessor.get("datasource_resolver_response")
+    if not resolver_response:
+        return "end"
+    if not resolver_response.resolved_datasources or not resolver_response.allowed_datasource_ids:
+        return "end"
+    return "continue"
+
+
+def build_scan_layer_router(ctx: NL2SQLContext):
+    def route_scan_layers(state: GraphState):
+        global_planner_response = state.global_planner_response
+        dag = global_planner_response.execution_dag if global_planner_response else None
+        decomposer_response = state.decomposer_response
+        sub_queries = decomposer_response.sub_queries if decomposer_response else []
+        sub_query_map = {sq.id: sq for sq in sub_queries}
+        artifact_refs = state.artifact_refs or {}
+        if not dag or not dag.layers:
+            return END
+
+        node_index = {n.node_id: n for n in dag.nodes}
+        target_ids = next_scan_layer_ids(dag, artifact_refs)
+        if not target_ids:
+            return [
+                Send("aggregator",state)
+            ]
+
+        branches = []
+        for node_id in target_ids:
+            if node_id in sub_query_map:
+                sq = sub_query_map[node_id]
+                datasource_id = sq.datasource_id
+            else:
+                node = node_index.get(node_id) if dag else None
+                if not node:
+                    continue
+                datasource_id = (node.attributes or {}).get("datasource_id")
+
+            target = resolve_subgraph(datasource_id, ctx)
+            if not target:
+                raise PipelineError(
+                    node="layer_router",
+                    message=f"No compatible subgraph found for datasource '{datasource_id}'.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.INVALID_STATE,
+                )
+            payload = build_scan_payload(state, target, node_id)
+            branches.append(Send(target, payload))
+
+        return branches
+
+    return route_scan_layers

nl2sql/pipeline/runtime.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import concurrent.futures
+import signal
+import sys
+import threading
+import traceback
+from typing import Callable, Dict, List, Optional
+
+from nl2sql.auth import UserContext
+from nl2sql.common.cancellation import CancellationToken
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.settings import settings
+from nl2sql.context import NL2SQLContext
+from nl2sql.pipeline.graph import build_graph
+from nl2sql.pipeline.state import GraphState
+
+
+def _start_keyboard_cancel_listener(
+    token: CancellationToken,
+    done: threading.Event,
+) -> None:
+    """Cancel ``token`` when Ctrl+X is pressed, until ``done`` is set.
+
+    One listener per run: the thread exits once the run finishes, so a later run
+    starts its own listener bound to its own token.
+    """
+    if sys.platform != "win32":
+        return
+
+    if not sys.stdin or not sys.stdin.isatty():
+        return
+
+    try:
+        import msvcrt
+    except Exception:
+        return
+
+    def _listen():
+        while not done.is_set():
+            if msvcrt.kbhit():
+                char = msvcrt.getch()
+                if char == b"\x18":  # Ctrl+X
+                    token.cancel()
+                    return
+            done.wait(0.05)
+
+    threading.Thread(target=_listen, daemon=True).start()
+
+
+def _install_signal_handlers(token: CancellationToken) -> Callable[[], None]:
+    previous = {}
+
+    def _handler(signum, frame):
+        token.cancel()
+
+    for sig in (getattr(signal, "SIGINT", None), getattr(signal, "SIGTERM", None)):
+        if sig is None:
+            continue
+        previous[sig] = signal.getsignal(sig)
+        signal.signal(sig, _handler)
+
+    def _restore():
+        for sig, handler in previous.items():
+            signal.signal(sig, handler)
+
+    return _restore
+
+
+def run_with_graph(
+    ctx: NL2SQLContext,
+    user_query: str,
+    datasource_id: Optional[str] = None,
+    execute: bool = True,
+    callbacks: Optional[List] = None,
+    user_context: UserContext = None,
+) -> Dict:
+    """Convenience function to run the full pipeline."""
+    token = CancellationToken()
+    run_done = threading.Event()
+    restore_signals = _install_signal_handlers(token)
+    _start_keyboard_cancel_listener(token, run_done)
+
+    graph = build_graph(
+        ctx,
+        execute=execute,
+    )
+
+    initial_state = GraphState(
+        user_query=user_query,
+        user_context=user_context,
+        datasource_id=datasource_id,
+    )
+
+    timeout_sec = settings.global_timeout_sec
+
+    def _invoke():
+        return graph.invoke(
+            initial_state.model_dump(),
+            config={
+                "configurable": {"cancellation_token": token},
+                "callbacks": callbacks,
+            },
+        )
+
+    try:
+        # Use configured thread pool size for pipeline execution
+        with concurrent.futures.ThreadPoolExecutor(max_workers=settings.sandbox_exec_workers) as executor:
+            future = executor.submit(_invoke)
+            result = future.result(timeout=timeout_sec)
+
+        # Nodes observe the token and unwind, so a cancelled run returns normally.
+        if token.is_cancelled():
+            return {
+                "errors": [
+                    PipelineError(
+                        node="orchestrator",
+                        message="Pipeline cancelled by user.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.CANCELLED,
+                    )
+                ]
+            }
+        return result
+    except concurrent.futures.TimeoutError:
+        error_msg = f"Pipeline execution timed out after {timeout_sec} seconds."
+        return {
+            "errors": [
+                PipelineError(
+                    node="orchestrator",
+                    message=error_msg,
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.PIPELINE_TIMEOUT,
+                )
+            ],
+            "final_answer": "I apologize, but the request timed out. Please try again with a simpler query.",
+        }
+    except Exception as e:
+        # Fallback for other runtime crashes
+        return {
+            "errors": [
+                PipelineError(
+                    node="orchestrator",
+                    message=f"Pipeline crashed: {str(e)}",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.UNKNOWN_ERROR,
+                    stack_trace=traceback.format_exc(),
+                )
+            ]
+        }
+    finally:
+        run_done.set()
+        restore_signals()

nl2sql/pipeline/state.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Annotated
+import operator
+import uuid
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from nl2sql.common.errors import PipelineError
+from nl2sql.auth import UserContext
+from nl2sql.pipeline.nodes.datasource_resolver.schemas import DatasourceResolverResponse
+from nl2sql.pipeline.nodes.decomposer.schemas import DecomposerResponse, SubQuery
+from nl2sql.pipeline.nodes.global_planner.schemas import GlobalPlannerResponse
+from nl2sql.pipeline.nodes.aggregator.schemas import AggregatorResponse
+from nl2sql.pipeline.nodes.answer_synthesizer.schemas import AnswerSynthesizerResponse
+from nl2sql.pipeline.nodes.ast_planner.schemas import ASTPlannerResponse
+from nl2sql.pipeline.nodes.validator.schemas import LogicalValidatorResponse
+from nl2sql.pipeline.nodes.generator.schemas import GeneratorResponse
+from nl2sql.execution.contracts import ArtifactRef, ExecutorResponse
+from nl2sql.pipeline.nodes.refiner.schemas import RefinerResponse
+from nl2sql.pipeline.subgraphs.schemas import SubgraphOutput
+from nl2sql.pipeline.nodes.schema_retriever.schema import Table
+
+
+def update_results(current: Dict, new: Dict) -> Dict:
+    """Reducer to merge execution results from parallel branches."""
+    if current is None:
+        return new
+    return {**current, **new}
+
+
+class GraphState(BaseModel):
+    """Represents the shared state of the NL2SQL pipeline execution graph.
+
+    Attributes:
+        trace_id (str): Distributed unique trace ID.
+        user_query (str): Canonical user query.
+        user_context (UserContext): User identity and permissions context.
+        datasource_id (Optional[str]): Optional datasource override for resolution.
+        datasource_resolver_response (Optional[DatasourceResolverResponse]): Output of resolver node.
+        decomposer_response (Optional[DecomposerResponse]): Output of decomposer node.
+        global_planner_response (Optional[GlobalPlannerResponse]): Output of planner node.
+        aggregator_response (Optional[AggregatorResponse]): Output of aggregator node.
+        answer_synthesizer_response (Optional[AnswerSynthesizerResponse]): Output of synthesizer node.
+        artifact_refs (Dict[str, ArtifactRef]): Artifact refs keyed by ExecutionDAG node_id.
+        subgraph_outputs (Dict[str, SubgraphOutput]): Per-subgraph diagnostic outputs.
+        errors (List[PipelineError]): List of errors encountered during execution.
+        reasoning (List[Dict[str, Any]]): Log of reasoning steps from nodes.
+        warnings (List[Dict[str, Any]]): Warning messages emitted by nodes.
+        subgraph_id (Optional[str]): ID of the subgraph execution.
+        subgraph_name (Optional[str]): Name of the subgraph execution.
+    """
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)
+
+    trace_id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Distributed unique trace ID.")
+    user_query: str = Field(description="Canonical user query.")
+    user_context: UserContext = Field(default_factory=UserContext, description="User identity and permissions context.")
+    datasource_id: Optional[str] = Field(default=None, description="Optional datasource override for resolution.")
+    datasource_resolver_response: Optional[DatasourceResolverResponse] = Field(default=None)
+    decomposer_response: Optional[DecomposerResponse] = Field(default=None)
+    global_planner_response: Optional[GlobalPlannerResponse] = Field(default=None)
+    aggregator_response: Optional[AggregatorResponse] = Field(default=None)
+    answer_synthesizer_response: Optional[AnswerSynthesizerResponse] = Field(default=None)
+    artifact_refs: Annotated[Dict[str, ArtifactRef], update_results] = Field(default_factory=dict)
+    subgraph_outputs: Annotated[Dict[str, SubgraphOutput], update_results] = Field(default_factory=dict)
+    errors: Annotated[List[PipelineError], operator.add] = Field(default_factory=list)
+    reasoning: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)
+    warnings: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)
+    subgraph_id: Optional[str] = Field(default=None)
+    subgraph_name: Optional[str] = Field(default=None)
+
+
+class SubgraphExecutionState(BaseModel):
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)
+
+    trace_id: str
+    sub_query: Optional[SubQuery] = None
+    user_context: Optional[UserContext] = None
+    subgraph_id: Optional[str] = None
+    subgraph_name: Optional[str] = None
+    relevant_tables: List[Table] = Field(default_factory=list)
+
+    ast_planner_response: Optional[ASTPlannerResponse] = Field(default=None)
+    logical_validator_response: Optional[LogicalValidatorResponse] = Field(default=None)
+    generator_response: Optional[GeneratorResponse] = Field(default=None)
+    executor_response: Optional[ExecutorResponse] = Field(default=None)
+    refiner_response: Optional[RefinerResponse] = Field(default=None)
+
+    retry_count: int = 0
+    errors: Annotated[List[PipelineError], operator.add] = Field(default_factory=list)
+    reasoning: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)
+    warnings: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)

nl2sql/pipeline/subgraphs/__init__.py

@@ -0,0 +1,5 @@
+"""Subgraph implementations."""
+
+from .schemas import SubgraphOutput
+
+__all__ = ["SubgraphOutput"]

nl2sql/pipeline/subgraphs/schemas.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from nl2sql.common.errors import PipelineError
+from nl2sql.pipeline.nodes.ast_planner.schemas import PlanModel
+from nl2sql.execution.contracts import ArtifactRef
+from nl2sql.pipeline.nodes.decomposer.schemas import SubQuery
+
+
+class SubgraphOutput(BaseModel):
+    sub_query: Optional[SubQuery] = None
+    subgraph_id: str
+    subgraph_name: Optional[str] = None
+    retry_count: int = 0
+    plan: Optional[PlanModel] = None
+    sql_draft: Optional[str] = None
+    artifact: Optional[ArtifactRef] = None
+    errors: List[PipelineError] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)
+    status: Optional[str] = None

nl2sql/pipeline/subgraphs/sql_agent.py

@@ -0,0 +1,167 @@
+from typing import Dict, Optional
+from langchain_core.runnables import Runnable, RunnableConfig
+from langgraph.graph import END, StateGraph
+
+from nl2sql.common.cancellation import CancellationToken
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.settings import settings
+from nl2sql.pipeline.state import SubgraphExecutionState
+from nl2sql.pipeline.nodes.ast_planner import ASTPlannerNode
+from nl2sql.pipeline.nodes.schema_retriever import SchemaRetrieverNode
+from nl2sql.pipeline.nodes.validator import LogicalValidatorNode
+from nl2sql.pipeline.nodes.refiner import RefinerNode
+from nl2sql.pipeline.nodes.generator import GeneratorNode
+from nl2sql.pipeline.nodes.executor import ExecutorNode
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.context import NL2SQLContext
+
+
+
+import random
+import time
+
+
+def build_sql_agent_graph(
+    ctx: NL2SQLContext,
+):
+    """Builds the SQL Agent Subgraph.
+
+    Pipeline Flow:
+    Planner -> LogicalValidator -> Generator -> Executor
+
+    Feedback Loop:
+    (LogicalValidator Error) -> RetryHandler -> Refiner -> Planner
+    """
+    graph = StateGraph(SubgraphExecutionState)
+
+    schema_retriever = SchemaRetrieverNode(ctx)
+    ast_planner = ASTPlannerNode(ctx)
+    logical_validator = LogicalValidatorNode(ctx)
+    refiner = RefinerNode(ctx)
+    generator = GeneratorNode(ctx)
+    executor = ExecutorNode(ctx)
+
+    def _get_subgraph_id(state: SubgraphExecutionState) -> str:
+        if state.subgraph_id:
+            return state.subgraph_id
+        sub_query_id = state.sub_query.id if state.sub_query else None
+        return f"sql_agent:{sub_query_id}:{state.trace_id}"
+
+    def _get_retry_count(state: SubgraphExecutionState) -> int:
+        return state.retry_count
+
+    def _token(config: Optional[RunnableConfig]) -> Optional[CancellationToken]:
+        """Read the current run's cancellation token from the RunnableConfig.
+
+        Returns None when invoked without a config (e.g. a direct unit-test call),
+        which is treated as "not cancelled".
+        """
+        if not config:
+            return None
+        return (config.get("configurable") or {}).get("cancellation_token")
+
+    def _is_cancelled(config: Optional[RunnableConfig]) -> bool:
+        token = _token(config)
+        return bool(token and token.is_cancelled())
+
+    def retry_node(state: SubgraphExecutionState, config: Optional[RunnableConfig] = None) -> Dict:
+        """Increments retry count with exponential backoff and jitter."""
+        count = _get_retry_count(state)
+        if _is_cancelled(config):
+            return {
+                "errors": [
+                    PipelineError(
+                        node="retry_handler",
+                        message="Pipeline cancelled by user.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.CANCELLED,
+                    )
+                ]
+            }
+        if count >= settings.sql_agent_max_retries:
+            return {"retry_count": count}
+
+        base_delay = min(settings.sql_agent_retry_max_delay_sec, settings.sql_agent_retry_base_delay_sec * (2 ** count))
+        jitter = random.uniform(0.0, settings.sql_agent_retry_jitter_sec)
+        sleep_time = base_delay + jitter
+
+        token = _token(config)
+        if token is None:
+            time.sleep(sleep_time)
+        elif token.wait(timeout=sleep_time):
+            return {
+                "errors": [
+                    PipelineError(
+                        node="retry_handler",
+                        message="Pipeline cancelled by user.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.CANCELLED,
+                    )
+                ]
+            }
+
+        return {
+            "retry_count": count + 1,
+        }
+
+    def check_planner(state: SubgraphExecutionState, config: Optional[RunnableConfig] = None) -> str:
+        """Routes based on planner result."""
+        if _is_cancelled(config):
+            return "end"
+        if not (state.ast_planner_response and state.ast_planner_response.plan):
+            # If explicit errors exist, check retryability
+            if state.errors:
+                 if not all(e.is_retryable for e in state.errors):
+                     return "end"
+
+            if _get_retry_count(state) < settings.sql_agent_max_retries:
+                return "retry"
+            return "end"
+        return "ok"
+
+    def check_logical_validation(state: SubgraphExecutionState, config: Optional[RunnableConfig] = None) -> str:
+        """Routes based on logical validation result."""
+        if _is_cancelled(config):
+            return "end"
+        if state.logical_validator_response and state.logical_validator_response.errors:
+            # Critical/Fatal errors stop execution immediately
+            if not all(e.is_retryable for e in state.logical_validator_response.errors):
+                return "end"
+
+            if _get_retry_count(state) < settings.sql_agent_max_retries:
+                return "retry"
+            return "end"
+        return "ok"
+
+    graph.add_node("schema_retriever", schema_retriever)
+    graph.add_node("ast_planner", ast_planner)
+    graph.add_node("logical_validator", logical_validator)
+    graph.add_node("generator", generator)
+    graph.add_node("executor", executor)
+    graph.add_node("refiner", refiner)
+    graph.add_node("retry_handler", retry_node)
+
+    graph.set_entry_point("schema_retriever")
+
+    graph.add_edge("schema_retriever", "ast_planner")
+
+    graph.add_conditional_edges(
+        "ast_planner",
+        check_planner,
+        {"ok": "logical_validator", "retry": "retry_handler", "end": END},
+    )
+
+    graph.add_conditional_edges(
+        "logical_validator",
+        check_logical_validation,
+        {"ok": "generator", "retry": "retry_handler", "end": END},
+    )
+
+    graph.add_edge("generator", "executor")
+
+    graph.add_edge("executor", END)
+
+    graph.add_edge("retry_handler", "refiner")
+    graph.add_edge("refiner", "ast_planner")
+
+    return graph.compile()