nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/nodes/ast_planner/prompts.py

@@ -0,0 +1,138 @@
+"""Prompts and examples for the SQL Planner node."""
+
+PLANNER_EXAMPLES = """
+Examples:
+
+User Query: "Show me the names of users who placed orders in 2023"
+Semantic Context:
+{
+  "canonical_query": "List names of users with orders in 2023",
+  "keywords": ["users", "orders"]
+}
+
+Plan:
+{
+  "reasoning": "Filter orders by year 2023. Join users. Select user name.",
+  "tables": [
+    {"name": "users", "alias": "t1", "ordinal": 0},
+    {"name": "orders", "alias": "t2", "ordinal": 1}
+  ],
+  "joins": [
+    {
+      "left_alias": "t1",
+      "right_alias": "t2",
+      "join_type": "inner",
+      "ordinal": 0,
+      "condition": {
+        "kind": "binary",
+        "op": "=",
+        "left": {"kind": "column", "alias": "t1", "column_name": "id"},
+        "right": {"kind": "column", "alias": "t2", "column_name": "user_id"}
+      }
+    }
+  ],
+  "where": {
+    "kind": "binary",
+    "op": "AND",
+    "left": {
+      "kind": "binary",
+      "op": ">=",
+      "left": {"kind": "column", "alias": "t2", "column_name": "order_date"},
+      "right": {"kind": "literal", "value": "2023-01-01"}
+    },
+    "right": {
+      "kind": "binary",
+      "op": "<=",
+      "left": {"kind": "column", "alias": "t2", "column_name": "order_date"},
+      "right": {"kind": "literal", "value": "2023-12-31"}
+    }
+  },
+  "select_items": [
+    {
+      "ordinal": 0,
+      "expr": {"kind": "column", "alias": "t1", "column_name": "name"},
+      "alias": "user_name"
+    }
+  ]
+}
+
+User Query: "Total revenue by region last quarter"
+Expected Schema:
+[
+  {"name": "region", "dtype": "string"},
+  {"name": "total_revenue", "dtype": "float"}
+]
+
+Plan:
+{
+  "reasoning": "Group by region and sum revenue.",
+  "tables": [
+    {"name": "orders", "alias": "t1", "ordinal": 0}
+  ],
+  "joins": [],
+  "where": {
+    "kind": "binary",
+    "op": "=",
+    "left": {"kind": "column", "alias": "t1", "column_name": "quarter"},
+    "right": {"kind": "literal", "value": "last_quarter"}
+  },
+  "select_items": [
+    {
+      "ordinal": 0,
+      "expr": {"kind": "column", "alias": "t1", "column_name": "region"},
+      "alias": "region"
+    },
+    {
+      "ordinal": 1,
+      "expr": {"kind": "func", "func_name": "SUM", "args": [{"kind": "column", "alias": "t1", "column_name": "revenue"}], "is_aggregate": true},
+      "alias": "total_revenue"
+    }
+  ],
+  "group_by": [
+    {
+      "ordinal": 0,
+      "expr": {"kind": "column", "alias": "t1", "column_name": "region"}
+    }
+  ]
+}
+"""
+
+PLANNER_PROMPT = (
+    "[ROLE]\n"
+    "You are a SQL Planner. Your job is to create a structured, executable SQL plan"
+    " in the form of a deterministic Abstract Syntax Tree (AST).\n\n"
+
+    "[INSTRUCTIONS]\n"
+    "1. Analyze [USER_QUERY] and [SEMANTIC_CONTEXT].\n"
+    "2. Select ONLY tables from [RELEVANT_TABLES]. Assign strict 'ordinal' positions 0..N.\n"
+    "3. When joining, use relationships listed in [RELEVANT_TABLES]. If no relationship exists, do not join.\n"
+    "4. Define joins using ONLY table aliases (left_alias/right_alias).\n"
+    "5. Build Expr trees using:\n"
+    "   literal | column | func | binary | unary | case\n"
+    "6. Every list MUST contain `ordinal` fields in ascending order starting at 0.\n"
+    "7. Order lists to match ordinals (0..N) exactly.\n"
+    "8. For literal values on '=' or 'IN', choose values from column stats if available.\n"
+    "9. If no exact match is available, fall back to LIKE but keep the pattern derived from stats/synonyms.\n\n"
+
+    "[OUTPUT CONTRACT]\n"
+    "- If [EXPECTED_SCHEMA] is provided and non-empty:\n"
+    "  - select_items length MUST equal expected_schema length.\n"
+    "  - select_items aliases MUST match expected_schema names in the same order.\n"
+    "- All table/column references MUST come from [RELEVANT_TABLES].\n"
+    "- Joins MUST use relationships provided in [RELEVANT_TABLES].\n"
+    "- The [EXAMPLES] are illustrative; always follow [EXPECTED_SCHEMA] when provided.\n\n"
+
+    "[CONSTRAINTS]\n"
+    "- STRICTLY follow PlanModel schema.\n"
+    "- Do NOT hallucinate tables or columns.\n"
+    "- Do NOT output text, ONLY the JSON object.\n"
+    "- Use ISO 8601 dates.\n"
+    "- No extra keys beyond the schema.\n\n"
+
+    "[RELEVANT_TABLES]\n{relevant_tables}\n\n"
+    "[EXPECTED_SCHEMA]\n{expected_schema}\n\n"
+    "[SEMANTIC_CONTEXT]\n{semantic_context}\n\n"
+    "[EXAMPLES]\n{examples}\n\n"
+    "[FEEDBACK]\n{feedback}\n\n"
+    "[USER_QUERY]\n{user_query}"
+)

