Skip to content

Extending the Planner

The planner generates a structured AST (PlanModel) which is validated and compiled into SQL. Extending planner logic typically requires changes to schema models, prompts, validator, and generator.

1. Update AST schemas

Modify PlanModel or related AST models in:

  • packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/schemas.py

Keep these rules in mind:

  • Models are strict (extra="forbid").
  • Expression kinds must pass model_post_init validation.

2. Update planner prompt

Update the prompt and examples in:

  • packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/prompts.py

PLANNER_PROMPT is a chat template of two messages. PLANNER_SYSTEM_PROMPT holds what is the same for every question (instructions, PLANNER_EXAMPLES, the schema) and PLANNER_HUMAN_PROMPT holds what changes per call (expected schema, feedback, the question). Keep per-question text out of the system message, or the provider's prompt cache stops hitting. See AST Planner: prompt layout.

Ensure structured output still matches PlanModel.

3. Update validator logic

The logical validator enforces:

  • query type (READ-only)
  • table/column existence
  • alias uniqueness
  • join relationships
  • expected schema alignment

Update LogicalValidatorNode to recognize any new AST constructs.

4. Update SQL generator

GeneratorNode and SqlVisitor translate the AST into sqlglot expressions:

  • Extend SqlVisitor to handle new expression kinds or operators.
  • Update GeneratorNode._generate_sql() if new clauses are added.

Source references

  • AST schemas: packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/schemas.py
  • Planner node: packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/node.py
  • Logical validator: packages/nl2sql/src/nl2sql/pipeline/nodes/validator/node.py
  • SQL generator: packages/nl2sql/src/nl2sql/pipeline/nodes/generator/node.py