From 56e6ad54b241241f3b8a9763146fd169b50b9e13 Mon Sep 17 00:00:00 2001 From: Nadeem Khan Date: Wed, 23 Sep 2026 15:11:37 -0500 Subject: [PATCH] refactor(pipeline): build the execution DAG in the decomposer instead of its own node `global_planner` read nothing but `state.decomposer_response`, so the DAG it built was a pure function of the decomposer's output: a format conversion with a graph node, a response model, a state field and a package around it. It is now `decomposer/dag.py`, a function the decomposer calls, and `GraphState` carries the `ExecutionDAG` itself rather than a one-field wrapper around it. Most of what the node wrote was never read, and is gone with it: * `LogicalNode.output_schema` (and `RelationSchema`/`ColumnSpec` with it) -- the only reader was the planner copying a scan's schema onto the combine node; * `ExecutionDAG.dag_id` and `.content_hash`, and the hash that produced them -- no readers anywhere; * all six scan-node attribute keys. Five had no reader; the sixth, `datasource_id`, was read on a branch that cannot be taken, because a scan node's `node_id` *is* its sub-query's id and so is always found in the router's `sub_query_map` on the line above. The router now looks the sub-query up and uses its datasource, with no fallback to invent. `layer_router` is untouched: its empty body is the fan-in barrier for parallel sub-queries, and `decomposer -> layer_router` replaces `decomposer -> global_planner -> layer_router`. The import cycle the review predicted did appear. `graph_utils` imported `ExecutionDAG` through `global_planner.schemas`, which loaded the whole `pipeline.nodes` package first and so happened to import `plan_cache` in the one order that works. Importing `nl2sql.execution.dag` directly exposed the real cycle: `plan_cache` needs `PlanModel` from `ast_planner.schemas`, and `ast_planner.node` needs `PlanCache`, while both packages' `__init__` eagerly re-exported the node class -- so importing the schemas ran the node. Fixed at the root: neither `__init__` re-exports a node class any more (nothing imported those re-exports; `sql_agent.py` now names `.node`), and a test imports `nl2sql.pipeline.plan_cache` on its own to keep it that way. One behaviour is preserved deliberately. A decomposition can be well-formed and still not describe a runnable graph -- two sub-queries the model wrote identically get the same content-addressed id -- so the DAG build keeps its own `PLANNER_FAILED` path, now on the decomposer: the decomposition is returned, no `execution_dag` is, and the layer router ends the run with that error as its cause. `test_ops_after_a_real_combine_stay_post_combine` fed exactly that case and never noticed, because nothing downstream of the decomposer ran in it. `nl2sql/pipeline/steps.py` loses the `global_planner` step and the decomposer's sentence gains what it said, so the Pipeline page's drift test agrees with the graph again. The trace recorder's `NODE_INPUTS` drops `global_planner` and the aggregator now reads `execution_dag`. --- docs/architecture/determinism.md | 6 +- docs/architecture/failure_recovery.md | 2 +- docs/architecture/graph_state.md | 24 ++- docs/architecture/invariants.md | 4 +- docs/architecture/nodes/decomposer_node.md | 7 +- .../nodes/engine_aggregator_node.md | 2 +- .../architecture/nodes/global_planner_node.md | 167 ---------------- docs/architecture/nodes/index.md | 1 - docs/architecture/overview.md | 9 +- docs/architecture/pipeline.md | 9 +- .../subgraphs/main_pipeline_graph.md | 27 ++- docs/index.md | 3 +- mkdocs.yml | 1 - packages/nl2sql/src/nl2sql/execution/dag.py | 22 +-- packages/nl2sql/src/nl2sql/pipeline/graph.py | 8 +- .../nl2sql/src/nl2sql/pipeline/graph_utils.py | 2 +- .../src/nl2sql/pipeline/nodes/__init__.py | 33 +--- .../nl2sql/pipeline/nodes/aggregator/node.py | 9 +- .../pipeline/nodes/ast_planner/__init__.py | 11 +- .../nl2sql/pipeline/nodes/decomposer/dag.py | 121 ++++++++++++ .../nl2sql/pipeline/nodes/decomposer/node.py | 28 +++ .../pipeline/nodes/global_planner/__init__.py | 4 - .../pipeline/nodes/global_planner/node.py | 187 ------------------ .../pipeline/nodes/global_planner/schemas.py | 17 -- packages/nl2sql/src/nl2sql/pipeline/routes.py | 20 +- packages/nl2sql/src/nl2sql/pipeline/state.py | 6 +- packages/nl2sql/src/nl2sql/pipeline/steps.py | 8 +- .../nl2sql/pipeline/subgraphs/sql_agent.py | 2 +- .../nl2sql/src/nl2sql/tracing/recorder.py | 3 +- .../tests/architecture/test_boundaries.py | 11 ++ .../nl2sql/tests/e2e/test_trace_fake_llm.py | 2 +- .../integration/test_aggregator_real_data.py | 12 +- .../test_answer_synthesizer_real_data.py | 12 +- .../tests/unit/test_boundary_tidy_ups.py | 8 +- ...obal_planner.py => test_decomposer_dag.py} | 106 +++++----- .../nl2sql/tests/unit/test_graph_layers.py | 12 +- .../nl2sql/tests/unit/test_node_aggregator.py | 51 ++--- .../nl2sql/tests/unit/test_pipeline_steps.py | 2 +- .../tests/unit/test_playground_pipeline.py | 2 +- .../tests/unit/test_scan_router_execute.py | 4 +- .../tests/unit/test_subgraph_resolution.py | 7 +- 41 files changed, 339 insertions(+), 633 deletions(-) delete mode 100644 docs/architecture/nodes/global_planner_node.md create mode 100644 packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py delete mode 100644 packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/__init__.py delete mode 100644 packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py delete mode 100644 packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/schemas.py rename packages/nl2sql/tests/unit/{test_node_global_planner.py => test_decomposer_dag.py} (76%) diff --git a/docs/architecture/determinism.md b/docs/architecture/determinism.md index 1903ec79..b388e9e9 100644 --- a/docs/architecture/determinism.md +++ b/docs/architecture/determinism.md @@ -7,7 +7,7 @@ answer". **Model output is not reproducible.** What is stable is the shape of a run: - **Stable sub-query and DAG ids.** Both are SHA-256 content hashes over - canonical JSON, not counters or UUIDs ([`decomposer/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py), [`global_planner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py)). The same decomposition always yields the same ids. + canonical JSON, not counters or UUIDs ([`decomposer/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py), [`decomposer/dag.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py)). The same decomposition always yields the same ids. - **Sorted layer order.** The topological sort sorts each ready set and each dependent set, so the execution layers of a given DAG are fixed ([`execution/dag.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/execution/dag.py)). @@ -101,7 +101,7 @@ gives every result a total row order. - Non-deterministic: User query interpretation depends on LLMs in the decomposer, planner, and refiner nodes. These nodes set nothing themselves; `temperature` (default `0.0`, omitted when configured as `null`) and `seed=42` are applied once, where the client is built in [`llm/registry.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/llm/registry.py), and neither makes the output reproducible ([`pipeline/nodes/decomposer/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py), [`pipeline/nodes/ast_planner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/node.py), [`pipeline/nodes/refiner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/refiner/node.py)). ### Planner and DAG Construction -- Deterministic: Global planner sorts nodes and edges by IDs and roles before constructing the DAG, then hashes a sorted JSON payload to produce a stable `dag_id` for a given logical plan ([`pipeline/nodes/global_planner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py)). +- Deterministic: The decomposer sorts nodes and edges by IDs and roles before constructing the DAG, so the same decomposition always gives the same layers ([`pipeline/nodes/decomposer/dag.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py)). - Deterministic: DAG layers are computed with a topological sort that sorts ready nodes and dependents, yielding stable layer ordering given the same node/edge sets ([`execution/dag.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/execution/dag.py)). - Non-deterministic: The AST planner is LLM-driven; the PlanModel content is not stabilized inside the node ([`pipeline/nodes/ast_planner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/node.py)). - Deterministic (conditional): Once a sub-query's plan has validated and executed, the plan cache pins it: the same normalised intent, datasource and schema version reuse that plan without a planner call, and it is validated again on every use (see [The plan cache](#the-plan-cache-determinism-from-the-architecture)). @@ -146,7 +146,7 @@ gives every result a total row order. - Potentially non-deterministic: GraphState merges dict fields with a last-write-wins reducer and concatenates lists; in parallel branches, merge order is not constrained in this code ([`pipeline/state.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/state.py)). ### Hashing/Fingerprinting -- Deterministic: Stable hashing is consistently performed with sorted JSON and fixed separators for sub-query IDs, DAG hashes, schema fingerprints, and artifact content hashes ([`pipeline/nodes/decomposer/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py), [`pipeline/nodes/global_planner/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py), [`schema/protocol.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/schema/protocol.py), [`execution/artifacts/store.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/execution/artifacts/store.py)). +- Deterministic: Stable hashing is consistently performed with sorted JSON and fixed separators for sub-query IDs, schema fingerprints, and artifact content hashes ([`pipeline/nodes/decomposer/node.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py), [`schema/protocol.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/schema/protocol.py), [`execution/artifacts/store.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/execution/artifacts/store.py)). ### Time-Based Logic and Runtime Controls - Non-deterministic: Schema versions and artifact creation timestamps are derived from wall-clock time ([`schema/in_memory_store.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/schema/in_memory_store.py), [`schema/sqlite_store.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/schema/sqlite_store.py), [`execution/artifacts/store.py`](https://github.com/nadeem4/nl2sql/blob/main/packages/nl2sql/src/nl2sql/execution/artifacts/store.py)). diff --git a/docs/architecture/failure_recovery.md b/docs/architecture/failure_recovery.md index d820fea0..8b6909d8 100644 --- a/docs/architecture/failure_recovery.md +++ b/docs/architecture/failure_recovery.md @@ -21,7 +21,7 @@ Failure in this system is represented as structured `PipelineError` objects accu ### Planning - Decomposer LLM failures return `ORCHESTRATOR_CRASH` (critical) and empty responses. - AST planner LLM failures return `PLANNING_FAILURE` and a `None` plan. -- Global planner failures return `PLANNER_FAILED` and no `global_planner_response`; the layer router ends the run. +- A decomposition that cannot be turned into an execution DAG returns `PLANNER_FAILED` and no `execution_dag`; the layer router ends the run. ### Validation - Logical validation returns structured errors for missing tables, columns, invalid plan structure, or security violations. diff --git a/docs/architecture/graph_state.md b/docs/architecture/graph_state.md index e4f4b661..bb52ac32 100644 --- a/docs/architecture/graph_state.md +++ b/docs/architecture/graph_state.md @@ -17,7 +17,7 @@ Fields (exact names and types from code): - `datasource_id: Optional[str]` - `datasource_resolver_response: Optional[DatasourceResolverResponse]` - `decomposer_response: Optional[DecomposerResponse]` -- `global_planner_response: Optional[GlobalPlannerResponse]` +- `execution_dag: Optional[ExecutionDAG]` - `aggregator_response: Optional[AggregatorResponse]` - `answer_synthesizer_response: Optional[AnswerSynthesizerResponse]` - `artifact_refs: Annotated[Dict[str, ArtifactRef], update_results]` @@ -105,11 +105,11 @@ The lifecycle below lists creation, mutation, reads, and resets based strictly o ### `decomposer_response` - Creation/mutation: returned by `DecomposerNode` as `decomposer_response`. -- Read points: `GlobalPlannerNode`, `build_scan_payload`, and `wrap_subgraph` (to find `sub_query` by id). +- Read points: `build_scan_layer_router`, `build_scan_payload`, and `wrap_subgraph` (to find `sub_query` by id). - Reset: none in code. -### `global_planner_response` -- Creation/mutation: returned by `GlobalPlannerNode` as `global_planner_response`. +### `execution_dag` +- Creation/mutation: returned by `DecomposerNode`, which builds it from its own response (`decomposer/dag.py`). - Read points: `build_scan_layer_router` and `EngineAggregatorNode`. - Reset: none in code. @@ -165,7 +165,7 @@ The lifecycle below lists creation, mutation, reads, and resets based strictly o Ownership is defined by which node returns updates for a field: - `datasource_resolver_response`: `DatasourceResolverNode` - `decomposer_response`: `DecomposerNode` -- `global_planner_response`: `GlobalPlannerNode` +- `execution_dag`: `DecomposerNode` - `aggregator_response`: `EngineAggregatorNode` - `answer_synthesizer_response`: `AnswerSynthesizerNode` - `artifact_refs`: `wrap_subgraph` (subgraph wrapper in `graph_utils.py`) @@ -187,16 +187,15 @@ Step-by-step execution flow as defined in `build_graph` and routing: 1. `run_with_graph` constructs `GraphState` with `user_query`, `user_context`, and optional `datasource_id`, then calls `graph.invoke(initial_state.model_dump())`. 2. `DatasourceResolverNode` runs first and populates `datasource_resolver_response`, `reasoning`, and `errors`. 3. `resolver_route` decides whether to continue based on `datasource_resolver_response`. -4. `DecomposerNode` produces `decomposer_response` and reasoning. -5. `GlobalPlannerNode` produces `global_planner_response` (including the `ExecutionDAG`). -6. `build_scan_layer_router` emits `Send` branches using `build_scan_payload` for each pending scan node. The payload contains `subgraph_id`, `subgraph_name`, `trace_id`, `user_context`, `decomposer_response`, and `datasource_resolver_response`. -7. Each subgraph is wrapped by `wrap_subgraph`, which: +4. `DecomposerNode` produces `decomposer_response`, the `execution_dag` built from it, and reasoning. +5. `build_scan_layer_router` emits `Send` branches using `build_scan_payload` for each pending scan node. The payload contains `subgraph_id`, `subgraph_name`, `trace_id`, `user_context`, `decomposer_response`, and `datasource_resolver_response`. +6. Each subgraph is wrapped by `wrap_subgraph`, which: - Builds a `SubgraphExecutionState` using the payload and a `sub_query` resolved from `decomposer_response`. - Invokes the subgraph and validates the result into `SubgraphExecutionState`. - Returns updates for `artifact_refs`, `subgraph_outputs`, `errors`, and `reasoning`. -8. The router checks `artifact_refs` to decide which scan nodes are still pending. When none remain, it routes to `aggregator`. -9. `EngineAggregatorNode` consumes `global_planner_response` and `artifact_refs` to produce `aggregator_response`. -10. `AnswerSynthesizerNode` consumes `aggregator_response` and `decomposer_response` to produce `answer_synthesizer_response`. +7. The router checks `artifact_refs` to decide which scan nodes are still pending. When none remain, it routes to `aggregator`. +8. `EngineAggregatorNode` consumes `execution_dag` and `artifact_refs` to produce `aggregator_response`. +9. `AnswerSynthesizerNode` consumes `aggregator_response` and `decomposer_response` to produce `answer_synthesizer_response`. Subgraph internal flow uses `SubgraphExecutionState` and is defined in `build_sql_agent_graph`: `schema_retriever` -> `ast_planner` -> `logical_validator` -> `generator` -> `executor`, with a retry loop via `retry_handler` and `refiner`. @@ -274,7 +273,6 @@ Subgraph state and execution: Mutators / consumers: - `packages/nl2sql/src/nl2sql/pipeline/nodes/datasource_resolver/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py` -- `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/answer_synthesizer/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/schema_retriever/node.py` diff --git a/docs/architecture/invariants.md b/docs/architecture/invariants.md index e2645409..d031221b 100644 --- a/docs/architecture/invariants.md +++ b/docs/architecture/invariants.md @@ -224,7 +224,7 @@ Prevents invalid execution requests and ensures capability compatibility. Post-combine ops must target known combine groups, all edges must reference existing nodes, and the DAG must be acyclic. ### Enforcement Points -- `GlobalPlannerNode.__call__()` in `nl2sql.pipeline.nodes.global_planner.node` +- `build_execution_dag()` in `nl2sql.pipeline.nodes.decomposer.dag`, called by `DecomposerNode.__call__()` - `ExecutionDAG._layered_toposort()` in `nl2sql.execution.dag` ### Failure Behavior @@ -417,7 +417,7 @@ The validator resolves columns against a throw-away query in which every table i - `packages/nl2sql/src/nl2sql/pipeline/nodes/generator/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/executor/node.py` - `packages/nl2sql/src/nl2sql/execution/executor/sql_executor.py` -- `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` +- `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py` - `packages/nl2sql/src/nl2sql/execution/dag.py` - `packages/nl2sql/src/nl2sql/aggregation/aggregator.py` - `packages/nl2sql/src/nl2sql/context.py` diff --git a/docs/architecture/nodes/decomposer_node.md b/docs/architecture/nodes/decomposer_node.md index aa939cd8..444e003c 100644 --- a/docs/architecture/nodes/decomposer_node.md +++ b/docs/architecture/nodes/decomposer_node.md @@ -4,7 +4,7 @@ - Decomposes the user query into datasource‑scoped sub‑queries and combination operations. - Produces deterministic IDs for sub‑queries and post‑combine ops. -- Sits after `DatasourceResolverNode` and before `GlobalPlannerNode`. +- Sits after `DatasourceResolverNode` and before the layer router. - Class: `DecomposerNode` - Source: `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py` @@ -16,6 +16,7 @@ - Filter sub‑queries by resolved/allowed/unsupported datasources. - Stabilize IDs using a hash of sub‑query content. - Normalize and sort combine groups and post‑combine operations. +- Build the `ExecutionDAG` the layer router walks and the aggregator runs (`decomposer/dag.py`). The DAG is a pure function of the response above, so it is code here rather than a node of its own. --- @@ -25,14 +26,14 @@ Upstream: - `DatasourceResolverNode` Downstream: -- `GlobalPlannerNode` +- `layer_router` Trigger conditions: - Executed only when `resolver_route` returns `continue`. ```mermaid flowchart LR - Resolver[DatasourceResolverNode] --> Decomposer[DecomposerNode] --> Planner[GlobalPlannerNode] + Resolver[DatasourceResolverNode] --> Decomposer[DecomposerNode] --> Router[layer_router] ``` --- diff --git a/docs/architecture/nodes/engine_aggregator_node.md b/docs/architecture/nodes/engine_aggregator_node.md index fece36cf..5a096649 100644 --- a/docs/architecture/nodes/engine_aggregator_node.md +++ b/docs/architecture/nodes/engine_aggregator_node.md @@ -40,7 +40,7 @@ flowchart LR From `GraphState`: -- `global_planner_response.execution_dag` (required) +- `execution_dag` (required) - `artifact_refs` (required) Validation performed: diff --git a/docs/architecture/nodes/global_planner_node.md b/docs/architecture/nodes/global_planner_node.md deleted file mode 100644 index d90d7a74..00000000 --- a/docs/architecture/nodes/global_planner_node.md +++ /dev/null @@ -1,167 +0,0 @@ -# GlobalPlannerNode - -## Overview - -- Builds a deterministic `ExecutionDAG` from decomposer output. -- Creates scan, combine, and post‑combine nodes with explicit edges. -- Sits between `DecomposerNode` and `layer_router`. -- Class: `GlobalPlannerNode` -- Source: `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` - ---- - -## Responsibilities - -- Translate `SubQuery` objects into `scan` nodes. -- Construct `combine` nodes for join/union/compare groups. -- Construct `post_*` nodes for filter/aggregate/project/sort/limit operations. -- Create edges between nodes and validate references. -- Compute deterministic `ExecutionDAG` ordering, content hash, and dag_id. - ---- - -## Position in Execution Graph - -Upstream: -- `DecomposerNode` - -Downstream: -- `layer_router` - -Trigger conditions: -- Always executed after successful decomposition. - -```mermaid -flowchart LR - Decomposer[DecomposerNode] --> Planner[GlobalPlannerNode] --> Router[layer_router] -``` - ---- - -## Inputs - -From `GraphState`: - -- `decomposer_response.sub_queries` -- `decomposer_response.combine_groups` -- `decomposer_response.post_combine_ops` - -Validation performed: - -- Validates edges reference existing node IDs. -- Fails if post‑combine ops reference unknown combine groups. - ---- - -## Outputs - -Mutations to `GraphState`: - -- `global_planner_response` (`GlobalPlannerResponse` with `ExecutionDAG`) -- `reasoning` on success -- `errors` on failure - -Side effects: - -- None beyond in‑memory DAG construction. - ---- - -## Internal Flow (Step-by-Step) - -1. Extract `sub_queries`, `combine_groups`, `post_combine_ops`. -2. For each sub‑query, build a `scan` `LogicalNode` with `RelationSchema` from expected schema. -3. For each combine group, create a `combine` node and edges from inputs. -4. For each post‑combine op, create a `post_*` node and edge from its target combine. -5. Validate that all edges reference known nodes. -6. Build `ExecutionDAG` with sorted nodes/edges. -7. Compute `content_hash` and `dag_id`. -8. Return `GlobalPlannerResponse`. -9. On exception, emit a `PLANNER_FAILED` error carrying the cause, and no - `GlobalPlannerResponse`. The layer router finds no DAG and ends the run, so - this error is what the caller sees. - ---- - -## Contracts & Interfaces - -Implements a LangGraph node callable: - -``` -def __call__(self, state: GraphState) -> Dict[str, Any] -``` - -Key contracts: - -- `ExecutionDAG` -- `LogicalNode`, `LogicalEdge` -- `RelationSchema`, `ColumnSpec` -- `GlobalPlannerResponse` - ---- - -## Determinism Guarantees - -- Nodes and edges are sorted before DAG creation. -- `ExecutionDAG._layered_toposort()` produces stable layer ordering. -- `content_hash` is a deterministic hash of nodes/edges/version. - ---- - -## Error Handling - -Emits `PipelineError` with: - -- `PLANNER_FAILED` on exceptions - -Logs failures via `logger.error`. - ---- - -## Retry + Idempotency - -- No internal retry logic. -- Deterministic for a given decomposer response. - ---- - -## Performance Characteristics - -- Pure in‑memory graph construction and hashing. -- Complexity scales with number of sub‑queries and combine ops. - ---- - -## Observability - -- Logger: `global_planner` -- Adds a reasoning entry on success or failure. - ---- - -## Configuration - -- No direct settings read in this node. - ---- - -## Extension Points - -- Extend DAG semantics by modifying `kind_map` and `attributes` mapping. -- Replace node in `build_graph()` for custom DAG construction. - ---- - -## Known Limitations - -- Assumes `expected_schema` is present to build `RelationSchema`. -- No schema validation at DAG construction time. -- No support for custom node kinds without code changes. - ---- - -## Related Code - -- `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` -- `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/schemas.py` -- `packages/nl2sql/src/nl2sql/execution/dag.py` (the DAG models, shared with the aggregation service) diff --git a/docs/architecture/nodes/index.md b/docs/architecture/nodes/index.md index f7877fba..f03e5dd4 100644 --- a/docs/architecture/nodes/index.md +++ b/docs/architecture/nodes/index.md @@ -8,7 +8,6 @@ For pipeline-level wiring, see `../pipeline.md`. For subgraph wiring and lifecyc - [DatasourceResolverNode](datasource_resolver_node.md) - [DecomposerNode](decomposer_node.md) -- [GlobalPlannerNode](global_planner_node.md) - [EngineAggregatorNode](engine_aggregator_node.md) - [AnswerSynthesizerNode](answer_synthesizer_node.md) diff --git a/docs/architecture/overview.md b/docs/architecture/overview.md index fec7209c..7fef0711 100644 --- a/docs/architecture/overview.md +++ b/docs/architecture/overview.md @@ -17,8 +17,7 @@ flowchart TD Runtime --> Graph[build_graph()] Graph --> Resolver[DatasourceResolverNode] Resolver --> Decomposer[DecomposerNode] - Decomposer --> Planner[GlobalPlannerNode] - Planner --> Router[Scan Layer Router] + Decomposer --> Router[Scan Layer Router] Router --> Subgraph[SQL Agent Subgraph] Subgraph --> Router Router --> Aggregator[EngineAggregatorNode] @@ -61,7 +60,7 @@ flowchart LR ## Major subsystems (and responsibilities) -- **Planner / Decomposer**: `DecomposerNode` produces stable, semantically-scoped sub-queries; `GlobalPlannerNode` produces a deterministic `ExecutionDAG`. +- **Planner / Decomposer**: `DecomposerNode` produces stable, semantically-scoped sub-queries and, from them, a deterministic `ExecutionDAG`. - **Schema Store**: `SchemaStore` persists versioned schema snapshots with fingerprints. - **Chunking + Retrieval**: `SchemaChunkBuilder` produces typed chunks; `VectorStore` provides staged retrieval for routing and planning context. - **Validation layer**: `LogicalValidatorNode` enforces schema correctness and RBAC. @@ -78,7 +77,6 @@ sequenceDiagram participant Runtime as run_with_graph participant Resolver as DatasourceResolverNode participant Decomposer as DecomposerNode - participant Planner as GlobalPlannerNode participant Router as Scan Layer Router participant Subgraph as SQL Agent Subgraph participant Agg as EngineAggregatorNode @@ -87,8 +85,7 @@ sequenceDiagram User->>Runtime: user_query Runtime->>Resolver: GraphState Resolver->>Decomposer: resolved datasources - Decomposer->>Planner: sub_queries + combine groups - Planner->>Router: ExecutionDAG + Decomposer->>Router: ExecutionDAG over the sub-queries Router->>Subgraph: Send(sub_query) Subgraph-->>Router: ArtifactRef + diagnostics Router->>Agg: all scan artifacts diff --git a/docs/architecture/pipeline.md b/docs/architecture/pipeline.md index fd0de638..2303c07c 100644 --- a/docs/architecture/pipeline.md +++ b/docs/architecture/pipeline.md @@ -7,8 +7,7 @@ This document consolidates the LangGraph pipeline architecture as implemented in ```mermaid flowchart TD Resolver[DatasourceResolverNode] --> Decomposer[DecomposerNode] - Decomposer --> Planner[GlobalPlannerNode] - Planner --> Router[layer_router] + Decomposer --> Router[layer_router] Router --> SqlAgent[SQL Agent Subgraph] SqlAgent --> Router Router --> Aggregator[EngineAggregatorNode] @@ -19,7 +18,6 @@ flowchart TD - [DatasourceResolverNode](nodes/datasource_resolver_node.md) - [DecomposerNode](nodes/decomposer_node.md) -- [GlobalPlannerNode](nodes/global_planner_node.md) - `layer_router` (routing function, not a node class) - [EngineAggregatorNode](nodes/engine_aggregator_node.md) - [AnswerSynthesizerNode](nodes/answer_synthesizer_node.md) @@ -52,12 +50,11 @@ graph TD - `kind`: `scan`, `combine`, `post_filter`, `post_aggregate`, `post_project`, `post_sort`, `post_limit` - `inputs`: upstream node IDs -- `output_schema`: `RelationSchema` (column specs) -- `attributes`: operation-specific metadata (e.g., join keys, filters, metrics) +- `attributes`: operation-specific metadata for a combine or post-combine node (join keys, filters, metrics). A `scan` node carries none: its ID is its sub-query's ID, so everything about it is one lookup away in `decomposer_response`. ## Routing and subgraph selection -`build_scan_layer_router()` selects the next scan layer with missing artifacts and dispatches subgraphs in parallel using `Send()`. Scan node IDs are matched against `SubQuery.id` for datasource routing; otherwise datasource ID is resolved from node attributes. +`build_scan_layer_router()` selects the next scan layer with missing artifacts and dispatches subgraphs in parallel using `Send()`. A scan node's ID is its sub-query's ID, so the datasource comes from that `SubQuery`. Subgraph selection is capability-based via `resolve_subgraph()`: diff --git a/docs/architecture/subgraphs/main_pipeline_graph.md b/docs/architecture/subgraphs/main_pipeline_graph.md index 3f71b59e..79566214 100644 --- a/docs/architecture/subgraphs/main_pipeline_graph.md +++ b/docs/architecture/subgraphs/main_pipeline_graph.md @@ -12,7 +12,7 @@ Source file path: `packages/nl2sql/src/nl2sql/pipeline/graph.py` ## Boundary Definition -This graph is the top-level orchestrator. It encapsulates datasource resolution, decomposition, global planning, routing to subgraphs, aggregation, and answer synthesis. It does NOT implement subgraph internals (SQL planning/execution), per-node business logic, or adapter-specific execution; those are delegated to nodes and registered subgraphs. +This graph is the top-level orchestrator. It encapsulates datasource resolution, decomposition (which builds the execution DAG), routing to subgraphs, aggregation, and answer synthesis. It does NOT implement subgraph internals (SQL planning/execution), per-node business logic, or adapter-specific execution; those are delegated to nodes and registered subgraphs. --- @@ -51,8 +51,7 @@ Partial completion behavior: Execution order (nominal path): - `datasource_resolver` — `DatasourceResolverNode` — `packages/nl2sql/src/nl2sql/pipeline/nodes/datasource_resolver/node.py` — resolve candidate datasources, RBAC, and the answerability check (an LLM call). -- `decomposer` — `DecomposerNode` — `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py` — decompose query into sub-queries. -- `global_planner` — `GlobalPlannerNode` — `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` — build execution DAG. +- `decomposer` — `DecomposerNode` — `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py` — decompose query into sub-queries, and build the `ExecutionDAG` from them (`decomposer/dag.py`). - `layer_router` — inline lambda + `routes.build_scan_layer_router` — `packages/nl2sql/src/nl2sql/pipeline/routes.py` — route each scan layer to a subgraph or aggregator. - `` — via `wrap_subgraph()` — `packages/nl2sql/src/nl2sql/pipeline/graph_utils.py` — invoke registered subgraphs per scan node. - `aggregator` — `EngineAggregatorNode` — `packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py` — execute aggregation DAG. @@ -63,8 +62,7 @@ Mermaid diagram (main pipeline only): flowchart TD datasource_resolver -->|continue| decomposer datasource_resolver -->|end| END - decomposer --> global_planner - global_planner --> layer_router + decomposer --> layer_router layer_router -->|subgraph| subgraph_exec subgraph_exec --> layer_router layer_router -->|aggregator| aggregator @@ -88,17 +86,16 @@ Field ownership, reducers, and lifecycle are defined in `../graph_state.md`. 2. `resolver_route` decides: - `continue` if allowed datasources exist and the question was judged answerable. - `end` if none exist, the question is not answerable, or the response is missing. -3. `decomposer` uses the LLM to produce `SubQuery` objects and combine groups. -4. `global_planner` builds a deterministic `ExecutionDAG` from sub-queries and combines. -5. `layer_router` inspects the DAG and current `artifact_refs`: +3. `decomposer` uses the LLM to produce `SubQuery` objects and combine groups, then builds a deterministic `ExecutionDAG` from them in code. The DAG is a pure function of that output, so it is not a node of its own; a DAG that cannot be built is reported as `PLANNER_FAILED` and the run ends at the router. +4. `layer_router` inspects the DAG and current `artifact_refs`: - If no DAG or layers, returns `END`. - If next scan layer is empty, routes to `aggregator` -- or to `END` when no scan produced an artifact or `execute` is False. - For each scan node, resolves a compatible subgraph and sends `build_scan_payload`. -6. Each subgraph execution returns `artifact_refs`, `subgraph_outputs`, and `errors` to `GraphState`. -7. `layer_router` is re-entered until all scan-layer nodes produce artifacts. -8. `aggregator` executes the DAG using the stored `artifact_refs`. -9. `answer_synthesizer` summarizes aggregated results into a final answer. -10. Graph reaches `END`. +5. Each subgraph execution returns `artifact_refs`, `subgraph_outputs`, and `errors` to `GraphState`. +6. `layer_router` is re-entered until all scan-layer nodes produce artifacts. +7. `aggregator` executes the DAG using the stored `artifact_refs`. +8. `answer_synthesizer` summarizes aggregated results into a final answer. +9. Graph reaches `END`. --- @@ -123,7 +120,7 @@ See `../failure_recovery.md` for retry scope and recovery behavior. ## Performance Characteristics - Blocking calls include LLM requests in `datasource_resolver`, `decomposer` and `answer_synthesizer`. -- `global_planner` and `aggregator` are CPU-bound (DAG construction and local aggregation). +- The decomposer's DAG construction and the `aggregator` are CPU-bound. - Subgraph executions are dispatched per scan layer and can run in parallel via LangGraph routing. - Overall pipeline is executed on a single-worker thread pool created per run and guarded by `settings.global_timeout_sec`. @@ -170,6 +167,6 @@ See `../failure_recovery.md` for retry scope and recovery behavior. - Node implementations: - `packages/nl2sql/src/nl2sql/pipeline/nodes/datasource_resolver/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py` - - `packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py` + - `packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py` - `packages/nl2sql/src/nl2sql/pipeline/nodes/answer_synthesizer/node.py` diff --git a/docs/index.md b/docs/index.md index 17b20aa8..1613747b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -30,8 +30,7 @@ Model output itself is not reproducible; see [Determinism](architecture/determin flowchart TD User[User Query] --> Resolver[DatasourceResolverNode] Resolver --> Decomposer[DecomposerNode] - Decomposer --> Planner[GlobalPlannerNode] - Planner --> Router[Scan Layer Router] + Decomposer --> Router[Scan Layer Router] Router --> Subgraph[SQL Agent Subgraph] Subgraph --> Router Router --> Aggregator[EngineAggregatorNode] diff --git a/mkdocs.yml b/mkdocs.yml index afc288b2..606250be 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -75,7 +75,6 @@ nav: - Engine Aggregator: architecture/nodes/engine_aggregator_node.md - Executor: architecture/nodes/executor_node.md - Generator: architecture/nodes/generator_node.md - - Global Planner: architecture/nodes/global_planner_node.md - Logical Validator: architecture/nodes/logical_validator_node.md - Refiner: architecture/nodes/refiner_node.md - Schema Retriever: architecture/nodes/schema_retriever_node.md diff --git a/packages/nl2sql/src/nl2sql/execution/dag.py b/packages/nl2sql/src/nl2sql/execution/dag.py index 482010ba..27a22091 100644 --- a/packages/nl2sql/src/nl2sql/execution/dag.py +++ b/packages/nl2sql/src/nl2sql/execution/dag.py @@ -1,4 +1,4 @@ -"""The execution DAG: the logical plan the global planner builds and the aggregation service runs. +"""The execution DAG: the logical plan the decomposer builds and the aggregation service runs. It lives here, below both, so ``nl2sql.aggregation`` never imports ``nl2sql.pipeline``. """ @@ -7,23 +7,6 @@ from pydantic import BaseModel, Field, model_validator -JsonLiteral = str | int | float | bool | None - - -class ColumnSpec(BaseModel): - name: str - dtype: Optional[str] = None - -class RelationSchema(BaseModel): - columns: List[ColumnSpec] - - @model_validator(mode="after") - def validate_unique_columns(self): - names = [c.name for c in self.columns] - if len(names) != len(set(names)): - raise ValueError(f"Duplicate columns in schema: {names}") - return self - class LogicalNode(BaseModel): node_id: str kind: Literal[ @@ -36,7 +19,6 @@ class LogicalNode(BaseModel): "post_limit", ] inputs: List[str] = Field(default_factory=list) - output_schema: RelationSchema attributes: Dict[str, Any] = Field(default_factory=dict) @@ -48,8 +30,6 @@ class LogicalEdge(BaseModel): class ExecutionDAG(BaseModel): - dag_id: Optional[str] = None - content_hash: Optional[str] = None nodes: List[LogicalNode] edges: List[LogicalEdge] layers: List[List[str]] = Field(default_factory=list) diff --git a/packages/nl2sql/src/nl2sql/pipeline/graph.py b/packages/nl2sql/src/nl2sql/pipeline/graph.py index 65e52017..d509414d 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/graph.py +++ b/packages/nl2sql/src/nl2sql/pipeline/graph.py @@ -6,7 +6,6 @@ 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 @@ -33,13 +32,11 @@ def build_graph( 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, execute=execute) 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, execute=execute), @@ -56,10 +53,11 @@ def build_graph( {"continue": "decomposer", "end": END}, ) - graph.add_edge("decomposer", "global_planner") + # The decomposer builds the execution DAG itself, so its output is all the + # router needs: there is no node between them. route_scan_layers = build_scan_layer_router(ctx, execute=execute) - graph.add_edge("global_planner", "layer_router") + graph.add_edge("decomposer", "layer_router") graph.add_conditional_edges( "layer_router", route_scan_layers, diff --git a/packages/nl2sql/src/nl2sql/pipeline/graph_utils.py b/packages/nl2sql/src/nl2sql/pipeline/graph_utils.py index ee64acb3..385f6920 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/graph_utils.py +++ b/packages/nl2sql/src/nl2sql/pipeline/graph_utils.py @@ -6,7 +6,7 @@ from nl2sql.common.errors import ErrorCode, ErrorSeverity, PipelineError from nl2sql.context import NL2SQLContext -from nl2sql.pipeline.nodes.global_planner.schemas import ExecutionDAG +from nl2sql.execution.dag import ExecutionDAG from nl2sql.pipeline.plan_cache import PlanCache from nl2sql.pipeline.state import GraphState, SubgraphExecutionState from nl2sql.pipeline.subgraphs import SubgraphOutput diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/__init__.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/__init__.py index 1ef9125b..9d48d74b 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/__init__.py +++ b/packages/nl2sql/src/nl2sql/pipeline/nodes/__init__.py @@ -1,25 +1,10 @@ -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 +"""The pipeline's nodes. Import the one you need from its own module. -__all__ = [ - "ASTPlannerNode", - "SchemaRetrieverNode", - "GeneratorNode", - "ExecutorNode", - "RefinerNode", - "DecomposerNode", - "DatasourceResolverNode", - "GlobalPlannerNode", - "LogicalValidatorNode", - "EngineAggregatorNode", - "AnswerSynthesizerNode", -] +This package deliberately re-exports nothing. It used to import every node +class, which meant that importing any leaf module under it -- say +``nodes.ast_planner.schemas`` for ``PlanModel`` -- first ran every node module +in the package. ``plan_cache`` needs exactly that one schema, and +``ast_planner.node`` needs ``plan_cache``, so the two formed a cycle that only +stayed hidden while some earlier import happened to load them in the lucky +order. Nothing in the tree imported the re-exports. +""" diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py index 4f91bc37..3a26c773 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py +++ b/packages/nl2sql/src/nl2sql/pipeline/nodes/aggregator/node.py @@ -6,7 +6,6 @@ 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 @@ -25,12 +24,8 @@ def __init__(self, ctx: NL2SQLContext): 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) + dag = state.execution_dag + terminal_results = self.service.execute(dag, state.artifact_refs) aggregator_response = AggregatorResponse( terminal_results=terminal_results, computed_artifacts={}, diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/__init__.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/__init__.py index 0c207a8a..e3c68af0 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/__init__.py +++ b/packages/nl2sql/src/nl2sql/pipeline/nodes/ast_planner/__init__.py @@ -1,4 +1,9 @@ -from .node import ASTPlannerNode -from .schemas import ASTPlannerResponse +"""The AST planner. Its node class lives in ``.node``, imported from there. -__all__ = ["ASTPlannerNode", "ASTPlannerResponse"] +The package re-exports the schemas only. ``plan_cache`` needs ``PlanModel`` +from ``.schemas``, and ``.node`` needs ``PlanCache``: re-exporting the node +here made that pair a cycle, because importing the schemas ran ``.node`` first. +""" +from .schemas import ASTPlannerResponse, PlanModel + +__all__ = ["ASTPlannerResponse", "PlanModel"] diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py new file mode 100644 index 00000000..2c2c5ad9 --- /dev/null +++ b/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/dag.py @@ -0,0 +1,121 @@ +"""The execution DAG, built from the decomposition it is a pure function of. + +The decomposer says which sub-queries answer the question and how their results +combine. Turning that into the graph the layer router walks and the aggregation +service runs is a format conversion, not a decision: it reads nothing but the +:class:`~nl2sql.pipeline.nodes.decomposer.schemas.DecomposerResponse`. It used +to be a graph node of its own (``global_planner``), which bought a node, a +response model and a state field for a function call. + +Node ids are the decomposition's own ids -- a scan node *is* its sub-query, by +id -- so nothing has to map between the two. Nodes and edges are sorted before +the DAG is built, so the same decomposition always gives the same layers. +""" +from __future__ import annotations + +from typing import Dict, List + +from nl2sql.execution.dag import ExecutionDAG, LogicalEdge, LogicalNode + +from .schemas import DecomposerResponse + +# The logical node kind each post-combine operation becomes. +_POST_OP_KINDS = { + "filter": "post_filter", + "aggregate": "post_aggregate", + "project": "post_project", + "sort": "post_sort", + "limit": "post_limit", +} + + +def build_execution_dag(response: DecomposerResponse) -> ExecutionDAG: + """The execution DAG for a decomposition. + + Args: + response: The decomposer's own output. ``DecomposerResponse`` has + already checked that every combine group names a real sub-query and + every post-combine op a real group, so the references here hold. + + Returns: + The DAG, with its layers populated by ``ExecutionDAG``. + + Raises: + ValueError: If the graph is not acyclic, or an edge names a node that + is not in it. Neither is reachable from a valid response; they are + the last line of defence before the router walks the layers. + """ + nodes: List[LogicalNode] = [] + edges: List[LogicalEdge] = [] + + # A scan node carries no attributes: its id is its sub-query's id, so + # everything about it is one lookup away in the decomposer response. + for sq in response.sub_queries: + nodes.append(LogicalNode(node_id=sq.id, kind="scan", inputs=[])) + + combine_node_ids: Dict[str, str] = {} + for cg in response.combine_groups: + combine_id = f"combine_{cg.group_id}" + combine_node_ids[cg.group_id] = combine_id + nodes.append( + LogicalNode( + node_id=combine_id, + kind="combine", + inputs=[i.subquery_id for i in cg.inputs], + attributes={ + "operation": cg.operation, + "group_id": cg.group_id, + "inputs": [i.model_dump() for i in cg.inputs], + "join_keys": [jk.model_dump() for jk in cg.join_keys], + }, + ) + ) + for inp in cg.inputs: + edges.append( + LogicalEdge( + edge_id=f"edge_{inp.subquery_id}_{combine_id}", + from_id=inp.subquery_id, + to_id=combine_id, + role=inp.role, + ) + ) + + for op in response.post_combine_ops: + target_combine_id = combine_node_ids.get(op.target_group_id) + if not target_combine_id: + raise ValueError(f"PostCombineOp references unknown combine group: {op.target_group_id}") + nodes.append( + LogicalNode( + node_id=op.op_id, + kind=_POST_OP_KINDS[op.operation], + inputs=[target_combine_id], + attributes={ + "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, + }, + ) + ) + edges.append( + LogicalEdge( + edge_id=f"edge_{target_combine_id}_{op.op_id}", + from_id=target_combine_id, + to_id=op.op_id, + ) + ) + + node_ids = {n.node_id for n in nodes} + for edge in edges: + if edge.from_id not in node_ids or edge.to_id not in node_ids: + raise ValueError(f"Edge references unknown node: {edge}") + + return ExecutionDAG( + nodes=sorted(nodes, key=lambda n: n.node_id), + edges=sorted(edges, key=lambda e: (e.from_id, e.to_id, e.role or "")), + ) diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py index 8ad45ac8..e202d812 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py +++ b/packages/nl2sql/src/nl2sql/pipeline/nodes/decomposer/node.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: from nl2sql.pipeline.state import GraphState +from .dag import build_execution_dag from .schemas import DecomposerResponse, SubQuery, UnmappedSubQuery, PostCombineOp from .prompts import DECOMPOSER_PROMPT from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode @@ -241,8 +242,35 @@ def __call__(self, state: GraphState) -> Dict[str, Any]: unmapped_subqueries=unmapped, ) + # The execution DAG is a pure function of this response, so it is + # built here rather than in a node of its own. It has its own + # failure: a decomposition can be well-formed and still not make a + # runnable graph (two identical sub-queries share one stable id, so + # the graph they describe has a node with two of everything). The + # response is returned either way -- the layer router ends a run + # with no DAG, leaving this error as its cause. + try: + execution_dag = build_execution_dag(response) + except Exception as exc: + logger.error(f"Execution DAG generation failed: {exc}") + return { + "decomposer_response": response, + "reasoning": [{"node": self.node_name, + "content": f"Execution DAG generation failed: {exc}", + "type": "error"}], + "errors": [ + PipelineError( + node=self.node_name, + message=f"Execution DAG generation failed: {exc}", + severity=ErrorSeverity.ERROR, + error_code=ErrorCode.PLANNER_FAILED, + ) + ], + } + return { "decomposer_response": response, + "execution_dag": execution_dag, "reasoning": [{"node": self.node_name, "content": "Decomposition completed."}], } diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/__init__.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/__init__.py deleted file mode 100644 index f3cf0cda..00000000 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .node import GlobalPlannerNode -from .schemas import GlobalPlannerResponse - -__all__ = ["GlobalPlannerNode", "GlobalPlannerResponse"] diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py deleted file mode 100644 index f0b61e3e..00000000 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/node.py +++ /dev/null @@ -1,187 +0,0 @@ -from __future__ import annotations -from typing import Dict, Any, TYPE_CHECKING, List - -if TYPE_CHECKING: - from nl2sql.pipeline.state import GraphState - -from .schemas import ( - RelationSchema, - ColumnSpec, - ExecutionDAG, - LogicalNode, - LogicalEdge, -) -from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode -from nl2sql.common.logger import get_logger -from nl2sql.context import NL2SQLContext -from .schemas import GlobalPlannerResponse - -logger = get_logger("global_planner") - - -class GlobalPlannerNode: - """Generates a deterministic execution plan for aggregating sub-query results. - - This node runs AFTER the Decomposer and BEFORE the SQL Agents. - It does not execute SQL, but produces a blueprint (ResultPlan) for the Aggregator. - """ - - def __init__(self, ctx: NL2SQLContext): - self.node_name = self.__class__.__name__.lower().replace("node", "") - - def __call__(self, state: GraphState) -> Dict[str, Any]: - """Executes the planning logic.""" - decomposer_response = state.decomposer_response - sub_queries = decomposer_response.sub_queries if decomposer_response else [] - combine_groups = decomposer_response.combine_groups if decomposer_response else [] - post_combine_ops = decomposer_response.post_combine_ops if decomposer_response else [] - - try: - nodes: List[LogicalNode] = [] - edges: List[LogicalEdge] = [] - node_index: Dict[str, LogicalNode] = {} - - for sq in sub_queries: - schema = RelationSchema( - columns=[ColumnSpec(name=c.name, dtype=c.dtype) for c in sq.expected_schema] - ) - node = LogicalNode( - node_id=sq.id, - kind="scan", - inputs=[], - output_schema=schema, - attributes={ - "datasource_id": sq.datasource_id, - "intent": sq.intent, - "metrics": [m.model_dump() for m in sq.metrics], - "filters": [f.model_dump() for f in sq.filters], - "group_by": [g.model_dump() for g in sq.group_by], - "expected_schema": [c.model_dump() for c in sq.expected_schema], - }, - ) - nodes.append(node) - node_index[node.node_id] = node - - combine_node_ids: Dict[str, str] = {} - for cg in combine_groups: - combine_id = f"combine_{cg.group_id}" - combine_node_ids[cg.group_id] = combine_id - input_ids = [i.subquery_id for i in cg.inputs] - output_schema = ( - node_index[input_ids[0]].output_schema - if input_ids and input_ids[0] in node_index - else RelationSchema(columns=[]) - ) - node = LogicalNode( - node_id=combine_id, - kind="combine", - inputs=input_ids, - output_schema=output_schema, - attributes={ - "operation": cg.operation, - "group_id": cg.group_id, - "inputs": [i.model_dump() for i in cg.inputs], - "join_keys": [jk.model_dump() for jk in cg.join_keys], - }, - ) - nodes.append(node) - node_index[node.node_id] = node - - for inp in cg.inputs: - edges.append( - LogicalEdge( - edge_id=f"edge_{inp.subquery_id}_{combine_id}", - from_id=inp.subquery_id, - to_id=combine_id, - role=inp.role, - ) - ) - - for op in post_combine_ops: - target_combine_id = combine_node_ids.get(op.target_group_id) - if not target_combine_id: - raise ValueError(f"PostCombineOp references unknown combine group: {op.target_group_id}") - output_schema = RelationSchema( - columns=[ColumnSpec(name=c.name, dtype=c.dtype) for c in op.expected_schema] - ) - kind_map = { - "filter": "post_filter", - "aggregate": "post_aggregate", - "project": "post_project", - "sort": "post_sort", - "limit": "post_limit", - } - node = LogicalNode( - node_id=op.op_id, - kind=kind_map[op.operation], - inputs=[target_combine_id], - output_schema=output_schema, - attributes={ - "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, - }, - ) - nodes.append(node) - node_index[node.node_id] = node - edges.append( - LogicalEdge( - edge_id=f"edge_{target_combine_id}_{op.op_id}", - from_id=target_combine_id, - to_id=op.op_id, - ) - ) - - node_ids = {n.node_id for n in nodes} - for edge in edges: - if edge.from_id not in node_ids or edge.to_id not in node_ids: - raise ValueError(f"Edge references unknown node: {edge}") - - execution_dag = ExecutionDAG( - nodes=sorted(nodes, key=lambda n: n.node_id), - edges=sorted(edges, key=lambda e: (e.from_id, e.to_id, e.role or "")), - ) - - execution_dag.content_hash = self._hash_execution_dag(execution_dag) - execution_dag.dag_id = f"dag_{execution_dag.content_hash[:12]}" - - result = GlobalPlannerResponse(execution_dag=execution_dag) - - return { - "global_planner_response": result, - "reasoning": [{"node": self.node_name, "content": "Built explicit execution DAG."}], - } - - except Exception as e: - logger.error(f"GlobalPlanner failed: {e}") - # No response: there is no DAG, and the layer router ends the run - # when it finds none, leaving this error as the run's cause. - return { - "reasoning": [{"node": self.node_name, "content": f"Planning failed: {e}", "type": "error"}], - "errors": [ - PipelineError( - node=self.node_name, - message=f"Global Plan generation failed: {str(e)}", - severity=ErrorSeverity.ERROR, - error_code=ErrorCode.PLANNER_FAILED, - ) - ] - } - - def _hash_execution_dag(self, dag: ExecutionDAG) -> str: - import hashlib - import json - - payload = { - "nodes": [n.model_dump() for n in dag.nodes], - "edges": [e.model_dump() for e in dag.edges], - "version": dag.version, - } - data = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True) - return hashlib.sha256(data.encode("utf-8")).hexdigest() diff --git a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/schemas.py b/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/schemas.py deleted file mode 100644 index 26c26262..00000000 --- a/packages/nl2sql/src/nl2sql/pipeline/nodes/global_planner/schemas.py +++ /dev/null @@ -1,17 +0,0 @@ -from __future__ import annotations - -from pydantic import BaseModel - -# The DAG models live in nl2sql.execution.dag, shared with the aggregation service. -from nl2sql.execution.dag import ( # noqa: F401 - ColumnSpec, - ExecutionDAG, - JsonLiteral, - LogicalEdge, - LogicalNode, - RelationSchema, -) - - -class GlobalPlannerResponse(BaseModel): - execution_dag: ExecutionDAG diff --git a/packages/nl2sql/src/nl2sql/pipeline/routes.py b/packages/nl2sql/src/nl2sql/pipeline/routes.py index 06d71752..0cf35757 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/routes.py +++ b/packages/nl2sql/src/nl2sql/pipeline/routes.py @@ -42,8 +42,7 @@ def build_scan_layer_router(ctx: NL2SQLContext, execute: bool = True): """ def route_scan_layers(state: GraphState): - global_planner_response = state.global_planner_response - dag = global_planner_response.execution_dag if global_planner_response else None + dag = state.execution_dag decomposer_response = state.decomposer_response sub_queries = decomposer_response.sub_queries if decomposer_response else [] sub_query_map = {sq.id: sq for sq in sub_queries} @@ -51,7 +50,6 @@ def route_scan_layers(state: GraphState): if not dag or not dag.layers: return END - node_index = {n.node_id: n for n in dag.nodes} target_ids = next_scan_layer_ids( dag, artifact_refs, completed_scan_ids(state.subgraph_outputs) ) @@ -62,21 +60,17 @@ def route_scan_layers(state: GraphState): branches = [] for node_id in target_ids: - if node_id in sub_query_map: - sq = sub_query_map[node_id] - datasource_id = sq.datasource_id - else: - node = node_index.get(node_id) if dag else None - if not node: - continue - datasource_id = (node.attributes or {}).get("datasource_id") + # A scan node's id is its sub-query's id, so the lookup always hits. + sq = sub_query_map.get(node_id) + if sq is None: + continue - target = resolve_subgraph(datasource_id, ctx) + target = resolve_subgraph(sq.datasource_id, ctx) if not target: raise PipelineExecutionError( PipelineError( node="layer_router", - message=f"No compatible subgraph found for datasource '{datasource_id}'.", + message=f"No compatible subgraph found for datasource '{sq.datasource_id}'.", severity=ErrorSeverity.ERROR, error_code=ErrorCode.INVALID_STATE, ) diff --git a/packages/nl2sql/src/nl2sql/pipeline/state.py b/packages/nl2sql/src/nl2sql/pipeline/state.py index 20610b73..23ce645f 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/state.py +++ b/packages/nl2sql/src/nl2sql/pipeline/state.py @@ -10,13 +10,13 @@ from nl2sql.auth import UserContext from nl2sql.pipeline.nodes.datasource_resolver.schemas import DatasourceResolverResponse from nl2sql.pipeline.nodes.decomposer.schemas import DecomposerResponse, SubQuery -from nl2sql.pipeline.nodes.global_planner.schemas import GlobalPlannerResponse from nl2sql.pipeline.nodes.aggregator.schemas import AggregatorResponse from nl2sql.pipeline.nodes.answer_synthesizer.schemas import AnswerSynthesizerResponse from nl2sql.pipeline.nodes.ast_planner.schemas import ASTPlannerResponse from nl2sql.pipeline.nodes.validator.schemas import LogicalValidatorResponse from nl2sql.pipeline.nodes.generator.schemas import GeneratorResponse from nl2sql.execution.contracts import ArtifactRef, ExecutorResponse +from nl2sql.execution.dag import ExecutionDAG from nl2sql.pipeline.nodes.refiner.schemas import RefinerResponse from nl2sql.pipeline.subgraphs.schemas import SubgraphOutput from nl2sql.pipeline.nodes.schema_retriever.schema import Table @@ -39,7 +39,7 @@ class GraphState(BaseModel): datasource_id (Optional[str]): Optional datasource override for resolution. datasource_resolver_response (Optional[DatasourceResolverResponse]): Output of resolver node. decomposer_response (Optional[DecomposerResponse]): Output of decomposer node. - global_planner_response (Optional[GlobalPlannerResponse]): Output of planner node. + execution_dag (Optional[ExecutionDAG]): The graph of sub-queries the decomposer built. aggregator_response (Optional[AggregatorResponse]): Output of aggregator node. answer_synthesizer_response (Optional[AnswerSynthesizerResponse]): Output of synthesizer node. artifact_refs (Dict[str, ArtifactRef]): Artifact refs keyed by ExecutionDAG node_id. @@ -58,7 +58,7 @@ class GraphState(BaseModel): datasource_id: Optional[str] = Field(default=None, description="Optional datasource override for resolution.") datasource_resolver_response: Optional[DatasourceResolverResponse] = Field(default=None) decomposer_response: Optional[DecomposerResponse] = Field(default=None) - global_planner_response: Optional[GlobalPlannerResponse] = Field(default=None) + execution_dag: Optional[ExecutionDAG] = Field(default=None) aggregator_response: Optional[AggregatorResponse] = Field(default=None) answer_synthesizer_response: Optional[AnswerSynthesizerResponse] = Field(default=None) artifact_refs: Annotated[Dict[str, ArtifactRef], update_results] = Field(default_factory=dict) diff --git a/packages/nl2sql/src/nl2sql/pipeline/steps.py b/packages/nl2sql/src/nl2sql/pipeline/steps.py index 5563614e..ee4dbdb0 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/steps.py +++ b/packages/nl2sql/src/nl2sql/pipeline/steps.py @@ -60,11 +60,9 @@ def _step(node: str, label: str, does: str, parent: Optional[str] = None) -> Pip "Decides whether the connected databases can answer the question at all, and which one " "it is about, before anything else runs."), _step("decomposer", "Question splitter", - "Breaks the question into the sub-queries that have to be answered, and says how their " - "answers combine."), - _step("global_planner", "Execution plan", - "Arranges those sub-queries into a dependency graph: which can run together, which wait " - "on another's rows."), + "Breaks the question into the sub-queries that have to be answered, says how their " + "answers combine, and arranges them into a dependency graph: which can run together, " + "which wait on another's rows."), _step("layer_router", "Layer router", "Sends the next layer of sub-queries to the SQL agent, and stops when every one of them " "has an answer."), diff --git a/packages/nl2sql/src/nl2sql/pipeline/subgraphs/sql_agent.py b/packages/nl2sql/src/nl2sql/pipeline/subgraphs/sql_agent.py index fd1d6da6..e1f345cd 100644 --- a/packages/nl2sql/src/nl2sql/pipeline/subgraphs/sql_agent.py +++ b/packages/nl2sql/src/nl2sql/pipeline/subgraphs/sql_agent.py @@ -6,7 +6,7 @@ from nl2sql.common.errors import PipelineError, ErrorSeverity, ErrorCode from nl2sql.common.settings import settings from nl2sql.pipeline.state import SubgraphExecutionState -from nl2sql.pipeline.nodes.ast_planner import ASTPlannerNode +from nl2sql.pipeline.nodes.ast_planner.node import ASTPlannerNode from nl2sql.pipeline.nodes.schema_retriever import SchemaRetrieverNode from nl2sql.pipeline.nodes.validator import LogicalValidatorNode from nl2sql.pipeline.nodes.refiner import RefinerNode diff --git a/packages/nl2sql/src/nl2sql/tracing/recorder.py b/packages/nl2sql/src/nl2sql/tracing/recorder.py index 64e5c2a7..c3be76c4 100644 --- a/packages/nl2sql/src/nl2sql/tracing/recorder.py +++ b/packages/nl2sql/src/nl2sql/tracing/recorder.py @@ -42,10 +42,9 @@ NODE_INPUTS: Dict[str, Tuple[str, ...]] = { "datasource_resolver": ("user_query", "datasource_id", "user_context"), "decomposer": ("user_query", "datasource_resolver_response"), - "global_planner": ("decomposer_response",), "layer_router": ("artifact_refs",), "sql_agent": ("subgraph_id",), - "aggregator": ("global_planner_response", "artifact_refs"), + "aggregator": ("execution_dag", "artifact_refs"), "answer_synthesizer": ("user_query", "aggregator_response"), "schema_retriever": ("sub_query",), "ast_planner": ("sub_query", "relevant_tables", "errors", "retry_count"), diff --git a/packages/nl2sql/tests/architecture/test_boundaries.py b/packages/nl2sql/tests/architecture/test_boundaries.py index 7f7079e7..be5b063c 100644 --- a/packages/nl2sql/tests/architecture/test_boundaries.py +++ b/packages/nl2sql/tests/architecture/test_boundaries.py @@ -439,6 +439,17 @@ def test_the_engine_builds_its_benchmark_api_lazily(): assert isinstance(NL2SQL.__dict__["benchmark"], property) +def test_the_plan_cache_imports_without_the_nodes_package_importing_it_back(): + """``plan_cache`` needs ``PlanModel``; ``ast_planner.node`` needs ``PlanCache``. + + While ``nodes/__init__`` and ``nodes/ast_planner/__init__`` re-exported the + node classes, importing the schemas ran the node first, and the pair was a + cycle that only stayed hidden because some earlier import happened to load + them in the lucky order. Importing the cache on its own is the test. + """ + assert _run("import nl2sql.pipeline.plan_cache; print('ok')") == "ok", f"nl2sql.pipeline.plan_cache cannot be imported on its own. The rule is written down in {RULES}." + + def test_aggregation_never_imports_the_pipeline(): """The pipeline's aggregator node calls the aggregation service; the reverse was a cycle.""" offenders = [f"{_where(path)}:{lineno} {name}" diff --git a/packages/nl2sql/tests/e2e/test_trace_fake_llm.py b/packages/nl2sql/tests/e2e/test_trace_fake_llm.py index f9ad78b7..feaa6c70 100644 --- a/packages/nl2sql/tests/e2e/test_trace_fake_llm.py +++ b/packages/nl2sql/tests/e2e/test_trace_fake_llm.py @@ -30,7 +30,7 @@ # Every node a counted, executed question passes through. NODES = { - "datasource_resolver", "decomposer", "global_planner", "layer_router", "sql_agent", + "datasource_resolver", "decomposer", "layer_router", "sql_agent", "schema_retriever", "ast_planner", "logical_validator", "generator", "executor", "aggregator", "answer_synthesizer", } diff --git a/packages/nl2sql/tests/integration/test_aggregator_real_data.py b/packages/nl2sql/tests/integration/test_aggregator_real_data.py index 4589cc38..c818dc6f 100644 --- a/packages/nl2sql/tests/integration/test_aggregator_real_data.py +++ b/packages/nl2sql/tests/integration/test_aggregator_real_data.py @@ -12,13 +12,10 @@ from nl2sql.common.settings import settings from nl2sql.context import NL2SQLContext from nl2sql.pipeline.nodes.aggregator.node import EngineAggregatorNode -from nl2sql.pipeline.nodes.global_planner.schemas import ( +from nl2sql.execution.dag import ( ExecutionDAG, LogicalNode, LogicalEdge, - RelationSchema, - ColumnSpec, - GlobalPlannerResponse, ) from nl2sql.pipeline.state import GraphState from nl2sql.execution.artifacts import ArtifactStore, ArtifactStoreConfig @@ -103,19 +100,16 @@ def demo_env() -> SimpleNamespace: shutil.rmtree(tmp_dir, ignore_errors=True) -def _schema(columns): - return RelationSchema(columns=[ColumnSpec(name=c) for c in columns]) def test_aggregator_real_data(demo_env) -> None: node = EngineAggregatorNode(demo_env.ctx) - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["value"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) post_filter = LogicalNode( node_id="op_filter", kind="post_filter", inputs=["sq_1"], - output_schema=_schema(["value"]), attributes={"operation": "filter", "filters": [{"attribute": "value", "operator": ">", "value": 10}]}, ) dag = ExecutionDAG( @@ -142,7 +136,7 @@ def test_aggregator_real_data(demo_env) -> None: ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_1": artifact}, ) diff --git a/packages/nl2sql/tests/integration/test_answer_synthesizer_real_data.py b/packages/nl2sql/tests/integration/test_answer_synthesizer_real_data.py index d9308368..3e61d677 100644 --- a/packages/nl2sql/tests/integration/test_answer_synthesizer_real_data.py +++ b/packages/nl2sql/tests/integration/test_answer_synthesizer_real_data.py @@ -13,13 +13,10 @@ from nl2sql.context import NL2SQLContext from nl2sql.pipeline.nodes.aggregator.node import EngineAggregatorNode from nl2sql.pipeline.nodes.answer_synthesizer.node import AnswerSynthesizerNode -from nl2sql.pipeline.nodes.global_planner.schemas import ( +from nl2sql.execution.dag import ( ExecutionDAG, LogicalNode, LogicalEdge, - RelationSchema, - ColumnSpec, - GlobalPlannerResponse, ) from nl2sql.pipeline.nodes.aggregator.schemas import AggregatorResponse from nl2sql.pipeline.state import GraphState @@ -109,17 +106,14 @@ def demo_env() -> SimpleNamespace: shutil.rmtree(tmp_dir, ignore_errors=True) -def _schema(columns): - return RelationSchema(columns=[ColumnSpec(name=c) for c in columns]) def _build_aggregator_response() -> AggregatorResponse: - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["value"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) post_filter = LogicalNode( node_id="op_filter", kind="post_filter", inputs=["sq_1"], - output_schema=_schema(["value"]), attributes={"operation": "filter", "filters": [{"attribute": "value", "operator": ">", "value": 10}]}, ) dag = ExecutionDAG( @@ -146,7 +140,7 @@ def _build_aggregator_response() -> AggregatorResponse: ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_1": artifact}, ) return EngineAggregatorNode(SimpleNamespace())(state)["aggregator_response"] diff --git a/packages/nl2sql/tests/unit/test_boundary_tidy_ups.py b/packages/nl2sql/tests/unit/test_boundary_tidy_ups.py index e55e388f..198bbd21 100644 --- a/packages/nl2sql/tests/unit/test_boundary_tidy_ups.py +++ b/packages/nl2sql/tests/unit/test_boundary_tidy_ups.py @@ -32,12 +32,12 @@ def test_the_wheel_carries_the_chinook_database(): # F18: the DAG models are neutral, so the aggregation service need not import # the pipeline. The import rule itself lives in tests/architecture/. -def test_the_global_planner_uses_the_neutral_dag_models(): +def test_the_decomposer_builds_the_neutral_dag_models(): from nl2sql.execution import dag - from nl2sql.pipeline.nodes.global_planner import schemas + from nl2sql.pipeline.nodes.decomposer import dag as builder - assert schemas.ExecutionDAG is dag.ExecutionDAG - assert schemas.LogicalNode is dag.LogicalNode + assert builder.ExecutionDAG is dag.ExecutionDAG + assert builder.LogicalNode is dag.LogicalNode # F16: SQL is compared in the datasource's dialect, not sqlglot's default. diff --git a/packages/nl2sql/tests/unit/test_node_global_planner.py b/packages/nl2sql/tests/unit/test_decomposer_dag.py similarity index 76% rename from packages/nl2sql/tests/unit/test_node_global_planner.py rename to packages/nl2sql/tests/unit/test_decomposer_dag.py index 10418589..bd9f33e0 100644 --- a/packages/nl2sql/tests/unit/test_node_global_planner.py +++ b/packages/nl2sql/tests/unit/test_decomposer_dag.py @@ -1,6 +1,11 @@ +"""The execution DAG the decomposer builds from its own decomposition. + +It was a graph node of its own until the DAG turned out to be a pure function +of ``DecomposerResponse``; these are that node's tests, against the function. +""" import pytest -from nl2sql.pipeline.nodes.global_planner.node import GlobalPlannerNode +from nl2sql.pipeline.nodes.decomposer.dag import build_execution_dag from nl2sql.pipeline.nodes.decomposer.schemas import ( DecomposerResponse, SubQuery, @@ -9,10 +14,9 @@ PostCombineOp, ExpectedColumn, ) -from nl2sql.pipeline.state import GraphState -def test_global_planner_builds_execution_dag(): +def test_the_dag_builds_execution_dag(): # Validates DAG construction because aggregation depends on correct edges. # Arrange sub_queries = [ @@ -63,12 +67,7 @@ def test_global_planner_builds_execution_dag(): unmapped_subqueries=[], ) - node = GlobalPlannerNode(ctx=None) - state = GraphState(user_query="q", decomposer_response=response) - - # Act - result = node(state) - dag = result["global_planner_response"].execution_dag + dag = build_execution_dag(response) # Assert node_ids = {n.node_id for n in dag.nodes} @@ -76,10 +75,9 @@ def test_global_planner_builds_execution_dag(): assert "sq2" in node_ids assert "combine_g1" in node_ids assert any(n.kind.startswith("post_") for n in dag.nodes) - assert dag.dag_id -def test_global_planner_dag_layers_and_edges(): +def test_the_dag_dag_layers_and_edges(): # Validates layer rules and acyclic DAG for deterministic planning. sub_queries = [ SubQuery( @@ -142,11 +140,7 @@ def test_global_planner_dag_layers_and_edges(): unmapped_subqueries=[], ) - node = GlobalPlannerNode(ctx=None) - state = GraphState(user_query="q", decomposer_response=response) - - result = node(state) - dag = result["global_planner_response"].execution_dag + dag = build_execution_dag(response) node_index = {n.node_id: n for n in dag.nodes} assert set(node_index) == {"sq_base", "sq_left", "sq_right", "combine_g_join", "op_stub"} @@ -164,11 +158,12 @@ def test_global_planner_dag_layers_and_edges(): assert edge.from_id in node_index assert edge.to_id in node_index - scan_schema = node_index["sq_base"].output_schema - assert [c.name for c in scan_schema.columns] == ["base_id"] + # A scan node carries nothing of its own: its id is its sub-query's id, + # so everything about it is one lookup away in the decomposer response. + assert node_index["sq_base"].attributes == {} -def test_global_planner_scan_only_dag(): +def test_the_dag_scan_only_dag(): # Validates scan-only plans because some queries skip combine/post stages. sub_queries = [ SubQuery( @@ -197,10 +192,7 @@ def test_global_planner_scan_only_dag(): unmapped_subqueries=[], ) - node = GlobalPlannerNode(ctx=None) - state = GraphState(user_query="q", decomposer_response=response) - result = node(state) - dag = result["global_planner_response"].execution_dag + dag = build_execution_dag(response) assert len(dag.edges) == 0 assert dag.layers @@ -209,7 +201,7 @@ def test_global_planner_scan_only_dag(): assert node_obj.kind == "scan" -def test_global_planner_multiple_combine_groups(): +def test_the_dag_multiple_combine_groups(): # Validates multi-group plans because complex queries may have multiple combines. sub_queries = [ SubQuery( @@ -265,17 +257,14 @@ def test_global_planner_multiple_combine_groups(): unmapped_subqueries=[], ) - node = GlobalPlannerNode(ctx=None) - state = GraphState(user_query="q", decomposer_response=response) - result = node(state) - dag = result["global_planner_response"].execution_dag + dag = build_execution_dag(response) node_ids = {n.node_id for n in dag.nodes} assert "combine_g1" in node_ids assert "combine_g2" in node_ids -def test_global_planner_unknown_post_combine_group_returns_error(): +def test_the_dag_unknown_post_combine_group_returns_error(): # Validates error handling when post-ops reference unknown groups. with pytest.raises(ValueError, match="PostCombineOp references unknown combine group"): DecomposerResponse( @@ -309,7 +298,7 @@ def test_global_planner_unknown_post_combine_group_returns_error(): ) -def test_global_planner_unknown_subquery_in_combine_returns_error(): +def test_the_dag_unknown_subquery_in_combine_returns_error(): # Validates error handling for edges pointing to unknown nodes. with pytest.raises(ValueError, match="CombineGroup references unknown subquery"): DecomposerResponse( @@ -336,32 +325,57 @@ def test_global_planner_unknown_subquery_in_combine_returns_error(): ) -def test_a_planner_failure_reports_its_own_error_instead_of_crashing(): - # Duplicate expected_schema names are a real way the decomposer's output - # breaks the DAG. The error handler used to build a response the schema - # rejects, so the run surfaced "Pipeline crashed" instead of this error. +def test_a_dag_that_cannot_be_built_ends_the_run_with_its_error(): + """A DAG failure is reported, never raised through the graph. + + Two sub-queries the model wrote identically get the same content-addressed + id, so the graph they describe has one node where the combine group expects + two. That was the global planner's own error path; it is now the + decomposer's, which keeps the decomposition, emits no ``execution_dag`` -- + the layer router ends a run without one -- and reports the cause. + """ + from types import SimpleNamespace + from unittest.mock import MagicMock + from nl2sql.api.query_api import result_from_state from nl2sql.common.errors import ErrorCode + from nl2sql.pipeline.nodes.datasource_resolver.schemas import ( + DatasourceResolverResponse, + ResolvedDatasource, + ) + from nl2sql.pipeline.nodes.decomposer.node import DecomposerNode + from nl2sql.pipeline.state import GraphState - response = DecomposerResponse( + twins = DecomposerResponse( sub_queries=[ - SubQuery( - id="sq1", - intent="customers per country", - datasource_id="chinook", - expected_schema=[ExpectedColumn(name="country"), ExpectedColumn(name="country")], - ) + SubQuery(id="a", intent="customers per country", datasource_id="chinook"), + SubQuery(id="b", intent="customers per country", datasource_id="chinook"), ], - combine_groups=[], + combine_groups=[CombineGroup( + group_id="g1", operation="union", + inputs=[CombineInput(subquery_id="a", role="left"), + CombineInput(subquery_id="b", role="right")], + )], ) - result = GlobalPlannerNode(ctx=None)(GraphState(user_query="q", decomposer_response=response)) + node = DecomposerNode(SimpleNamespace(llm_registry=MagicMock())) + node.chain = MagicMock() + node.chain.invoke.return_value = twins + + result = node(GraphState( + user_query="q", + datasource_resolver_response=DatasourceResolverResponse( + resolved_datasources=[ResolvedDatasource(datasource_id="chinook", schema_version="v1")], + allowed_datasource_ids=["chinook"], + ), + )) [error] = result["errors"] assert error.error_code == ErrorCode.PLANNER_FAILED - assert error.node == "globalplanner" - assert "Duplicate columns in schema" in error.message - assert result.get("global_planner_response") is None + assert error.node == "decomposer" + assert result.get("execution_dag") is None + # The decomposition itself survives, so the failure names what was decomposed. + assert result["decomposer_response"].sub_queries query_result = result_from_state({**result, "trace_id": "t"}) assert query_result.status == "error" diff --git a/packages/nl2sql/tests/unit/test_graph_layers.py b/packages/nl2sql/tests/unit/test_graph_layers.py index 18c2b20b..b187d317 100644 --- a/packages/nl2sql/tests/unit/test_graph_layers.py +++ b/packages/nl2sql/tests/unit/test_graph_layers.py @@ -1,27 +1,22 @@ from nl2sql.pipeline.graph_utils import next_scan_layer_ids -from nl2sql.pipeline.nodes.global_planner.schemas import ( +from nl2sql.execution.dag import ( ExecutionDAG, LogicalNode, LogicalEdge, - RelationSchema, - ColumnSpec, ) -def _schema(columns): - return RelationSchema(columns=[ColumnSpec(name=c) for c in columns]) def test_next_scan_layer_ids_respects_existing_results(): # Validates DAG routing because scan layers must honor completed nodes. # Arrange - scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[], output_schema=_schema(["id"])) - scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[], output_schema=_schema(["id"])) + scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[]) + scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[]) combine = LogicalNode( node_id="combine_cg_1", kind="combine", inputs=["sq_left", "sq_right"], - output_schema=_schema(["id"]), ) dag = ExecutionDAG( nodes=[scan_left, scan_right, combine], @@ -47,7 +42,6 @@ def test_completed_scan_without_artifact_is_not_pending(): kind="scan", attributes={}, inputs=[], - output_schema=_schema(["id"]), ) ], edges=[], diff --git a/packages/nl2sql/tests/unit/test_node_aggregator.py b/packages/nl2sql/tests/unit/test_node_aggregator.py index fc489b0c..5dbb7e64 100644 --- a/packages/nl2sql/tests/unit/test_node_aggregator.py +++ b/packages/nl2sql/tests/unit/test_node_aggregator.py @@ -2,14 +2,11 @@ import tempfile from nl2sql.pipeline.nodes.aggregator.node import EngineAggregatorNode -from nl2sql.pipeline.nodes.global_planner.schemas import ( +from nl2sql.execution.dag import ( ExecutionDAG, LogicalNode, LogicalEdge, - RelationSchema, - ColumnSpec, ) -from nl2sql.pipeline.nodes.global_planner.schemas import GlobalPlannerResponse from nl2sql.pipeline.state import GraphState from nl2sql_adapter_sdk.contracts import ResultFrame from nl2sql.common.errors import ErrorCode @@ -17,8 +14,6 @@ from nl2sql.common.settings import settings -def _schema(columns): - return RelationSchema(columns=[ColumnSpec(name=c) for c in columns]) def test_aggregator_filters_rows_deterministically(monkeypatch): @@ -27,12 +22,11 @@ def test_aggregator_filters_rows_deterministically(monkeypatch): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["id", "value"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) post_filter = LogicalNode( node_id="op_filter", kind="post_filter", inputs=["sq_1"], - output_schema=_schema(["id", "value"]), attributes={"operation": "filter", "filters": [{"attribute": "value", "operator": ">", "value": 10}]}, ) dag = ExecutionDAG( @@ -67,7 +61,7 @@ def test_aggregator_filters_rows_deterministically(monkeypatch): ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_1": artifact}, ) @@ -85,13 +79,12 @@ def test_aggregator_join_requires_keys(monkeypatch): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[], output_schema=_schema(["id"])) - scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[], output_schema=_schema(["id"])) + scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[]) + scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[]) combine = LogicalNode( node_id="combine_cg_1", kind="combine", inputs=["sq_left", "sq_right"], - output_schema=_schema(["id"]), attributes={"operation": "join", "join_keys": []}, ) dag = ExecutionDAG( @@ -139,7 +132,7 @@ def test_aggregator_join_requires_keys(monkeypatch): ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_left": left_artifact, "sq_right": right_artifact}, ) @@ -155,7 +148,7 @@ def test_aggregator_requires_execution_dag(): # Arrange ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - state = GraphState(user_query="q", global_planner_response=None) + state = GraphState(user_query="q", execution_dag=None) # Act result = node(state) @@ -169,12 +162,11 @@ def test_aggregator_post_aggregate_sum(monkeypatch): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["value"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) post_agg = LogicalNode( node_id="op_agg", kind="post_aggregate", inputs=["sq_1"], - output_schema=_schema(["total_value"]), attributes={ "operation": "aggregate", "metrics": [{"name": "value", "aggregation": "sum"}], @@ -213,7 +205,7 @@ def test_aggregator_post_aggregate_sum(monkeypatch): ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_1": artifact}, ) @@ -228,12 +220,11 @@ def test_aggregator_post_sort_limit(monkeypatch): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["value"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) post_sort = LogicalNode( node_id="op_sort", kind="post_sort", inputs=["sq_1"], - output_schema=_schema(["value"]), attributes={ "operation": "sort", "order_by": [{"attribute": "value", "direction": "desc"}], @@ -243,7 +234,6 @@ def test_aggregator_post_sort_limit(monkeypatch): node_id="op_limit", kind="post_limit", inputs=["op_sort"], - output_schema=_schema(["value"]), attributes={ "operation": "limit", "limit": 1, @@ -284,7 +274,7 @@ def test_aggregator_post_sort_limit(monkeypatch): ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_1": artifact}, ) @@ -299,13 +289,12 @@ def test_aggregator_union_combines_rows(monkeypatch): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[], output_schema=_schema(["id"])) - scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[], output_schema=_schema(["id"])) + scan_left = LogicalNode(node_id="sq_left", kind="scan", inputs=[]) + scan_right = LogicalNode(node_id="sq_right", kind="scan", inputs=[]) combine = LogicalNode( node_id="combine_union", kind="combine", inputs=["sq_left", "sq_right"], - output_schema=_schema(["id"]), attributes={"operation": "union", "join_keys": []}, ) dag = ExecutionDAG( @@ -353,7 +342,7 @@ def test_aggregator_union_combines_rows(monkeypatch): ) state = GraphState( user_query="q", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs={"sq_left": left_artifact, "sq_right": right_artifact}, ) @@ -368,9 +357,9 @@ def test_aggregator_missing_artifact_reference(): ctx = SimpleNamespace() node = EngineAggregatorNode(ctx) - scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[], output_schema=_schema(["id"])) + scan = LogicalNode(node_id="sq_1", kind="scan", inputs=[]) dag = ExecutionDAG(nodes=[scan], edges=[]) - state = GraphState(user_query="q", global_planner_response=GlobalPlannerResponse(execution_dag=dag)) + state = GraphState(user_query="q", execution_dag=dag) result = node(state) @@ -387,13 +376,12 @@ def test_the_jazz_but_never_rock_dag_is_refused_with_the_reason(monkeypatch): # Arrange node = EngineAggregatorNode(SimpleNamespace()) - jazz = LogicalNode(node_id="sq_jazz", kind="scan", inputs=[], output_schema=_schema(["customer"])) - rock = LogicalNode(node_id="sq_rock", kind="scan", inputs=[], output_schema=_schema(["customer"])) + jazz = LogicalNode(node_id="sq_jazz", kind="scan", inputs=[]) + rock = LogicalNode(node_id="sq_rock", kind="scan", inputs=[]) combine = LogicalNode( node_id="combine_cg_1", kind="combine", inputs=["sq_jazz", "sq_rock"], - output_schema=_schema(["customer"]), attributes={"operation": "join", "join_keys": [{"left": "customer", "right": "right.customer"}]}, ) @@ -401,7 +389,6 @@ def test_the_jazz_but_never_rock_dag_is_refused_with_the_reason(monkeypatch): node_id="op_not_rock", kind="post_filter", inputs=["combine_cg_1"], - output_schema=_schema(["customer"]), attributes={"operation": "filter", "filters": [{"attribute": "right.customer", "operator": "=", "value": None}]}, ) @@ -438,7 +425,7 @@ def test_the_jazz_but_never_rock_dag_is_refused_with_the_reason(monkeypatch): ) state = GraphState( user_query="Which customers bought jazz tracks but never rock?", - global_planner_response=GlobalPlannerResponse(execution_dag=dag), + execution_dag=dag, artifact_refs=refs, ) diff --git a/packages/nl2sql/tests/unit/test_pipeline_steps.py b/packages/nl2sql/tests/unit/test_pipeline_steps.py index 8bc31d67..6f3eaa24 100644 --- a/packages/nl2sql/tests/unit/test_pipeline_steps.py +++ b/packages/nl2sql/tests/unit/test_pipeline_steps.py @@ -26,7 +26,7 @@ def graph_nodes(monkeypatch): from nl2sql.pipeline.subgraphs import sql_agent for name in ("DatasourceResolverNode", "DecomposerNode", "EngineAggregatorNode", - "AnswerSynthesizerNode", "GlobalPlannerNode"): + "AnswerSynthesizerNode"): monkeypatch.setattr(control, name, _blank) for name in ("SchemaRetrieverNode", "ASTPlannerNode", "LogicalValidatorNode", "GeneratorNode", "ExecutorNode", "RefinerNode"): diff --git a/packages/nl2sql/tests/unit/test_playground_pipeline.py b/packages/nl2sql/tests/unit/test_playground_pipeline.py index 19b83193..eadda8bd 100644 --- a/packages/nl2sql/tests/unit/test_playground_pipeline.py +++ b/packages/nl2sql/tests/unit/test_playground_pipeline.py @@ -78,7 +78,7 @@ def test_the_route_answers_on_the_hosted_demo_too(): def test_an_engine_that_cannot_report_its_models_still_lists_the_steps(): body = _client(_Broken()).get("/api/pipeline").json() - assert len(body["steps"]) == 14 + assert len(body["steps"]) == 13 assert all(s["model"] is None for s in body["steps"]) diff --git a/packages/nl2sql/tests/unit/test_scan_router_execute.py b/packages/nl2sql/tests/unit/test_scan_router_execute.py index fe4cc0bb..439d8a0d 100644 --- a/packages/nl2sql/tests/unit/test_scan_router_execute.py +++ b/packages/nl2sql/tests/unit/test_scan_router_execute.py @@ -8,7 +8,7 @@ def test_router_ends_instead_of_aggregating_when_execute_is_false(monkeypatch): monkeypatch.setattr("nl2sql.pipeline.routes.next_scan_layer_ids", lambda dag, refs, completed=frozenset(): []) dag = SimpleNamespace(layers=[["sq1"]], nodes=[]) - state = SimpleNamespace(global_planner_response=SimpleNamespace(execution_dag=dag), + state = SimpleNamespace(execution_dag=dag, decomposer_response=None, artifact_refs={"sq1": object()}, subgraph_outputs={}) assert build_scan_layer_router(SimpleNamespace(), execute=False)(state) == END sends = build_scan_layer_router(SimpleNamespace(), execute=True)(state) @@ -17,7 +17,7 @@ def test_router_ends_instead_of_aggregating_when_execute_is_false(monkeypatch): def _finished_state(artifact_refs): dag = SimpleNamespace(layers=[["sq1"]], nodes=[]) - return SimpleNamespace(global_planner_response=SimpleNamespace(execution_dag=dag), + return SimpleNamespace(execution_dag=dag, decomposer_response=None, artifact_refs=artifact_refs, subgraph_outputs={"sql_agent:sq1:t": object()}) diff --git a/packages/nl2sql/tests/unit/test_subgraph_resolution.py b/packages/nl2sql/tests/unit/test_subgraph_resolution.py index 4b0d8411..fb638420 100644 --- a/packages/nl2sql/tests/unit/test_subgraph_resolution.py +++ b/packages/nl2sql/tests/unit/test_subgraph_resolution.py @@ -3,11 +3,9 @@ from nl2sql.common.errors import ErrorCode, ErrorSeverity, PipelineError from nl2sql.common.exceptions import NL2SQLError, PipelineExecutionError from nl2sql.pipeline.graph_utils import resolve_subgraph -from nl2sql.pipeline.nodes.global_planner.schemas import ( - ColumnSpec, +from nl2sql.execution.dag import ( ExecutionDAG, LogicalNode, - RelationSchema, ) from nl2sql.pipeline.routes import build_scan_layer_router from nl2sql.pipeline.subgraphs.sql_agent import build_sql_agent_graph @@ -65,7 +63,6 @@ def _scan_state(datasource_id): node_id="sq_1", kind="scan", inputs=[], - output_schema=RelationSchema(columns=[ColumnSpec(name="id")]), ) dag = ExecutionDAG(nodes=[node], edges=[]) return SimpleNamespace( @@ -74,7 +71,7 @@ def _scan_state(datasource_id): datasource_resolver_response=None, artifact_refs={}, subgraph_outputs={}, - global_planner_response=SimpleNamespace(execution_dag=dag), + execution_dag=dag, decomposer_response=SimpleNamespace( sub_queries=[SimpleNamespace(id="sq_1", datasource_id=datasource_id)] ),