nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/nodes/decomposer/node.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+from typing import Dict, Any, TYPE_CHECKING
+
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import Runnable
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import GraphState
+
+from .schemas import DecomposerResponse, SubQuery, UnmappedSubQuery, PostCombineOp
+from .prompts import DECOMPOSER_PROMPT
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+import hashlib
+import json
+
+logger = get_logger("decomposer")
+
+
+
+class DecomposerNode:
+    """Orchestrates query decomposition and routing.
+
+    Analyzes the user query to generate semantic sub-queries and combine groups.
+
+    Attributes:
+        llm (ChatOpenAI): The language model to use.
+        prompt (ChatPromptTemplate): The prompt template.
+        chain (Runnable): The execution chain.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        """Initializes the DecomposerNode.
+
+        Args:
+            llm (LLMCallable): The LLM instance or runnable.
+            vector_store (Optional[VectorStore]): Vector store for RAG.
+        """
+        self.node_name = self.__class__.__name__.lower().replace('node', '')
+        self.llm = ctx.llm_registry.get_llm(self.node_name)
+        self.prompt = ChatPromptTemplate.from_template(DECOMPOSER_PROMPT)
+        self.chain = self.prompt | self.llm.with_structured_output(
+            DecomposerResponse, method="function_calling"
+        )
+
+    def _stable_id(self, prefix: str, payload: Dict[str, Any]) -> str:
+        data = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
+        digest = hashlib.sha256(data.encode("utf-8")).hexdigest()[:12]
+        return f"{prefix}_{digest}"
+
+    def __call__(self, state: GraphState) -> Dict[str, Any]:
+        """Executes the decomposer node.
+
+        Invokes the LLM to produce semantic sub-queries and combine groups.
+
+        Args:
+            state (GraphState): The current execution state.
+
+        Returns:
+            Dict[str, Any]: Dictionary containing 'sub_queries', confidence, reasoning, etc.
+        """
+        try:
+            resolver_response = state.datasource_resolver_response
+            if not resolver_response or not resolver_response.resolved_datasources:
+                raise ValueError("Unable to resolve any datasource for the current user query.")
+            
+            resolved_datasources = resolver_response.resolved_datasources 
+            resolved_payload = []
+            resolved_ids = set()
+            schema_version_map = {}
+            for datasource in resolved_datasources:
+                resolved_payload.append(datasource.model_dump())
+                resolved_ids.add(datasource.datasource_id)
+                schema_version_map[datasource.datasource_id] = datasource.schema_version
+
+            llm_response: DecomposerResponse = self.chain.invoke(
+                {
+                    "user_query": state.user_query,
+                    "resolved_datasources": resolved_payload,
+                }
+            )
+
+            final_sub_queries = []
+            unmapped = []
+            allowed_ids = set(resolver_response.allowed_datasource_ids)
+            unsupported_ids = set(resolver_response.unsupported_datasource_ids)
+            id_map: Dict[str, str] = {}
+
+            for llm_sq in llm_response.sub_queries:
+                datasource_id = (llm_sq.datasource_id or "").strip() or None
+                if not datasource_id or datasource_id not in resolved_ids:
+                    unmapped.append(
+                        UnmappedSubQuery(
+                            intent=llm_sq.intent,
+                            reason="no_datasource",
+                            datasource_id=datasource_id,
+                            detail="Datasource is missing or not resolved for this sub-query.",
+                        )
+                    )
+                    continue
+                if datasource_id not in allowed_ids:
+                    unmapped.append(
+                        UnmappedSubQuery(
+                            intent=llm_sq.intent,
+                            reason="restricted_datasource",
+                            datasource_id=datasource_id,
+                            detail="Datasource is not allowed for the current user context.",
+                        )
+                    )
+                    continue
+                if datasource_id in unsupported_ids:
+                    unmapped.append(
+                        UnmappedSubQuery(
+                            intent=llm_sq.intent,
+                            reason="unsupported_datasource",
+                            datasource_id=datasource_id,
+                            detail="Datasource does not match any supported adapter capabilities.",
+                        )
+                    )
+                    continue
+                stable_id = self._stable_id(
+                    "sq",
+                    {
+                        "datasource_id": datasource_id,
+                        "intent": llm_sq.intent,
+                        "metrics": [m.model_dump() for m in llm_sq.metrics],
+                        "filters": [f.model_dump() for f in llm_sq.filters],
+                        "group_by": [g.model_dump() for g in llm_sq.group_by],
+                        "expected_schema": [c.model_dump() for c in llm_sq.expected_schema],
+                    },
+                )
+                id_map[llm_sq.id] = stable_id
+                sq = SubQuery(
+                    id=stable_id,
+                    intent=llm_sq.intent,
+                    datasource_id=datasource_id,
+                    metrics=llm_sq.metrics,
+                    filters=llm_sq.filters,
+                    group_by=llm_sq.group_by,
+                    expected_schema=llm_sq.expected_schema,
+                    schema_version=schema_version_map.get(datasource_id),
+                )
+                final_sub_queries.append(sq)
+
+            valid_ids = {sq.id for sq in final_sub_queries}
+            combine_groups = []
+            for group in llm_response.combine_groups:
+                updated_inputs = []
+                for inp in group.inputs:
+                    mapped_id = id_map.get(inp.subquery_id, inp.subquery_id)
+                    if mapped_id in valid_ids:
+                        updated_inputs.append(inp.model_copy(update={"subquery_id": mapped_id}))
+                if not updated_inputs:
+                    continue
+                combine_groups.append(group.model_copy(update={"inputs": updated_inputs}))
+
+            combine_groups = sorted(combine_groups, key=lambda g: g.group_id)
+            final_sub_queries = sorted(final_sub_queries, key=lambda s: s.id)
+
+            post_combine_ops = []
+            for op in llm_response.post_combine_ops or []:
+                op_id = self._stable_id(
+                    "op",
+                    {
+                        "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,
+                    },
+                )
+                post_combine_ops.append(op.model_copy(update={"op_id": op_id}))
+
+            post_combine_ops = sorted(post_combine_ops, key=lambda o: o.op_id)
+
+            response = DecomposerResponse(
+                sub_queries=final_sub_queries,
+                combine_groups=combine_groups,
+                post_combine_ops=post_combine_ops,
+                unmapped_subqueries=unmapped,
+            )
+
+            return {
+                "decomposer_response": response,
+                "reasoning": [{"node": self.node_name, "content": "Decomposition completed."}],
+            }
+
+        except Exception as e:
+            logger.error(f"Node {self.node_name} failed: {e}")
+
+            return {
+                "decomposer_response": DecomposerResponse(
+                    sub_queries=[],
+                    combine_groups=[],
+                    post_combine_ops=[],
+                    unmapped_subqueries=[],
+                ),
+                "reasoning": [
+                    {
+                        "node": self.node_name,
+                        "content": f"Decomposition failed: {str(e)}",
+                        "type": "error",
+                    }
+                ],
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message=f"Decomposition failed: {str(e)}",
+                        severity=ErrorSeverity.CRITICAL,
+                        error_code=ErrorCode.ORCHESTRATOR_CRASH,
+                        stack_trace=str(e),
+                    )
+                ],
+            }

nl2sql/pipeline/nodes/decomposer/prompts.py

@@ -0,0 +1,96 @@
+"""Prompts for the Query Decomposer node."""
+
+DECOMPOSER_PROMPT = """
+SYSTEM:
+You are a Semantic Query Decomposer. You output ONLY structured semantic intent.
+
+TASK:
+Decompose the user query into semantic sub-queries and combine groups.
+Your output must be deterministic and strictly follow the JSON contract.
+
+INPUTS:
+User Query:
+{user_query}
+
+Resolved Datasources (id + semantic metadata):
+{resolved_datasources}
+
+RULES:
+1) Use resolved_datasources metadata to select the most appropriate datasource for each subquery.
+2) If an intent cannot be mapped to any resolved datasource, emit it under unmapped_subqueries.
+3) SubQueries must contain ONLY semantic intent:
+   - metrics
+   - filters
+   - group_by
+4) Do NOT emit:
+   - SQL
+   - table names
+   - column names
+   - joins
+   - physical schema
+5) Define combine_groups explicitly using:
+   - standalone
+   - compare
+   - join
+   - union
+6) For join or compare:
+   - include join_keys as left/right semantic attribute pairs.
+7) Any filters, metrics, group_by, order_by, or limits that apply AFTER a combine must be emitted in post_combine_ops.
+8) expected_schema must be derived strictly from semantic intent (metrics + group_by) and be minimal.
+   It defines the semantic output contract for downstream aggregation, not physical columns.
+9) Do not invent attributes not implied by the user query or datasource metadata.
+10) Output JSON only. No commentary.
+
+OUTPUT FORMAT:
+Return JSON exactly matching this structure:
+
+{{
+  "sub_queries": [
+    {{
+      "id": "sq_1",
+      "datasource_id": "ds_sales",
+      "intent": "total revenue by region last quarter",
+      "metrics": [{{"name": "total_revenue", "aggregation": "sum"}}],
+      "filters": [{{"attribute": "time_period", "operator": "=", "value": "last_quarter"}}],
+      "group_by": [{{"attribute": "region"}}],
+      "expected_schema": [
+        {{"name": "region", "dtype": "string"}},
+        {{"name": "total_revenue", "dtype": "float"}}
+      ]
+    }}
+  ],
+  "combine_groups": [
+    {{
+      "group_id": "cg_1",
+      "operation": "standalone",
+      "inputs": [
+        {{"subquery_id": "sq_1", "role": "base"}}
+      ],
+      "join_keys": []
+    }}
+  ],
+  "post_combine_ops": [
+    {{
+      "op_id": "op_1",
+      "target_group_id": "cg_1",
+      "operation": "filter",
+      "filters": [{{"attribute": "total_revenue", "operator": ">", "value": 1000}}],
+      "metrics": [],
+      "group_by": [],
+      "order_by": [{{"attribute": "total_revenue", "direction": "desc"}}],
+      "limit": 10,
+      "expected_schema": [
+        {{"name": "region", "dtype": "string"}},
+        {{"name": "total_revenue", "dtype": "float"}}
+      ]
+    }}
+  ],
+  "unmapped_subqueries": []
+}}
+
+VALIDATION:
+- sub_queries must be non-empty unless all intents are unmapped.
+- combine_groups must reference valid subquery_id values.
+- post_combine_ops.target_group_id must reference an existing group_id.
+- expected_schema must match semantic outputs only.
+"""

nl2sql/pipeline/nodes/decomposer/schemas.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+from typing import List, Optional, Literal, Dict, Any
+from pydantic import BaseModel, Field, model_validator
+
+
+def _contains_physical_tokens(text: str) -> bool:
+    forbidden = ["select", "from", "join", "where", "group by", "order by", ";", "--", "/*", "*/"]
+    lowered = text.lower()
+    return any(token in lowered for token in forbidden)
+
+
+class MetricSpec(BaseModel):
+    name: str
+    aggregation: Optional[Literal["count", "sum", "avg", "min", "max"]] = None
+    description: Optional[str] = None
+
+
+class FilterSpec(BaseModel):
+    attribute: str
+    operator: Literal["=", "!=", ">", ">=", "<", "<=", "between", "in", "contains"]
+    value: str | int | float | bool | List[str] | List[int] | List[float]
+
+
+class GroupBySpec(BaseModel):
+    attribute: str
+
+
+class OrderBySpec(BaseModel):
+    attribute: str
+    direction: Literal["asc", "desc"] = "asc"
+
+
+class ExpectedColumn(BaseModel):
+    name: str
+    dtype: Optional[Literal["string", "int", "float", "bool", "date", "datetime"]] = None
+
+
+class SubQuery(BaseModel):
+    id: str
+    datasource_id: str
+    intent: str
+    kind: Literal["scan"] = "scan"
+    metrics: List[MetricSpec] = Field(default_factory=list)
+    filters: List[FilterSpec] = Field(default_factory=list)
+    group_by: List[GroupBySpec] = Field(default_factory=list)
+    expected_schema: List[ExpectedColumn] = Field(default_factory=list)
+    schema_version: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_semantic_only(self):
+        content = " ".join(
+            [
+                self.intent,
+                " ".join(m.name for m in self.metrics),
+                " ".join(f.attribute for f in self.filters),
+                " ".join(g.attribute for g in self.group_by),
+                " ".join(c.name for c in self.expected_schema),
+            ]
+        ).lower()
+        if _contains_physical_tokens(content):
+            raise ValueError("SubQuery contains SQL or physical schema tokens.")
+        return self
+
+
+class CombineInput(BaseModel):
+    subquery_id: str
+    role: Optional[Literal["left", "right", "base", "compare", "primary", "secondary"]] = None
+
+
+class JoinKeyPair(BaseModel):
+    left: str
+    right: str
+
+
+class CombineGroup(BaseModel):
+    group_id: str
+    operation: Literal["standalone", "compare", "join", "union"]
+    inputs: List[CombineInput]
+    join_keys: List[JoinKeyPair] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def validate_roles(self):
+        if self.operation in {"compare", "join"}:
+            if any(i.role is None for i in self.inputs):
+                raise ValueError("Compare/join combine groups require roles.")
+            if not self.join_keys:
+                raise ValueError("Compare/join combine groups require join_keys.")
+        return self
+
+
+class PostCombineOp(BaseModel):
+    op_id: str
+    target_group_id: str
+    operation: Literal["filter", "aggregate", "project", "sort", "limit"]
+    filters: List[FilterSpec] = Field(default_factory=list)
+    metrics: List[MetricSpec] = Field(default_factory=list)
+    group_by: List[GroupBySpec] = Field(default_factory=list)
+    order_by: List[OrderBySpec] = Field(default_factory=list)
+    limit: Optional[int] = None
+    expected_schema: List[ExpectedColumn] = Field(default_factory=list)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def validate_semantic_only(self):
+        content = " ".join(
+            [
+                " ".join(m.name for m in self.metrics),
+                " ".join(f.attribute for f in self.filters),
+                " ".join(g.attribute for g in self.group_by),
+                " ".join(o.attribute for o in self.order_by),
+                " ".join(c.name for c in self.expected_schema),
+            ]
+        )
+        if _contains_physical_tokens(content):
+            raise ValueError("PostCombineOp contains SQL or physical schema tokens.")
+        return self
+
+
+class UnmappedSubQuery(BaseModel):
+    intent: str
+    reason: Literal["no_datasource", "restricted_datasource", "unsupported_datasource"]
+    datasource_id: Optional[str] = None
+    detail: Optional[str] = None
+
+
+class DecomposerResponse(BaseModel):
+    sub_queries: List[SubQuery]
+    combine_groups: List[CombineGroup]
+    post_combine_ops: List[PostCombineOp] = Field(default_factory=list)
+    unmapped_subqueries: List[UnmappedSubQuery] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def validate_references(self):
+        ids = {sq.id for sq in self.sub_queries}
+        for group in self.combine_groups:
+            for inp in group.inputs:
+                if inp.subquery_id not in ids:
+                    raise ValueError(f"CombineGroup references unknown subquery: {inp.subquery_id}")
+        group_ids = {g.group_id for g in self.combine_groups}
+        for op in self.post_combine_ops:
+            if op.target_group_id not in group_ids:
+                raise ValueError(f"PostCombineOp references unknown combine group: {op.target_group_id}")
+        return self

nl2sql/pipeline/nodes/executor/__init__.py

@@ -0,0 +1,3 @@
+from .node import ExecutorNode
+
+__all__ = ["ExecutorNode"]

nl2sql/pipeline/nodes/executor/node.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Any
+
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from nl2sql.execution.contracts import ExecutorRequest
+from nl2sql.execution.executor import SqlExecutorService
+from nl2sql_adapter_sdk.capabilities import DatasourceCapability
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import SubgraphExecutionState
+
+logger = get_logger("executor")
+
+
+class ExecutorNode:
+    """Thin wrapper that delegates to executor services based on capabilities."""
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+        self.ds_registry = ctx.ds_registry
+        self.executor = SqlExecutorService(ctx.ds_registry)
+        self.tenant_id = ctx.tenant_id
+
+    def _supports_sql(self, ds_id: str) -> bool:
+        """Reports whether a datasource may be executed as SQL.
+
+        The SQL executor is the only executor, so a datasource that does not
+        declare SUPPORTS_SQL -- or whose adapter cannot report capabilities --
+        gets no executor at all rather than being run as SQL anyway.
+        """
+        adapter = self.ds_registry.get_adapter(ds_id)
+        try:
+            capabilities = adapter.capabilities()
+        except Exception as exc:
+            logger.error(f"Failed to get capabilities for datasource '{ds_id}'. {exc}")
+            return False
+
+        normalized = {
+            cap.value if isinstance(cap, DatasourceCapability) else str(cap)
+            for cap in capabilities
+        }
+        return DatasourceCapability.SUPPORTS_SQL.value in normalized
+
+    def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
+        try:
+            ds_id = state.sub_query.datasource_id 
+            sql = state.generator_response.sql_draft 
+
+            if not sql:
+                error = PipelineError(
+                    node=self.node_name,
+                    message="No SQL to execute.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.MISSING_SQL,
+                )
+                return {"executor_response": None, "errors": [error]}
+
+            if not ds_id:
+                error = PipelineError(
+                    node=self.node_name,
+                    message="No datasource_id in state.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.MISSING_DATASOURCE_ID,
+                )
+                return {"executor_response": None, "errors": [error]}
+
+            executor = self.executor if self._supports_sql(ds_id) else None
+            if executor is None:
+                error = PipelineError(
+                    node=self.node_name,
+                    message=f"No executor available for datasource '{ds_id}'.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.INVALID_STATE,
+                )
+                return {"executor_response": None, "errors": [error]}
+
+            request = ExecutorRequest(
+                node_id=state.sub_query.id,
+                trace_id=state.trace_id,
+                subgraph_name=state.subgraph_name,
+                datasource_id=ds_id,
+                schema_version=state.sub_query.schema_version,
+                sql=sql,
+                user_context=state.user_context,
+                tenant_id=self.tenant_id,
+            )
+            response = executor.execute(request)
+            return {
+                "executor_response": response,
+                "errors": response.errors,
+                "reasoning": response.reasoning,
+            }
+        except Exception as exc:
+            logger.error(f"Node {self.node_name} failed: {exc}")
+            error = PipelineError(
+                node=self.node_name,
+                message=f"Executor crash: {exc}",
+                severity=ErrorSeverity.CRITICAL,
+                error_code=ErrorCode.EXECUTOR_CRASH,
+            )
+            return {
+                "executor_response": None,
+                "errors": [error],
+            }

nl2sql/pipeline/nodes/generator/__init__.py

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