nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/pipeline/nodes/refiner/node.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import json
+from typing import Callable, Optional, Union, Dict, Any, TYPE_CHECKING
+from langchain_core.runnables import Runnable
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import SubgraphExecutionState
+from .prompts import REFINER_PROMPT
+from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode
+from nl2sql.pipeline.nodes.refiner.schemas import RefinerResponse
+
+from nl2sql.common.logger import get_logger
+
+logger = get_logger("refiner")
+
+LLMCallable = Union[Callable[[str], str], Runnable]
+
+from langchain_core.prompts import ChatPromptTemplate
+
+from langchain_core.output_parsers import StrOutputParser
+from nl2sql.context import NL2SQLContext
+
+class RefinerNode:
+    """
+    Analyzes validation errors and generates constructive feedback for the Planner.
+
+    Uses an LLM to look at the failed plan, the schema, and the errors to suggest fixes.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        """
+        Initializes the RefinerNode.
+
+        Args:
+            llm: The language model to use for refinement.
+        """
+        self.node_name = self.__class__.__name__.lower().replace('node', '')
+        self.llm = ctx.llm_registry.get_llm(self.node_name)
+        self.prompt = ChatPromptTemplate.from_template(REFINER_PROMPT)
+        self.chain = None
+        if self.llm is not None:
+            self.chain = self.prompt | self.llm | StrOutputParser()
+
+    def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
+        """
+        Executes the summarization step.
+
+        Args:
+            state: The current graph state.
+
+        Returns:
+            Dictionary updates for the graph state with refined error messages (feedback).
+        """
+        try:
+            if not self.chain:
+                error = PipelineError(
+                    node=self.node_name,
+                    message="No LLM configured for refiner.",
+                    severity=ErrorSeverity.ERROR,
+                    error_code=ErrorCode.MISSING_LLM,
+                )
+                return {
+                    "refiner_response": RefinerResponse(errors=[error]),
+                    "errors": [error],
+                }
+            relevant_tables = ""
+            if state.relevant_tables:
+                lines = []
+                for tbl in state.relevant_tables:
+                    lines.append(tbl.model_dump_json(indent=2))
+                    lines.append("---")
+                relevant_tables = "\n".join(lines)
+
+            failed_plan_str = "No plan generated."
+            if state.ast_planner_response and state.ast_planner_response.plan:
+                try:
+                    failed_plan_str = json.dumps(state.ast_planner_response.plan, indent=2)
+                except:
+                    failed_plan_str = str(state.ast_planner_response.plan)
+
+            # Extract messages from PipelineError objects
+            errors_str = "\n".join(f"- {e.message}" for e in state.errors)
+
+            reasoning_str = "No reasoning history."
+            if state.reasoning:
+                reasoning_str = "\n".join(
+                    f"[{r.get('node', 'unknown')}]: {r.get('content')}" for r in state.reasoning
+                )
+
+            try:
+                feedback = self.chain.invoke({
+                    "user_query": state.sub_query.intent if state.sub_query else "",
+                    "relevant_tables": relevant_tables,
+                    "failed_plan": failed_plan_str,
+                    "errors": errors_str,
+                    "reasoning": reasoning_str
+                })
+                
+                warning = PipelineError(
+                    node=self.node_name,
+                    message=feedback,
+                    severity=ErrorSeverity.WARNING, # Feedback for retry
+                    error_code=ErrorCode.PLAN_FEEDBACK,
+                )
+                response = RefinerResponse(
+                    feedback=feedback,
+                    errors=[warning],
+                    reasoning=[{"node": self.node_name, "content": feedback}],
+                )
+                return {
+                    "refiner_response": response,
+                    "errors": [warning],
+                    "reasoning": response.reasoning,
+                }
+            except Exception as e:
+                    raise e
+                     
+        except Exception as e:
+            logger.error(f"Node {self.node_name} failed: {e}")
+            error = PipelineError(
+                node=self.node_name,
+                message=f"Refiner failed: {e}",
+                severity=ErrorSeverity.ERROR,
+                error_code=ErrorCode.REFINER_FAILED,
+                stack_trace=str(e),
+            )
+            return {
+                "refiner_response": RefinerResponse(errors=[error]),
+                "reasoning": [{"node": self.node_name, "content": f"Refiner failed: {e}", "type": "error"}],
+                "errors": [error],
+            }

nl2sql/pipeline/nodes/refiner/prompts.py

@@ -0,0 +1,28 @@
+REFINER_PROMPT = """You are an expert SQL debugger and schema analyst.
+Your goal is to analyze a failed SQL generation attempt and provide actionable, schema-aware feedback to the Planner.
+
+### Context
+User Query: "{user_query}"
+
+### Database Schema
+{relevant_tables}
+
+### Failed Plan
+{failed_plan}
+
+### Errors (Validation or Execution)
+{errors}
+
+### Previous Reasoning (Trace)
+{reasoning}
+
+### Instructions
+1. Analyze the Errors in the context of the Schema and User Query.
+2. If the error is an Execution Error (e.g., syntax error, runtime error), analyze why the SQL failed and suggest a fix in the Plan.
+2. If the error is about a missing table or column, check the Schema for the correct name.
+   - Example: Error "Column 'revenue' not found", Schema has "total_revenue". -> Suggest "Use 'total_revenue' instead of 'revenue'".
+3. If the plan is empty or missing, analyze the User Query and suggest which tables/columns to use.
+4. Provide a concise, numbered list of specific fixes. Do not generate SQL. Focus on correcting the Plan.
+
+### Feedback
+"""

nl2sql/pipeline/nodes/refiner/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 RefinerResponse(BaseModel):
+    feedback: Optional[str] = None
+    errors: List[PipelineError] = Field(default_factory=list)
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)

nl2sql/pipeline/nodes/schema_retriever/__init__.py

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

nl2sql/pipeline/nodes/schema_retriever/node.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, Any, List, Optional, TYPE_CHECKING, Set
+
+if TYPE_CHECKING:
+    from nl2sql.pipeline.state import SubgraphExecutionState
+    from nl2sql.pipeline.nodes.decomposer.schemas import SubQuery
+
+from nl2sql.common.logger import get_logger
+from nl2sql.context import NL2SQLContext
+from .schema import Table, Column
+from nl2sql_adapter_sdk.schema import SchemaSnapshot
+
+
+logger = get_logger("schema_retriever")
+
+
+class SchemaRetrieverNode:
+    """Retrieves relevant schema chunks for planning context."""
+
+    def __init__(self, ctx: NL2SQLContext):
+        self.node_name = self.__class__.__name__.lower().replace("node", "")
+        self.vector_store = ctx.vector_store
+        self.schema_store = ctx.schema_store
+
+
+    def _build_semantic_query(self, sub_query: SubQuery) -> str:
+        parts: List[str] = []
+        if sub_query and sub_query.intent:
+            parts.append(sub_query.intent)
+
+        if sub_query and sub_query.filters:
+            filters_text = []
+            for f in sub_query.filters:
+                value = f.value
+                if isinstance(value, list):
+                    value = ", ".join(str(v) for v in value)
+                filters_text.append(f"{f.attribute}={value}")
+            if filters_text:
+                parts.append("filters: " + "; ".join(filters_text))
+
+        if sub_query and sub_query.group_by:
+            group_by = ", ".join(g.attribute for g in sub_query.group_by)
+            if group_by:
+                parts.append("group_by: " + group_by)
+
+        if sub_query and sub_query.expected_schema:
+            expected = ", ".join(c.name for c in sub_query.expected_schema)
+            if expected:
+                parts.append("expected_schema: " + expected)
+
+        if sub_query and sub_query.metrics:
+            metrics = ", ".join(m.name for m in sub_query.metrics)
+            if metrics:
+                parts.append("metrics: " + metrics)
+
+        return "\n".join(parts).strip()
+
+    def _resolve_snapshot(self, datasource_id: str, schema_version: Optional[str]) -> Optional[SchemaSnapshot]:
+        if not self.schema_store:
+            return None
+        if schema_version:
+            return self.schema_store.get_snapshot(datasource_id, schema_version)
+        return self.schema_store.get_latest_snapshot(datasource_id)
+
+    def _build_tables_from_snapshot(
+        self,
+        snapshot: SchemaSnapshot,
+        resolved_tables: Optional[Dict[str, Set[str]]] = None,
+        schema_version: Optional[str] = None,
+    ) -> List[Table]:
+        if not snapshot:
+            return []
+
+        tables_out: List[Table] = []
+        table_keys = (
+            list(snapshot.contract.tables.keys())
+            if not resolved_tables
+            else list(resolved_tables.keys())
+        )
+
+        for table_key in table_keys:
+            table_contract = snapshot.contract.tables.get(table_key)
+            table_metadata = snapshot.metadata.tables.get(table_key)
+            if not table_contract:
+                continue
+
+            resolved_columns = resolved_tables[table_key] if resolved_tables else set()
+            if not resolved_columns:
+                resolved_columns = set(table_contract.columns.keys())
+
+            columns: List[Column] = []
+            for col_key, col_contract in table_contract.columns.items():
+                if col_key not in resolved_columns:
+                    continue
+                col_metadata = table_metadata.columns.get(col_key) if table_metadata else None
+
+                columns.append(
+                    Column(
+                        name=col_key,
+                        type=col_contract.data_type,
+                        stats=col_metadata.statistics.model_dump() if col_metadata and col_metadata.statistics else {},
+                        description=col_metadata.description if col_metadata else ""
+                    )
+                )
+
+            relationships = []
+            for fk in table_contract.foreign_keys:
+                relationships.append(
+                    {
+                        "from_table": table_contract.table.full_name,
+                        "to_table": fk.referred_table.full_name,
+                        "from_columns": fk.constrained_columns,
+                        "to_columns": fk.referred_columns,
+                        "cardinality": fk.cardinality,
+                        "business_meaning": fk.business_meaning,
+                    }
+                )
+
+            primary_keys = [
+                col.name for col in table_contract.columns.values() if col.is_primary_key
+            ]
+
+            table = Table(
+                    name=table_contract.table.table_name,
+                    columns=columns,
+                    description=table_metadata.description if table_metadata else "",
+                    primary_key=primary_keys,
+                    schema_version=schema_version,
+                    relationships=relationships,
+                )
+
+            tables_out.append( table )        
+
+        return tables_out
+
+    def __call__(self, state: SubgraphExecutionState) -> Dict[str, Any]:
+        try:
+            sub_query = state.sub_query
+            if not sub_query:
+                return {"relevant_tables": []}
+
+            datasource_id = sub_query.datasource_id
+            schema_version = sub_query.schema_version
+            query = self._build_semantic_query(sub_query)
+
+            tables: Dict[str, Set[str]] = defaultdict(set)
+            schema_docs = []
+            column_docs = []
+
+            if self.vector_store:
+                schema_docs = self.vector_store.retrieve_schema_context(
+                    query, datasource_id, k=8
+                )
+                if schema_docs:
+                    for doc in schema_docs:
+                        table = doc.metadata.get("table")
+                        if table:
+                            tables[table].update([])
+                else:
+                    column_docs = self.vector_store.retrieve_column_candidates(
+                        query, datasource_id, k=8
+                    )
+                    for doc in column_docs:
+                        table = doc.metadata.get("table")
+                        column = doc.metadata.get("column")
+                        if not table:
+                            continue
+                        tables[table].update([])
+                        if column:
+                            tables[table].add(column)
+
+            planning_docs = []
+            if self.vector_store and tables:
+                planning_docs = self.vector_store.retrieve_planning_context(
+                    query, datasource_id, list(tables.keys()), k=12
+                )
+
+            for doc in planning_docs:
+                doc_type = doc.metadata.get("type")
+                if doc_type == "schema.column":
+                    table = doc.metadata.get("table")
+                    column = doc.metadata.get("column")
+                    if table and column:
+                        tables[table].add(column)
+
+                if doc_type == "schema.relationship":
+                    from_table = doc.metadata.get("from_table")
+                    to_table = doc.metadata.get("to_table")
+                    if from_table:
+                        tables[from_table].update(doc.metadata.get("from_columns"))
+                    if to_table:
+                        tables[to_table].update(doc.metadata.get("to_columns"))
+
+            if not tables:
+                snapshot = self._resolve_snapshot(datasource_id, schema_version)
+                relevant_tables = self._build_tables_from_snapshot(
+                    snapshot,
+                    resolved_tables=None,
+                    schema_version=schema_version,
+                )
+                return {
+                    "relevant_tables": relevant_tables,
+                    "reasoning": [
+                        {
+                            "node": self.node_name,
+                            "content": "Vector retrieval produced no candidates. Using full schema snapshot.",
+                            "type": "warning",
+                        }
+                    ],
+                    "warnings": [
+                        {
+                            "node": self.node_name,
+                            "content": "Vector retrieval produced no candidates. Using full schema snapshot.",
+                        }
+                    ],
+                }
+
+            snapshot = self._resolve_snapshot(datasource_id, schema_version)
+            relevant_tables = self._build_tables_from_snapshot(
+                snapshot,
+                resolved_tables=tables,
+                schema_version=schema_version,
+            )
+
+
+
+            return {
+                "relevant_tables": relevant_tables,
+                "reasoning": [
+                    {
+                        "node": self.node_name,
+                        "content": (
+                            f"Retrieved {len(relevant_tables)} tables "
+                            f"with {sum(len(t.columns) for t in relevant_tables)} columns."
+                        ),
+                    }
+                ],
+            }
+        except Exception as exc:
+            logger.error(f"Schema retrieval failed: {exc}")
+            return {
+                "relevant_tables": [],
+                "reasoning": [
+                    {
+                        "node": self.node_name,
+                        "content": f"Schema retrieval failed: {exc}",
+                        "type": "error",
+                    }
+                ],
+            }

nl2sql/pipeline/nodes/schema_retriever/schema.py

@@ -0,0 +1,27 @@
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+
+class Column(BaseModel):
+    """Lightweight column schema for routing/planning."""
+
+    name: str
+    type: Optional[str] = None
+    stats: Optional[Dict[str, Any]] = None
+    description: Optional[str] = None
+
+    model_config = ConfigDict(extra="allow")
+
+
+class Table(BaseModel):
+    """Lightweight table schema for routing/planning."""
+
+    name: str
+    columns: List[Column] = Field(default_factory=list)
+    description: Optional[str] = None
+    primary_key: Optional[List[str]] = None
+    foreign_keys: Optional[Dict[str, List[str]]] = None
+
+    model_config = ConfigDict(extra="allow")

nl2sql/pipeline/nodes/validator/__init__.py

@@ -0,0 +1,7 @@
+from .node import LogicalValidatorNode
+from .schemas import LogicalValidatorResponse
+
+__all__ = [
+    "LogicalValidatorNode",
+    "LogicalValidatorResponse",
+]