1+ from langgraph .graph import END , StateGraph , START
2+ from dsa_agent .state import AgentState
3+ from dsa_agent .nodes import (
4+ retrieve_question ,
5+ update_question ,
6+ code_generation_in_node ,
7+ visualize_code ,
8+ complexity_analysis ,
9+ explaination_code ,
10+ no_context
11+ )
12+ from dsa_agent .edges import (
13+ new_question ,
14+ decide_to_generate_code ,
15+ decide_to_generate_visualisation ,
16+ decide_to_generate_complexity ,
17+ decide_to_generate_explanation
18+ )
19+ from langgraph .checkpoint .memory import MemorySaver
20+
21+
22+ workflow = StateGraph (AgentState )
23+
24+ workflow .add_node ("update_question" , update_question )
25+ workflow .add_node ("retrieve_question" , retrieve_question )
26+ workflow .add_node ("code_generation_in_node" , code_generation_in_node )
27+ workflow .add_node ("visualize_code" , visualize_code )
28+ workflow .add_node ("complexity_analysis" , complexity_analysis )
29+ workflow .add_node ("explaination_code" , explaination_code )
30+ workflow .add_node ("no_context" , no_context )
31+
32+ workflow .add_conditional_edges (
33+ START ,
34+ new_question ,
35+ {
36+ "update_question" : "update_question" ,
37+ "retrieve_question" : "retrieve_question" ,
38+ },
39+ )
40+
41+
42+
43+ workflow .add_edge ("retrieve_question" , "code_generation_in_node" )
44+
45+
46+
47+ workflow .add_conditional_edges (
48+ "code_generation_in_node" ,
49+ decide_to_generate_code ,
50+ {
51+ "useful" : "visualize_code" ,
52+ "code_generation_in_node" : "code_generation_in_node" ,
53+ },
54+ )
55+
56+
57+
58+ workflow .add_edge ("visualize_code" , "retrieve_question" )
59+
60+
61+ workflow .add_conditional_edges (
62+ "visualize_code" ,
63+ decide_to_generate_visualisation ,
64+ {
65+ "useful" : "complexity_analysis" ,
66+ "visualize_code" : "visualize_code" ,
67+ },
68+ )
69+
70+ workflow .add_edge ("complexity_analysis" , "retrieve_question" )
71+
72+
73+ workflow .add_conditional_edges (
74+ "complexity_analysis" ,
75+ decide_to_generate_complexity ,
76+ {
77+ "useful" : "explaination_code" ,
78+ "complexity_analysis" : "complexity_analysis" ,
79+ },
80+ )
81+
82+ workflow .add_edge ("explaination_code" , "retrieve_question" )
83+
84+
85+ workflow .add_conditional_edges (
86+ "explaination_code" ,
87+ decide_to_generate_explanation ,
88+ {
89+ "useful" : END ,
90+ "explaination_code" : "explaination_code" ,
91+ },
92+ )
93+
94+ workflow .add_edge ("update_question" , END )
95+
96+ workflow .add_edge ("no_context" , END )
97+
98+ graph = workflow .compile (checkpointer = MemorySaver ())
0 commit comments