nl2sql/pipeline/nodes/ast_planner/schemas.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+from typing import List, Optional, Literal, Union
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class CaseWhen(BaseModel):
+    """Represents a single WHEN ... THEN ... clause in a CASE expression.
+
+    Attributes:
+        condition (Expr): The condition expression to evaluate.
+        result (Expr): The result expression if the condition is true.
+        ordinal (int): The position of this clause in the CASE statement.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    condition: "Expr"
+    result: "Expr"
+    ordinal: int
+
+
+class Expr(BaseModel):
+    """Unified AST for deterministic SQL expressions.
+
+    The 'kind' attribute determines which fields must or may be populated.
+    Strict validation rules are enforced in `model_post_init`.
+
+    Attributes:
+        kind (Literal): The type of expression (literal, column, func, binary, unary, case).
+        value (Optional[Union[str, int, float, bool]]): Value for literal expressions.
+        is_null (bool): Whether the literal is a NULL value.
+        alias (Optional[str]): Table alias for column expressions.
+        column_name (Optional[str]): Name of the column.
+        func_name (Optional[str]): Name of the function.
+        args (List[Expr]): Arguments for function expressions.
+        is_aggregate (bool): Whether the function is an aggregate function.
+        op (Optional[str]): Operator for binary or unary expressions.
+        left (Optional[Expr]): Left operand for binary expressions.
+        right (Optional[Expr]): Right operand for binary expressions.
+        expr (Optional[Expr]): Operand for unary expressions.
+        whens (List[CaseWhen]): List of WHEN clauses for CASE expressions.
+        else_expr (Optional[Expr]): The ELSE expression for CASE expressions.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["literal", "column", "func", "binary", "unary", "case"]
+
+    # LITERAL
+    value: Optional[Union[str, int, float, bool]] = None
+    is_null: bool = False
+
+    # COLUMN
+    alias: Optional[str] = None
+    column_name: Optional[str] = None
+
+    # FUNCTION
+    func_name: Optional[str] = None
+    args: List["Expr"] = Field(default_factory=list)
+    is_aggregate: bool = False
+
+    # OPERATORS
+    op: Optional[
+        Literal[
+            "=", "!=", ">", "<", ">=", "<=",
+            "+", "-", "*", "/", "%",
+            "AND", "OR", "LIKE", "IN",
+            "NOT", "IS", "IS NOT"
+        ]
+    ] = None
+
+    left: Optional["Expr"] = None
+    right: Optional["Expr"] = None
+
+    # UNARY
+    expr: Optional["Expr"] = None
+
+    # CASE
+    whens: List[CaseWhen] = Field(default_factory=list)
+    else_expr: Optional["Expr"] = None
+
+    def model_post_init(self, *_):
+        """Validates the expression based on its `kind`."""
+        k = self.kind
+
+        if k == "literal":
+            if self.value is None and not self.is_null:
+                raise ValueError("Literal must have value or is_null=True")
+
+        if k == "column" and not self.column_name:
+            raise ValueError("column_name is required for column expression")
+
+        if k == "func" and not self.func_name:
+            raise ValueError("func_name is required for func expression")
+
+        if k == "binary":
+            if not (self.left and self.right):
+                raise ValueError("Binary expression requires left and right")
+            if not self.op:
+                raise ValueError("Binary expression requires operator")
+
+        if k == "unary":
+            if not self.expr:
+                raise ValueError("Unary expression requires expr")
+            if self.op not in ("NOT", "+", "-"):
+                raise ValueError("Unary op must be NOT, +, or -")
+
+        if k == "case" and not self.whens:
+            raise ValueError("CASE expression must have at least one WHEN")
+
+
+class TableRef(BaseModel):
+    """Represents a reference to a database table.
+
+    Attributes:
+        name (str): The name of the table.
+        schema_name (Optional[str]): The schema of the table.
+        database (Optional[str]): The database name.
+        alias (str): The alias used for the table in the query.
+        ordinal (int): The strict ordinal position of the table.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    name: str
+    schema_name: Optional[str] = None
+    database: Optional[str] = None
+    alias: str
+    ordinal: int
+
+
+class JoinSpec(BaseModel):
+    """Specification for a table join.
+
+    Attributes:
+        left_alias (str): Alias of the left table.
+        right_alias (str): Alias of the right table.
+        join_type (Literal): Type of join (inner, left, right, full).
+        condition (Expr): The join condition expression.
+        ordinal (int): The strict ordinal position of the join.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    left_alias: str
+    right_alias: str
+    join_type: Literal["inner", "left", "right", "full"] = "inner"
+    condition: Expr
+    ordinal: int
+
+
+class SelectItem(BaseModel):
+    """Represents an item in the SELECT clause.
+
+    Attributes:
+        expr (Expr): The expression to select.
+        alias (Optional[str]): The alias for the selected expression.
+        ordinal (int): The strict ordinal position of the select item.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    expr: Expr
+    alias: Optional[str] = None
+    ordinal: int
+
+
+class OrderItem(BaseModel):
+    """Represents an item in the ORDER BY clause.
+
+    Attributes:
+        expr (Expr): The expression to order by.
+        direction (Literal): The sort direction (asc, desc).
+        ordinal (int): The strict ordinal position of the order item.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    expr: Expr
+    direction: Literal["asc", "desc"] = "asc"
+    ordinal: int
+
+
+class GroupByItem(BaseModel):
+    """Represents an item in the GROUP BY clause.
+
+    Attributes:
+        expr (Expr): The expression to group by.
+        ordinal (int): The strict ordinal position of the group by item.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    expr: Expr
+    ordinal: int
+
+
+class PlanModel(BaseModel):
+    """Standardized representation of a SQL execution plan.
+
+    Attributes:
+        query_type (Literal): The type of query (default: READ).
+        distinct (bool): Whether to select distinct rows.
+        tables (List[TableRef]): List of tables involved in the query.
+        joins (List[JoinSpec]): List of join specifications.
+        select_items (List[SelectItem]): List of items to select.
+        where (Optional[Expr]): The WHERE clause expression.
+        group_by (List[GroupByItem]): List of GROUP BY items.
+        having (Optional[Expr]): The HAVING clause expression.
+        order_by (List[OrderItem]): List of ORDER BY items.
+        limit (Optional[int]): The LIMIT count.
+        offset (Optional[int]): The OFFSET count.
+        reasoning (Optional[str]): Explanatory text for the plan generation.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    query_type: Literal["READ"] = "READ"
+    distinct: bool = False
+
+    tables: List[TableRef] = Field(default_factory=list)
+    joins: List[JoinSpec] = Field(default_factory=list)
+
+    select_items: List[SelectItem] = Field(default_factory=list)
+
+    where: Optional[Expr] = None
+    group_by: List[GroupByItem] = Field(default_factory=list)
+    having: Optional[Expr] = None
+    order_by: List[OrderItem] = Field(default_factory=list)
+
+    limit: Optional[int] = None
+    offset: Optional[int] = None
+
+    reasoning: Optional[str] = None
+
+
+class ASTPlannerResponse(BaseModel):
+    plan: Optional[PlanModel] = None
+
+
+# Fix forward references
+Expr.model_rebuild()
+CaseWhen.model_rebuild()
+PlanModel.model_rebuild()

