nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/graph.py

@@ -0,0 +1,73 @@
+from langgraph.graph import END, StateGraph
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.pipeline.graph_utils import SQL_AGENT_SUBGRAPH, wrap_subgraph
+from nl2sql.pipeline.nodes.aggregator import EngineAggregatorNode
+from nl2sql.pipeline.nodes.answer_synthesizer import AnswerSynthesizerNode
+from nl2sql.pipeline.nodes.datasource_resolver import DatasourceResolverNode
+from nl2sql.pipeline.nodes.decomposer import DecomposerNode
+from nl2sql.pipeline.nodes.global_planner import GlobalPlannerNode
+from nl2sql.pipeline.routes import build_scan_layer_router, resolver_route
+from nl2sql.pipeline.state import GraphState
+from nl2sql.pipeline.subgraphs.sql_agent import build_sql_agent_graph
+
+def build_graph(
+    ctx: NL2SQLContext,
+    execute: bool = True,
+) -> StateGraph:
+    """Builds the main LangGraph pipeline.
+
+    Constructs the graph with Semantic Analysis, Decomposer, Execution branches,
+    and Aggregator.
+
+    Args:
+        ctx (NL2SQLContext): The application context containing registries and services.
+        execute (bool): Whether to allow execution against real databases.
+
+    Returns:
+        StateGraph: The compiled LangGraph runnable.
+    """
+    graph = StateGraph(GraphState)
+
+    resolver_node = DatasourceResolverNode(ctx)
+    decomposer_node = DecomposerNode(ctx)
+    aggregator_node = EngineAggregatorNode(ctx)
+    synthesizer_node = AnswerSynthesizerNode(ctx)
+    global_planner_node = GlobalPlannerNode(ctx)
+
+    sql_agent_subgraph = build_sql_agent_graph(ctx)
+
+    graph.add_node("datasource_resolver", resolver_node)
+    graph.add_node("decomposer", decomposer_node)
+    graph.add_node("global_planner", global_planner_node)
+    graph.add_node(
+        SQL_AGENT_SUBGRAPH,
+        wrap_subgraph(sql_agent_subgraph, SQL_AGENT_SUBGRAPH, ctx),
+    )
+    graph.add_node("aggregator", aggregator_node)
+    graph.add_node("answer_synthesizer", synthesizer_node)
+    graph.add_node("layer_router", lambda state: {})
+
+    graph.set_entry_point("datasource_resolver")
+
+    graph.add_conditional_edges(
+        "datasource_resolver",
+        resolver_route,
+        {"continue": "decomposer", "end": END},
+    )
+
+    graph.add_edge("decomposer", "global_planner")
+    route_scan_layers = build_scan_layer_router(ctx)
+
+    graph.add_edge("global_planner", "layer_router")
+    graph.add_conditional_edges(
+        "layer_router",
+        route_scan_layers,
+        [SQL_AGENT_SUBGRAPH, "aggregator", END],
+    )
+
+    graph.add_edge(SQL_AGENT_SUBGRAPH, "layer_router")
+    graph.add_edge("aggregator", "answer_synthesizer")
+    graph.add_edge("answer_synthesizer", END)
+
+    return graph.compile()

