nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/nodes/generator/node.py

@@ -0,0 +1,267 @@
+from __future__ import annotations
+
+import sqlglot
+from sqlglot import expressions as exp
+from typing import Dict, Any, List, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import SubgraphExecutionState
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.common.logger import get_logger
+from nl2sql.pipeline.nodes.ast_planner.schemas import PlanModel, Expr
+from nl2sql.pipeline.nodes.generator.schemas import GeneratorResponse
+from nl2sql.context import NL2SQLContext
+
+logger = get_logger("generator")
+
+
+class SqlVisitor:
+    """Visits the PlanModel AST and converts it to sqlglot expressions.
+
+    This visitor traverses the deterministic AST (Expr) produced by the Planner
+    and builds a corresponding sqlglot expression tree, which can then be
+    transpiled to the target dialect.
+    """
+
+    def visit(self, expr: Expr) -> exp.Expression:
+        """Dispatches the visit to the appropriate method based on expression kind.
+
+        Args:
+            expr (Expr): The expression node to visit.
+
+        Returns:
+            exp.Expression: The corresponding sqlglot expression.
+
+        Raises:
+            ValueError: If the expression kind is unknown.
+        """
+        if expr.kind == "literal":
+            return self._visit_literal(expr)
+        elif expr.kind == "column":
+            return self._visit_column(expr)
+        elif expr.kind == "func":
+            return self._visit_func(expr)
+        elif expr.kind == "binary":
+            return self._visit_binary(expr)
+        elif expr.kind == "unary":
+            return self._visit_unary(expr)
+        elif expr.kind == "case":
+            return self._visit_case(expr)
+        raise ValueError(f"Unknown expression kind: {expr.kind}")
+
+    def _visit_literal(self, expr: Expr) -> exp.Expression:
+        """Converts a literal expression to sqlglot."""
+        val = expr.value
+        if val is None:
+            return exp.Null()
+        if isinstance(val, bool):
+            return exp.Boolean(this="TRUE" if val else "FALSE")
+        if isinstance(val, (int, float)):
+            return exp.Literal.number(str(val))
+        return exp.Literal.string(str(val))
+
+    def _visit_column(self, expr: Expr) -> exp.Column:
+        """Converts a column expression to sqlglot."""
+        ident = exp.Identifier(this=expr.column_name, quoted=False)
+        if expr.alias:
+            return exp.Column(this=ident, table=exp.Identifier(this=expr.alias, quoted=False))
+        return exp.Column(this=ident)
+
+    def _visit_func(self, expr: Expr) -> exp.Expression:
+        """Converts a function call expression to sqlglot."""
+        if str(expr.func_name).upper() in ("TUPLE", "LIST"):
+            return exp.Tuple(expressions=[self.visit(arg) for arg in expr.args])
+
+        return exp.Anonymous(
+            this=expr.func_name,
+            expressions=[self.visit(arg) for arg in expr.args]
+        )
+
+    def _visit_binary(self, expr: Expr) -> exp.Expression:
+        """Converts a binary operation expression to sqlglot."""
+        if not expr.left or not expr.right:
+            raise ValueError("Binary expression missing operands")
+
+        left = self.visit(expr.left)
+        right = self.visit(expr.right)
+        op = str(expr.op).upper()
+
+        if op == "=": return exp.EQ(this=left, expression=right)
+        if op == "!=": return exp.NEQ(this=left, expression=right)
+        if op == ">": return exp.GT(this=left, expression=right)
+        if op == "<": return exp.LT(this=left, expression=right)
+        if op == ">=": return exp.GTE(this=left, expression=right)
+        if op == "<=": return exp.LTE(this=left, expression=right)
+        if op == "AND":
+            if isinstance(left, exp.Or): left = exp.Paren(this=left)
+            if isinstance(right, exp.Or): right = exp.Paren(this=right)
+            return exp.And(this=left, expression=right)
+        if op == "OR": return exp.Or(this=left, expression=right)
+        if op == "LIKE": return exp.Like(this=left, expression=right)
+        if op == "IN":
+            values = right.expressions if isinstance(right, exp.Tuple) else [right]
+            return exp.In(this=left, expressions=values)
+
+        return exp.Anonymous(this=op, expressions=[left, right])
+
+    def _visit_unary(self, expr: Expr) -> exp.Expression:
+        """Converts a unary operation expression to sqlglot."""
+        target = expr.expr
+        if not target:
+            raise ValueError("Unary expression missing target")
+
+        node = self.visit(target)
+        op = str(expr.op).upper()
+
+        if op == "NOT":
+            return exp.Not(this=node)
+        if op == "-":
+            return exp.Neg(this=node)
+
+        return exp.Paren(this=node)
+
+    def _visit_case(self, expr: Expr) -> exp.Case:
+        """Converts a CASE expression to sqlglot."""
+        when_list = []
+
+        if expr.whens:
+            for w in expr.whens:
+                when_list.append(
+                    exp.When(
+                        this=self.visit(w.condition),
+                        then=self.visit(w.result)
+                    )
+                )
+
+        default = self.visit(expr.else_expr) if expr.else_expr else None
+        return exp.Case(ifs=when_list, default=default)
+
+
+class GeneratorNode:
+    """Generates the final SQL string from the PlanModel using sqlglot.
+
+    Attributes:
+        registry (DatasourceRegistry): The registry to fetch datasource dialects.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        """Initializes the GeneratorNode.
+
+        Args:
+            registry (DatasourceRegistry): The registry of datasources.
+        """
+        self.node_name = self.__class__.__name__.lower().replace('node', '')
+        self.ds_registry = ctx.ds_registry
+
+    def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
+        """Executes the generator node.
+
+        Converts the PlanModel into a SQL string tailored for the target
+        datasource's dialect.
+
+        Args:
+            state (SubgraphExecutionState): The current state containing the execution plan.
+
+        Returns:
+            Dict[str, Any]: A dictionary with the generated 'sql_draft' and reasoning.
+        """
+        try:
+            datasource_id = state.sub_query.datasource_id if state.sub_query else None
+            if not datasource_id:
+                raise ValueError("No datasource selected")
+            plan = state.ast_planner_response.plan if state.ast_planner_response else None
+            if not plan:
+                raise ValueError("No plan provided")
+
+            adapter = self.ds_registry.get_adapter(datasource_id)
+            dialect = adapter.get_dialect()
+
+            row_limit = adapter.row_limit or 1000
+            limit = min(int(plan.limit or row_limit), row_limit)
+
+            sql = self._generate_sql(plan, limit, dialect)
+
+            response = GeneratorResponse(
+                sql_draft=sql,
+                reasoning=[{"node": self.node_name, "content": ["Generated SQL", sql]}],
+            )
+            return {
+                "generator_response": response,
+                "reasoning": response.reasoning,
+            }
+
+        except Exception as exc:
+            logger.exception(exc)
+            error = PipelineError(
+                node=self.node_name,
+                message=str(exc),
+                error_code=ErrorCode.SQL_GEN_FAILED,
+                severity=ErrorSeverity.ERROR,
+                stack_trace=str(exc),
+            )
+            return {
+                "generator_response": GeneratorResponse(errors=[error]),
+                "errors": [error],
+            }
+
+    def _generate_sql(self, plan: PlanModel, limit: int, dialect: str) -> str:
+        """Internal helper to build and optimize the SQL query."""
+        visitor = SqlVisitor()
+        query = exp.select()
+
+        for s in sorted(plan.select_items, key=lambda x: x.ordinal):
+            e = visitor.visit(s.expr)
+            if s.alias:
+                e = exp.Alias(this=e, alias=exp.Identifier(this=s.alias, quoted=False))
+            query = query.select(e)
+
+        tables = sorted(plan.tables, key=lambda x: x.ordinal)
+        if not tables:
+            raise ValueError("Plan has no tables")
+
+        primary = tables[0]
+        tbl = exp.Table(this=exp.Identifier(this=primary.name, quoted=False))
+
+        if primary.schema_name:
+            tbl.set("db", exp.Identifier(this=primary.schema_name, quoted=False))
+        if primary.database:
+            tbl.set("catalog", exp.Identifier(this=primary.database, quoted=False))
+        if primary.alias:
+            tbl.set("alias", exp.TableAlias(this=exp.Identifier(this=primary.alias, quoted=False)))
+
+        query = query.from_(tbl)
+
+        alias_map = {t.alias: t.name for t in tables}
+
+        for j in sorted(plan.joins, key=lambda x: x.ordinal):
+            name = alias_map.get(j.right_alias)
+            if not name:
+                raise ValueError(f"Join references unknown alias {j.right_alias}")
+
+            right = exp.Table(this=exp.Identifier(this=name, quoted=False))
+            right.set("alias", exp.TableAlias(this=exp.Identifier(this=j.right_alias, quoted=False)))
+
+            condition = visitor.visit(j.condition)
+
+            query = query.join(
+                right,
+                on=condition,
+                join_type=j.join_type
+            )
+
+        if plan.where:
+            query = query.where(visitor.visit(plan.where))
+
+        for g in sorted(plan.group_by, key=lambda x: x.ordinal):
+            query = query.group_by(visitor.visit(g.expr))
+
+        if plan.having:
+            query = query.having(visitor.visit(plan.having))
+
+        for o in sorted(plan.order_by, key=lambda x: x.ordinal):
+            query = query.order_by(visitor.visit(o.expr), desc=(o.direction == "desc"))
+
+        query = query.limit(limit)
+
+        return query.sql(dialect=dialect)

nl2sql/pipeline/nodes/generator/schemas.py

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

nl2sql/pipeline/nodes/global_planner/__init__.py

@@ -0,0 +1,4 @@
+from .node import GlobalPlannerNode
+from .schemas import GlobalPlannerResponse
+
+__all__ = ["GlobalPlannerNode", "GlobalPlannerResponse"]

nl2sql/pipeline/nodes/global_planner/node.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+from typing import Dict, Any, TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import GraphState
+
+from .schemas import (
+    RelationSchema,
+    ColumnSpec,
+    ExecutionDAG,
+    LogicalNode,
+    LogicalEdge,
+)
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from .schemas import GlobalPlannerResponse
+
+logger = get_logger("global_planner")
+
+
+class GlobalPlannerNode:
+    """Generates a deterministic execution plan for aggregating sub-query results.
+    
+    This node runs AFTER the Decomposer and BEFORE the SQL Agents.
+    It does not execute SQL, but produces a blueprint (ResultPlan) for the Aggregator.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+
+    def __call__(self, state: GraphState) -> Dict[str, Any]:
+        """Executes the planning logic."""
+        decomposer_response = state.decomposer_response
+        sub_queries = decomposer_response.sub_queries if decomposer_response else []
+        combine_groups = decomposer_response.combine_groups if decomposer_response else []
+        post_combine_ops = decomposer_response.post_combine_ops if decomposer_response else []
+
+        try:
+            nodes: List[LogicalNode] = []
+            edges: List[LogicalEdge] = []
+            node_index: Dict[str, LogicalNode] = {}
+
+            for sq in sub_queries:
+                schema = RelationSchema(
+                    columns=[ColumnSpec(name=c.name, dtype=c.dtype) for c in sq.expected_schema]
+                )
+                node = LogicalNode(
+                    node_id=sq.id,
+                    kind="scan",
+                    inputs=[],
+                    output_schema=schema,
+                    attributes={
+                        "datasource_id": sq.datasource_id,
+                        "intent": sq.intent,
+                        "metrics": [m.model_dump() for m in sq.metrics],
+                        "filters": [f.model_dump() for f in sq.filters],
+                        "group_by": [g.model_dump() for g in sq.group_by],
+                        "expected_schema": [c.model_dump() for c in sq.expected_schema],
+                    },
+                )
+                nodes.append(node)
+                node_index[node.node_id] = node
+
+            combine_node_ids: Dict[str, str] = {}
+            for cg in combine_groups:
+                combine_id = f"combine_{cg.group_id}"
+                combine_node_ids[cg.group_id] = combine_id
+                input_ids = [i.subquery_id for i in cg.inputs]
+                output_schema = (
+                    node_index[input_ids[0]].output_schema
+                    if input_ids and input_ids[0] in node_index
+                    else RelationSchema(columns=[])
+                )
+                node = LogicalNode(
+                    node_id=combine_id,
+                    kind="combine",
+                    inputs=input_ids,
+                    output_schema=output_schema,
+                    attributes={
+                        "operation": cg.operation,
+                        "group_id": cg.group_id,
+                        "inputs": [i.model_dump() for i in cg.inputs],
+                        "join_keys": [jk.model_dump() for jk in cg.join_keys],
+                    },
+                )
+                nodes.append(node)
+                node_index[node.node_id] = node
+
+                for inp in cg.inputs:
+                    edges.append(
+                        LogicalEdge(
+                            edge_id=f"edge_{inp.subquery_id}_{combine_id}",
+                            from_id=inp.subquery_id,
+                            to_id=combine_id,
+                            role=inp.role,
+                        )
+                    )
+
+            for op in post_combine_ops:
+                target_combine_id = combine_node_ids.get(op.target_group_id)
+                if not target_combine_id:
+                    raise ValueError(f"PostCombineOp references unknown combine group: {op.target_group_id}")
+                output_schema = RelationSchema(
+                    columns=[ColumnSpec(name=c.name, dtype=c.dtype) for c in op.expected_schema]
+                )
+                kind_map = {
+                    "filter": "post_filter",
+                    "aggregate": "post_aggregate",
+                    "project": "post_project",
+                    "sort": "post_sort",
+                    "limit": "post_limit",
+                }
+                node = LogicalNode(
+                    node_id=op.op_id,
+                    kind=kind_map[op.operation],
+                    inputs=[target_combine_id],
+                    output_schema=output_schema,
+                    attributes={
+                        "target_group_id": op.target_group_id,
+                        "operation": op.operation,
+                        "filters": [f.model_dump() for f in op.filters],
+                        "metrics": [m.model_dump() for m in op.metrics],
+                        "group_by": [g.model_dump() for g in op.group_by],
+                        "order_by": [o.model_dump() for o in op.order_by],
+                        "limit": op.limit,
+                        "expected_schema": [c.model_dump() for c in op.expected_schema],
+                        "metadata": op.metadata,
+                    },
+                )
+                nodes.append(node)
+                node_index[node.node_id] = node
+                edges.append(
+                    LogicalEdge(
+                        edge_id=f"edge_{target_combine_id}_{op.op_id}",
+                        from_id=target_combine_id,
+                        to_id=op.op_id,
+                    )
+                )
+
+            node_ids = {n.node_id for n in nodes}
+            for edge in edges:
+                if edge.from_id not in node_ids or edge.to_id not in node_ids:
+                    raise ValueError(f"Edge references unknown node: {edge}")
+
+            execution_dag = ExecutionDAG(
+                nodes=sorted(nodes, key=lambda n: n.node_id),
+                edges=sorted(edges, key=lambda e: (e.from_id, e.to_id, e.role or "")),
+            )
+
+            execution_dag.content_hash = self._hash_execution_dag(execution_dag)
+            execution_dag.dag_id = f"dag_{execution_dag.content_hash[:12]}"
+
+            result = GlobalPlannerResponse(execution_dag=execution_dag)
+
+            return {
+                "global_planner_response": result,      
+                "reasoning": [{"node": self.node_name, "content": "Built explicit execution DAG."}],
+            }
+
+        except Exception as e:
+            logger.error(f"GlobalPlanner failed: {e}")
+            return {
+                "global_planner_response": GlobalPlannerResponse(execution_dag=None),
+                "reasoning": [{"node": self.node_name, "content": f"Planning failed: {e}", "type": "error"}],
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message=f"Global Plan generation failed: {str(e)}",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.PLANNER_FAILED 
+                    )
+                ]
+            }
+
+    def _hash_execution_dag(self, dag: ExecutionDAG) -> str:
+        import hashlib
+        import json
+
+        payload = {
+            "nodes": [n.model_dump() for n in dag.nodes],
+            "edges": [e.model_dump() for e in dag.edges],
+            "version": dag.version,
+        }
+        data = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
+        return hashlib.sha256(data.encode("utf-8")).hexdigest()

