Spaces:
Sleeping
Sleeping
| from IPython.display import Image, display | |
| from langgraph.graph import StateGraph, END, START | |
| from langgraph.graph.state import CompiledStateGraph | |
| from graph import State, Nodes, Edges | |
| class GraphBuilder: | |
| def __init__(self): | |
| """ | |
| Initializes the GraphBuilder. | |
| """ | |
| self.nodes = Nodes() | |
| self.edges = Edges() | |
| self.compiled_graph = None | |
| def build_agent_graph(self) -> CompiledStateGraph: | |
| """Build and return the agent graph.""" | |
| graph = StateGraph(State) | |
| graph.add_node("manager", self.nodes.manager_node) | |
| graph.add_node("final_answer", self.nodes.final_answer_node) | |
| graph.add_node("auditor", self.nodes.auditor_node) | |
| graph.add_node("solver", self.nodes.solver_node) | |
| graph.add_node("researcher", self.nodes.researcher_node) | |
| graph.add_node("encryption_expert", self.nodes.encryption_expert_node) | |
| graph.add_node("encryption_advisor", self.nodes.encryption_advisor_node) | |
| graph.add_node("math_expert", self.nodes.math_expert_node) | |
| graph.add_node("math_advisor", self.nodes.math_advisor_node) | |
| graph.add_node("reasoner", self.nodes.reasoner_node) | |
| graph.add_node("image_handler", self.nodes.image_handler_node) | |
| graph.add_node("video_handler", self.nodes.video_handler_node) | |
| graph.add_edge(START, "manager") | |
| graph.add_edge("final_answer", END) | |
| graph.add_edge("auditor", "manager") | |
| graph.add_edge("researcher", "solver") | |
| graph.add_edge("encryption_advisor", "encryption_expert") | |
| graph.add_edge("math_advisor", "math_expert") | |
| graph.add_edge("reasoner", "solver") | |
| graph.add_edge("image_handler", "solver") | |
| graph.add_edge("video_handler", "solver") | |
| graph.add_conditional_edges( | |
| "manager", | |
| self.edges.manager_edge, | |
| { | |
| "solver": "solver", "auditor": "auditor", "final_answer": "final_answer" | |
| } | |
| ) | |
| graph.add_conditional_edges( | |
| "solver", | |
| self.edges.solver_edge, | |
| { | |
| "manager": "manager", "researcher": "researcher", "encryption_expert": "encryption_expert", "math_expert": "math_expert", "reasoner": "reasoner", "image_handler": "image_handler", "video_handler": "video_handler" | |
| } | |
| ) | |
| graph.add_conditional_edges( | |
| "encryption_expert", | |
| self.edges.encryption_expert_edge, | |
| { | |
| "solver": "solver", "encryption_advisor": "encryption_advisor" | |
| } | |
| ) | |
| graph.add_conditional_edges( | |
| "math_expert", | |
| self.edges.math_expert_edge, | |
| { | |
| "solver": "solver", "math_advisor": "math_advisor" | |
| } | |
| ) | |
| return graph.compile() | |
| def display_graph(self): | |
| if self.compiled_graph is None: | |
| raise ValueError("Graph has not been compiled yet. Call build_agent_graph() first.") | |
| display(Image(self.compiled_graph.get_graph().draw_mermaid_png())) | |