nl2sql/pipeline/graph_utils.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional
+
+from langchain_core.runnables import Runnable, RunnableConfig
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.pipeline.nodes.global_planner.schemas import ExecutionDAG
+from nl2sql.pipeline.state import GraphState, SubgraphExecutionState
+from nl2sql.pipeline.subgraphs import SubgraphOutput
+from nl2sql_adapter_sdk.capabilities import DatasourceCapability
+import logging
+
+SQL_AGENT_SUBGRAPH = "sql_agent"
+SQL_AGENT_REQUIRED_CAPABILITIES = {DatasourceCapability.SUPPORTS_SQL.value}
+
+
+class StateAccessor:
+    """Adapter for GraphState/dict access to simplify routing logic."""
+
+    def __init__(self, state: Any):
+        self._state = state
+
+    def get(self, key: str, default: Any = None) -> Any:
+        if isinstance(self._state, dict):
+            return self._state.get(key, default)
+        return getattr(self._state, key, default)
+
+
+def next_scan_layer_ids(
+    dag: ExecutionDAG,
+    artifact_refs: Dict[str, Any],
+) -> List[str]:
+    node_index = {n.node_id: n for n in dag.nodes}
+    for layer in dag.layers or []:
+        pending_scan = [
+            node_id
+            for node_id in layer
+            if node_id in node_index
+            and node_index[node_id].kind == "scan"
+            and node_id not in artifact_refs
+        ]
+        if pending_scan:
+            return pending_scan
+    return []
+
+
+def resolve_subgraph(
+    datasource_id: str,
+    ctx: NL2SQLContext,
+) -> Optional[str]:
+    """Returns the subgraph able to serve a datasource, or None if none can.
+
+    The sql_agent subgraph is the only one, and it requires SUPPORTS_SQL. A
+    datasource without that capability resolves to None so the caller fails
+    instead of routing it into SQL generation.
+    """
+    try:
+        caps = ctx.ds_registry.get_capabilities(datasource_id)
+    except Exception:
+        return None
+
+    if SQL_AGENT_REQUIRED_CAPABILITIES.issubset(caps):
+        return SQL_AGENT_SUBGRAPH
+    return None
+
+
+def build_scan_payload(
+    state: GraphState,
+    subgraph_name: str,
+    node_id: str,
+) -> Dict[str, Any]:
+    trace_id = state.trace_id   
+    return {
+        "subgraph_id": f"{subgraph_name}:{node_id}:{trace_id}",
+        "subgraph_name": subgraph_name,
+        "trace_id": trace_id,
+        "user_context": state.user_context,
+        "decomposer_response": state.decomposer_response,
+        "datasource_resolver_response": state.datasource_resolver_response,
+    }
+
+def wrap_subgraph(
+    subgraph: Runnable,
+    subgraph_name: str,
+    ctx: NL2SQLContext,
+) -> Callable[Dict[str, Any]]:
+    def _wrapper(state_dict: dict, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
+        trace_id = state_dict.get("trace_id")
+        subgraph_id = state_dict.get("subgraph_id")
+        sub_query_id = subgraph_id.split(":")[1]
+        sub_query = None
+        decomposer_response = state_dict.get("decomposer_response")
+        for sq in decomposer_response.sub_queries:
+            if sq.id == sub_query_id:
+                sub_query = sq
+                break
+
+        sub_state = SubgraphExecutionState(
+            trace_id=trace_id,
+            user_context=state_dict.get("user_context"),
+            sub_query=sub_query,
+            subgraph_id=subgraph_id,
+            subgraph_name=subgraph_name,
+        )
+        # Propagate the run config (carrying the cancellation token) into the subgraph.
+        result = subgraph.invoke(sub_state.model_dump(), config=config)
+
+        returned_state = SubgraphExecutionState.model_validate(result)
+
+        executor_response = returned_state.executor_response
+        planner_response = returned_state.ast_planner_response
+        generator_response = returned_state.generator_response
+        sub_reasoning = returned_state.reasoning
+        artifact_refs: Dict[str, Any] = {}
+        artifact = executor_response.artifact
+        artifact_refs[sub_query.id] = artifact
+
+        retry_count = returned_state.retry_count
+        status = "error" if returned_state.errors else "success"
+        subgraph_output = SubgraphOutput(
+            sub_query=sub_query,
+            subgraph_name=subgraph_name,
+            subgraph_id=subgraph_id,
+            retry_count=retry_count,
+            plan=planner_response.plan,
+            sql_draft=generator_response.sql_draft if generator_response else None,
+            artifact=artifact,
+            errors=returned_state.errors,
+            reasoning=sub_reasoning,
+            status=status,
+        )
+
+        return {
+            "artifact_refs": artifact_refs,
+            "subgraph_outputs": {subgraph_id: subgraph_output},
+            "errors": returned_state.errors,
+            "reasoning": returned_state.reasoning,
+        }
+
+    return _wrapper

nl2sql/pipeline/nodes/__init__.py

@@ -0,0 +1,25 @@
+from .decomposer.node import DecomposerNode
+from .datasource_resolver.node import DatasourceResolverNode
+from .global_planner.node import GlobalPlannerNode
+from .aggregator.node import EngineAggregatorNode
+from .ast_planner.node import ASTPlannerNode
+from .schema_retriever.node import SchemaRetrieverNode
+from .generator.node import GeneratorNode
+from .executor.node import ExecutorNode
+from .refiner.node import RefinerNode
+from .validator.node import LogicalValidatorNode
+from .answer_synthesizer.node import AnswerSynthesizerNode
+
+__all__ = [
+    "ASTPlannerNode", 
+    "SchemaRetrieverNode",
+    "GeneratorNode", 
+    "ExecutorNode", 
+    "RefinerNode", 
+    "DecomposerNode", 
+    "DatasourceResolverNode",
+    "GlobalPlannerNode",
+    "LogicalValidatorNode",
+    "EngineAggregatorNode",
+    "AnswerSynthesizerNode",
+]

nl2sql/pipeline/nodes/aggregator/__init__.py

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

nl2sql/pipeline/nodes/aggregator/node.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+from typing import Dict, Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import GraphState
+
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.pipeline.nodes.aggregator.schemas import AggregatorResponse
+from nl2sql.pipeline.nodes.global_planner.schemas import GlobalPlannerResponse
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from nl2sql.aggregation import AggregationService
+from nl2sql.aggregation.engines.polars_duckdb import PolarsDuckdbEngine
+
+logger = get_logger("aggregator")
+
+
+class EngineAggregatorNode:
+    """Thin wrapper for aggregation service using ExecutionDAG layers."""
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+        self.ctx = ctx
+        self.service = AggregationService(PolarsDuckdbEngine())
+
+    def __call__(self, state: GraphState) -> Dict[str, Any]:
+        try:
+            planner_response = state.global_planner_response
+            artifact_refs = state.artifact_refs
+           
+            dag = planner_response.execution_dag
+            
+            terminal_results = self.service.execute(dag, artifact_refs)
+            aggregator_response = AggregatorResponse(
+                terminal_results=terminal_results,
+                computed_artifacts={},
+            )
+            return {
+                "aggregator_response": aggregator_response,
+                "reasoning": [{"node": self.node_name, "content": "ExecutionDAG aggregation executed successfully."}],
+            }
+        except Exception as exc:
+            logger.error(f"Node {self.node_name} failed: {exc}")
+            return {
+                "aggregator_response": AggregatorResponse(),
+                "reasoning": [{"node": self.node_name, "content": f"Error: {str(exc)}", "type": "error"}],
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message=f"Aggregator failed: {str(exc)}",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.AGGREGATOR_FAILED,
+                    )
+                ],
+            }