nl2sql/pipeline/nodes/datasource_resolver/__init__.py

@@ -0,0 +1,4 @@
+from .node import DatasourceResolverNode
+from .schemas import DatasourceResolverResponse, ResolvedDatasource
+
+__all__ = ["DatasourceResolverNode", "DatasourceResolverResponse", "ResolvedDatasource"]

nl2sql/pipeline/nodes/datasource_resolver/node.py

@@ -0,0 +1,253 @@
+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 langchain_core.documents import Document
+from nl2sql.auth import UserContext
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from nl2sql.common.settings import settings
+from .schemas import DatasourceResolverResponse, ResolvedDatasource
+
+
+logger = get_logger("datasource_resolver")
+
+
+class DatasourceResolverNode:
+    """Resolves candidate datasources using vector search over datasource chunks."""
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+        self.vector_store = ctx.vector_store
+        self.rbac = ctx.rbac
+        self.ds_registry = ctx.ds_registry
+        self.schema_store = ctx.schema_store
+
+    def _get_unsupported_datasources(self, datasource_ids: list[str]) -> list[str]:
+        available_ds_ids = self.ds_registry.list_ids()
+        unsupported = [ds_id for ds_id in datasource_ids if ds_id not in available_ds_ids]
+        return sorted(unsupported)
+
+    def _error_response(
+        self,
+        resolved_datasources: list[ResolvedDatasource],
+        allowed_ids: list[str],
+        unsupported_ids: list[str],
+        message: str,
+        severity: ErrorSeverity,
+        error_code: ErrorCode,
+    ) -> Dict[str, Any]:
+        return {
+            "datasource_resolver_response": DatasourceResolverResponse(
+                resolved_datasources=resolved_datasources,
+                allowed_datasource_ids=allowed_ids,
+                unsupported_datasource_ids=unsupported_ids,
+            ),
+            "errors": [
+                PipelineError(
+                    node=self.node_name,
+                    message=message,
+                    severity=severity,
+                    error_code=error_code,
+                )
+            ],
+        }
+
+    def _get_candidate_datasources(
+        self,
+        candidate_docs: list[Document],
+    ) -> Dict[str, ResolvedDatasource]:
+        candidate_datasources: Dict[str, ResolvedDatasource] = {}
+        schema_versions: Dict[str, str | None] = {}
+        for doc in candidate_docs:
+            ds_id = doc.metadata.get("datasource_id")
+            if not ds_id or ds_id in candidate_datasources:
+                continue
+            if ds_id not in schema_versions:
+                schema_versions[ds_id] = self.schema_store.get_latest_version(ds_id)
+            chunk_schema_version = doc.metadata.get("schema_version")
+            schema_version = schema_versions.get(ds_id)
+            candidate_datasources[ds_id] = ResolvedDatasource(
+                datasource_id=ds_id,
+                metadata=dict(doc.metadata),
+                schema_version=schema_version,
+                chunk_schema_version=chunk_schema_version,
+                schema_version_mismatch=bool(
+                    chunk_schema_version
+                    and schema_version
+                    and chunk_schema_version != schema_version
+                ),
+            )
+        return candidate_datasources
+
+    def _get_allowed_datasource_ids(
+        self,
+        user_context: UserContext,
+        candidate_ids: list[str],
+    ) -> list[str]:
+        allowed_ids = self.rbac.get_allowed_datasources(user_context)
+        if not allowed_ids:
+            return []
+        if "*" in allowed_ids:
+            return candidate_ids
+        return [ds_id for ds_id in candidate_ids if ds_id in allowed_ids]
+
+    
+    def _apply_schema_version_mismatch_policy(self, resolved_datasources: list[ResolvedDatasource], allowed_ids: list[str], unsupported_ids: list[str]):
+        mismatches = [
+                ds.datasource_id
+                for ds in resolved_datasources
+                if ds.schema_version_mismatch
+        ]
+        if mismatches:
+            policy = (settings.schema_version_mismatch_policy or "warn").lower()
+            message = (
+                "Schema version mismatch for datasources: "
+                + ", ".join(sorted(mismatches))
+            )
+            if policy == "fail":
+                return self._error_response(
+                    resolved_datasources=[],
+                    allowed_ids=[],
+                    unsupported_ids=[],
+                    message=message,
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.INVALID_STATE,
+                )
+            elif policy == "warn":
+                return {
+                    "datasource_resolver_response": DatasourceResolverResponse(
+                        resolved_datasources=resolved_datasources,
+                        allowed_datasource_ids=allowed_ids,
+                        unsupported_datasource_ids=unsupported_ids,
+                    ),
+                    "reasoning": [
+                        {
+                            "node": self.node_name,
+                            "content": message,
+                            "type": "warning",
+                        }
+                    ],
+                        "warnings": [
+                            {
+                                "node": self.node_name,
+                                "content": message,
+                            }
+                        ],
+                }
+        
+
+    def __call__(self, state: GraphState) -> Dict[str, Any]:
+        try:
+            if state.datasource_id:
+                unsupported_ids = self._get_unsupported_datasources([state.datasource_id])
+                if unsupported_ids:
+                    return self._error_response(
+                        resolved_datasources=[],
+                        allowed_ids=[],
+                        unsupported_ids=unsupported_ids,
+                        message=f"Datasource not found: {state.datasource_id}.",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.INVALID_STATE,
+                    )
+
+                resolved = ResolvedDatasource(
+                    datasource_id=state.datasource_id,
+                    metadata={},
+                    schema_version=self.schema_store.get_latest_version(state.datasource_id),
+                )
+                allowed_ids = self._get_allowed_datasource_ids(
+                    state.user_context,
+                    [state.datasource_id],
+                )
+                if not allowed_ids:
+                    return self._error_response(
+                        resolved_datasources=[resolved],
+                        allowed_ids=[],
+                        unsupported_ids=[],
+                        message="Datasource not allowed.",
+                        severity=ErrorSeverity.CRITICAL,
+                        error_code=ErrorCode.SECURITY_VIOLATION,
+                    )
+
+                return {
+                    "datasource_resolver_response": DatasourceResolverResponse(
+                        resolved_datasources=[resolved],
+                        allowed_datasource_ids=allowed_ids,
+                        unsupported_datasource_ids=[],
+                    ),
+                    "reasoning": [
+                        {
+                            "node": self.node_name,
+                            "content": "Using explicit datasource override.",
+                        }
+                    ],
+                }
+
+            if not self.vector_store:
+                return {
+                    "datasource_resolver_response": DatasourceResolverResponse(),
+                    "reasoning": [{"node": self.node_name, "content": "Vector store unavailable."}],
+                }
+
+            query = state.user_query
+            candidate_docs = self.vector_store.retrieve_datasource_candidates(query, k=5)
+            candidate_datasources = self._get_candidate_datasources(candidate_docs)
+            candidate_ids = list(candidate_datasources.keys())
+            if not candidate_ids:
+                return self._error_response(
+                    resolved_datasources=[],
+                    allowed_ids=[],
+                    unsupported_ids=[],
+                    message="No datasource candidates resolved.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.SCHEMA_RETRIEVAL_FAILED,
+                )
+            unsupported_ids = self._get_unsupported_datasources(candidate_ids)
+            allowed_ids = self._get_allowed_datasource_ids(state.user_context, candidate_ids)
+            if not allowed_ids:
+                return self._error_response(
+                    resolved_datasources=list(candidate_datasources.values()),
+                    allowed_ids=[],
+                    unsupported_ids=unsupported_ids,
+                    message="No allowed datasources.",
+                    severity=ErrorSeverity.CRITICAL,
+                    error_code=ErrorCode.SECURITY_VIOLATION,
+                )
+
+
+            schema_version_mismatch_response = self._apply_schema_version_mismatch_policy(
+                list(candidate_datasources.values()),
+                allowed_ids,
+                unsupported_ids,
+            )
+            if schema_version_mismatch_response:
+                return schema_version_mismatch_response
+            
+
+            result = DatasourceResolverResponse(
+                resolved_datasources=list(candidate_datasources.values()),
+                allowed_datasource_ids=allowed_ids,
+                unsupported_datasource_ids=unsupported_ids,
+            )               
+            return {
+                "datasource_resolver_response": result,
+                "reasoning": [{"node": self.node_name, "content": "Ranked by vector similarity."}],
+            }
+        except Exception as exc:
+            logger.error(f"Datasource resolver failed: {exc}")
+            return {
+                "datasource_resolver_response": DatasourceResolverResponse(),
+                "errors": [
+                    PipelineError(
+                        node=self.node_name,
+                        message=f"Datasource resolution failed: {exc}",
+                        severity=ErrorSeverity.ERROR,
+                        error_code=ErrorCode.SCHEMA_RETRIEVAL_FAILED,
+                    )
+                ],
+            }

nl2sql/pipeline/nodes/datasource_resolver/schemas.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class ResolvedDatasource(BaseModel):
+    datasource_id: str
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    schema_version: Optional[str] = None
+    chunk_schema_version: Optional[str] = None
+    schema_version_mismatch: bool = False
+
+
+class DatasourceResolverResponse(BaseModel):
+    resolved_datasources: List[ResolvedDatasource] = Field(default_factory=list)
+    allowed_datasource_ids: List[str] = Field(default_factory=list)
+    unsupported_datasource_ids: List[str] = Field(default_factory=list)
+
+

nl2sql/pipeline/nodes/decomposer/__init__.py

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