nl2sql/pipeline/nodes/global_planner/schemas.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+from typing import List, Literal, Optional, Dict, Any
+from pydantic import BaseModel, Field, model_validator
+
+
+JsonLiteral = str | int | float | bool | None
+
+
+class ColumnSpec(BaseModel):
+    name: str
+    dtype: Optional[str] = None
+
+class RelationSchema(BaseModel):
+    columns: List[ColumnSpec]
+
+    @model_validator(mode="after")
+    def validate_unique_columns(self):
+        names = [c.name for c in self.columns]
+        if len(names) != len(set(names)):
+            raise ValueError(f"Duplicate columns in schema: {names}")
+        return self
+
+class LogicalNode(BaseModel):
+    node_id: str
+    kind: Literal[
+        "scan",
+        "combine",
+        "post_filter",
+        "post_aggregate",
+        "post_project",
+        "post_sort",
+        "post_limit",
+    ]
+    inputs: List[str] = Field(default_factory=list)
+    output_schema: RelationSchema
+    attributes: Dict[str, Any] = Field(default_factory=dict)
+
+
+class LogicalEdge(BaseModel):
+    edge_id: str
+    from_id: str
+    to_id: str
+    role: Optional[str] = None
+
+
+class ExecutionDAG(BaseModel):
+    dag_id: Optional[str] = None
+    content_hash: Optional[str] = None
+    nodes: List[LogicalNode]
+    edges: List[LogicalEdge]
+    layers: List[List[str]] = Field(default_factory=list)
+    version: str = "v1"
+
+    @model_validator(mode="after")
+    def populate_layers(self):
+        if not self.layers and self.nodes:
+            self.layers = self._layered_toposort(self.nodes, self.edges)
+        return self
+
+    @staticmethod
+    def _layered_toposort(
+        nodes: List[LogicalNode],
+        edges: List[LogicalEdge],
+    ) -> List[List[str]]:
+        node_ids = {n.node_id for n in nodes}
+        indegree: Dict[str, int] = {n.node_id: 0 for n in nodes}
+        dependents: Dict[str, List[str]] = {n.node_id: [] for n in nodes}
+
+        for edge in edges:
+            if edge.from_id not in node_ids or edge.to_id not in node_ids:
+                continue
+            indegree[edge.to_id] += 1
+            dependents[edge.from_id].append(edge.to_id)
+
+        layers: List[List[str]] = []
+        ready = sorted([n_id for n_id, deg in indegree.items() if deg == 0])
+        processed = 0
+
+        while ready:
+            current_layer = ready
+            layers.append(current_layer)
+            processed += len(current_layer)
+
+            next_ready: List[str] = []
+            for node_id in current_layer:
+                for child in sorted(dependents.get(node_id, [])):
+                    indegree[child] -= 1
+                    if indegree[child] == 0:
+                        next_ready.append(child)
+            ready = sorted(set(next_ready))
+
+        if processed != len(nodes):
+            raise ValueError("ExecutionDAG contains a cycle; layered topo sort failed.")
+
+        return layers
+
+
+class GlobalPlannerResponse(BaseModel):
+    execution_dag: ExecutionDAG
+
+

nl2sql/pipeline/nodes/refiner/__init__.py

@@ -0,0 +1,4 @@
+from .node import RefinerNode
+from .schemas import RefinerResponse
+
+__all__ = ["RefinerNode", "RefinerResponse"]