nl2sql/pipeline/nodes/aggregator/prompts.py

@@ -0,0 +1,20 @@
+"""Prompts for the Aggregator node."""
+
+AGGREGATOR_PROMPT = """You are an expert data analyst. Your task is to synthesize results from multiple database queries into a single, coherent answer for the user.
+
+User Query: {user_query}
+
+Intermediate Results from Sub-Queries:
+{intermediate_results}
+
+Instructions:
+1. Analyze the intermediate results.
+2. Determine the best way to present the combined information to the user (Table, List, or Text).
+   - Use 'table' if comparing data or listing structured records with common fields.
+   - Use 'list' if enumerating items.
+   - Use 'text' if providing a summary or explanation.
+3. Generate a summary of the findings.
+4. Format the content accordingly.
+
+If the results contain error messages, explain them clearly in the summary.
+"""

nl2sql/pipeline/nodes/aggregator/schemas.py

@@ -0,0 +1,28 @@
+from typing import Literal, List, Dict, Any
+from pydantic import BaseModel, Field
+
+
+class AggregatedResponse(BaseModel):
+    """Structured response for the aggregator node.
+
+    Attributes:
+        summary (str): A concise summary of the aggregated results.
+        format_type (Literal): The format to present the data (table, list, text).
+        content (str): The aggregated content formatted according to format_type.
+    """
+    summary: str = Field(description="A concise summary of the aggregated results.")
+    format_type: Literal["table", "list", "text"] = Field(
+        description="The best format to present the data: 'table' for structured data, 'list' for items, 'text' for narrative."
+    )
+    content: str = Field(description="The aggregated content formatted according to format_type (e.g., Markdown table, bullet points, or paragraph).")
+    warnings: List[str] = Field(
+        default_factory=list,
+        description="Optional warnings to communicate missing or skipped results.",
+    )
+
+
+class AggregatorResponse(BaseModel):
+    terminal_results: Dict[str, List[Dict[str, Any]]] = Field(default_factory=dict)
+    computed_artifacts: Dict[str, Any] = Field(default_factory=dict)
+    errors: List[Any] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)

