Spaces:
Sleeping
Sleeping
File size: 3,114 Bytes
910ae58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
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()))
|