| from __future__ import annotations |
|
|
| from langgraph.graph import END, START, StateGraph |
|
|
| from graph.nodes import WorkflowNodes |
| from graph.routes import ( |
| route_after_analysis, |
| route_after_planner, |
| route_after_reflection, |
| route_after_sql, |
| route_after_visualization, |
| ) |
| from graph.state import BIState |
|
|
|
|
| class BIWorkflow: |
| def __init__(self, nodes: WorkflowNodes) -> None: |
| self.nodes = nodes |
| self.graph = self._build_graph() |
|
|
| def _build_graph(self): |
| graph = StateGraph(BIState) |
| graph.add_node("planner", self.nodes.run_planner) |
| graph.add_node("sql", self.nodes.run_sql) |
| graph.add_node("analysis", self.nodes.run_analysis) |
| graph.add_node("reflection", self.nodes.run_reflection) |
| graph.add_node("visualization", self.nodes.run_visualization) |
|
|
| graph.add_edge(START, "planner") |
| graph.add_conditional_edges("planner", route_after_planner, {"sql": "sql", "analysis": "analysis"}) |
| graph.add_conditional_edges("sql", route_after_sql, {"analysis": "analysis"}) |
| graph.add_conditional_edges("analysis", route_after_analysis, {"reflection": "reflection"}) |
| graph.add_conditional_edges("reflection", route_after_reflection, {"visualization": "visualization"}) |
| graph.add_conditional_edges("visualization", route_after_visualization, {END: END}) |
| return graph.compile() |
|
|
| def run(self, state: BIState) -> BIState: |
| return self.graph.invoke(state) |
|
|