nl2sql/pipeline/nodes/answer_synthesizer/__init__.py

@@ -0,0 +1,4 @@
+from .node import AnswerSynthesizerNode
+from .schemas import AggregatedResponse, AnswerSynthesizerResponse
+
+__all__ = ["AnswerSynthesizerNode", "AggregatedResponse", "AnswerSynthesizerResponse"]

nl2sql/pipeline/nodes/answer_synthesizer/node.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import json
+from typing import Dict, Any, TYPE_CHECKING
+
+from langchain_core.prompts import ChatPromptTemplate
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import GraphState
+
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from .schemas import AggregatedResponse, AnswerSynthesizerResponse
+from .prompts import ANSWER_SYNTHESIZER_PROMPT
+
+logger = get_logger("answer_synthesizer")
+
+
+class AnswerSynthesizerNode:
+    """Summarizes aggregated results into a user-facing answer."""
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+        self.llm = ctx.llm_registry.get_llm(self.node_name)
+        self.prompt = ChatPromptTemplate.from_template(ANSWER_SYNTHESIZER_PROMPT)
+        self.chain = self.prompt | self.llm.with_structured_output(
+            AggregatedResponse, method="function_calling"
+        )
+
+    def _serialize_result(self, result: Any) -> str:
+        try:
+            return json.dumps(result, indent=2, ensure_ascii=True)
+        except TypeError:
+            return str(result)
+
+    def __call__(self, state: GraphState) -> Dict[str, Any]:
+        aggregated_result = None
+        if state.aggregator_response:
+            aggregated_result = state.aggregator_response.terminal_results
+        elif state.answer_synthesizer_response and state.answer_synthesizer_response.final_answer is not None:
+            aggregated_result = state.answer_synthesizer_response.final_answer
+
+        if aggregated_result is None:
+            return {
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message="No aggregated result available for synthesis.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.INVALID_STATE,
+                    )
+                ]
+            }
+
+        try:
+            unmapped_subqueries = []
+            if state.decomposer_response:
+                unmapped_subqueries = [
+                    u.model_dump()
+                    for u in (state.decomposer_response.unmapped_subqueries or [])
+                ]
+            response: AggregatedResponse = self.chain.invoke(
+                {
+                    "user_query": state.user_query,
+                    "aggregated_result": self._serialize_result(aggregated_result),
+                    "unmapped_subqueries": json.dumps(
+                        unmapped_subqueries, indent=2, ensure_ascii=True
+                    ),
+                }
+            )
+
+
+            return {
+                "answer_synthesizer_response": AnswerSynthesizerResponse(
+                    final_answer=response.model_dump(),
+                ),
+                "reasoning": [
+                    {
+                        "node": self.node_name,
+                        "content": response.summary,
+                    }
+                ],
+            }
+
+        except Exception as exc:
+            logger.error(f"Node {self.node_name} failed: {exc}")
+            return {
+                "answer_synthesizer_response": AnswerSynthesizerResponse(),
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message=f"Answer synthesis failed: {exc}",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.AGGREGATOR_FAILED,
+                    )
+                ]
+            }

nl2sql/pipeline/nodes/answer_synthesizer/prompts.py

@@ -0,0 +1,19 @@
+"""Prompts for the AnswerSynthesizer node."""
+
+ANSWER_SYNTHESIZER_PROMPT = """You are a data analyst. Summarize the aggregated results for the user.
+
+User Query: {user_query}
+
+Aggregated Results (keyed by terminal node id):
+{aggregated_result}
+
+Unmapped Subqueries (if any):
+{unmapped_subqueries}
+
+Instructions:
+1. Provide a concise summary.
+2. Choose the best output format: table, list, or text.
+3. Produce the formatted content for the chosen format.
+4. If results contain error messages, explain them clearly.
+5. If there are unmapped subqueries, add user-facing warnings that explain what was skipped and why.
+"""

nl2sql/pipeline/nodes/answer_synthesizer/schemas.py

@@ -0,0 +1,24 @@
+from typing import Literal, List, Dict, Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class AggregatedResponse(BaseModel):
+    """Structured response for the answer synthesizer."""
+    summary: str = Field(description="A concise summary of the aggregated results.")
+    format_type: Literal["table", "list", "text"] = Field(
+        description="The best format to present the data: 'table' for structured data, 'list' for items, 'text' for narrative."
+    )
+    content: str = Field(
+        description="The aggregated content formatted according to format_type (e.g., Markdown table, bullet points, or paragraph)."
+    )
+    warnings: List[str] = Field(
+        default_factory=list,
+        description="Optional warnings to communicate missing or skipped results.",
+    )
+
+
+class AnswerSynthesizerResponse(BaseModel):
+    final_answer: Optional[Dict[str, Any]] = None
+    errors: List[Any] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)

nl2sql/pipeline/nodes/ast_planner/__init__.py

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

nl2sql/pipeline/nodes/ast_planner/node.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+import traceback
+from typing import Any, Dict, Optional, TYPE_CHECKING
+from langchain_core.runnables import Runnable
+from langchain_core.prompts import ChatPromptTemplate
+
+from .prompts import PLANNER_PROMPT, PLANNER_EXAMPLES
+from .schemas import PlanModel, ASTPlannerResponse
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import SubgraphExecutionState
+
+logger = get_logger("planner")
+
+
+class ASTPlannerNode:
+    """Generates a structured SQL execution plan (PlanModel).
+
+    Uses an LLM to interpret the user query and semantic context, producing a
+    deterministic Abstract Syntax Tree (AST) that represents the SQL query.
+
+    Attributes:
+        llm (Optional[Runnable]): The Language Model executable.
+        chain (Optional[Runnable]): The langchain chain for planning.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        """Initializes the PlannerNode.
+
+        Args:
+            ctx (NL2SQLContext): The context of the pipeline.
+        """
+        self.node_name = self.__class__.__name__.lower().replace('node', '')
+        self.llm = ctx.llm_registry.get_llm(self.node_name)
+
+        self.prompt = ChatPromptTemplate.from_template(PLANNER_PROMPT)
+        self.chain = self.prompt | self.llm.with_structured_output(PlanModel)
+
+    def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
+        """Executes the planning node.
+
+        Args:
+           state (GraphState): The current state of the execution graph.
+
+        Returns:s
+            Dict[str, Any]: A dictionary containing the generated 'plan', 'reasoning',
+                and any 'errors' encountered.
+        """
+        try:
+            relevant_tables = '\n'.join(
+                t.model_dump_json(indent=2) for t in state.relevant_tables
+            )
+
+
+            feedback = ""
+            if state.errors:
+                feedback = "\n".join(e.model_dump_json(indent=2) for e in state.errors)
+
+            query_text = state.sub_query.intent if state.sub_query else ""
+            expected_schema = []
+            if state.sub_query and state.sub_query.expected_schema:
+                expected_schema = [c.model_dump() for c in state.sub_query.expected_schema]
+            plan: PlanModel = self.chain.invoke(
+                {
+                    "relevant_tables": relevant_tables,
+                    "examples": PLANNER_EXAMPLES,
+                    "feedback": feedback,
+                    "expected_schema": expected_schema,
+                    "semantic_context": "",
+                    "user_query": query_text,
+                }
+            )
+
+            return {
+                "ast_planner_response": ASTPlannerResponse(plan=plan),
+                "reasoning": [
+                    {
+                        "node": self.node_name,
+                        "content": [
+                            f"Reasoning: {plan.reasoning or 'None'}",
+                            f"Tables: {', '.join(t.name for t in plan.tables)}",
+                        ],
+                    }
+                ],
+                "errors": [],
+            }
+
+        except Exception as exc:
+            logger.exception("Planner failed")
+            return {
+                "ast_planner_response": ASTPlannerResponse(plan=None),
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message="Planner failed.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.PLANNING_FAILURE,
+                        stack_trace=traceback.format_exc(),
+                    )
+                ],
+            }