Update agent.py

agent.py

@@ -148,4 +148,48 @@ def files_condition(state: AgentState) -> str:
     elif state["file_name"]=="":
         return "reformulate_question"
     else:
-        return "not_handled"
+        return "not_handled"
+
+
+def build_agent()
+    ## The graph
+    builder = StateGraph(AgentState)
+    
+    # Define nodes: these do the work
+    builder.add_node("get_question", get_question)
+    builder.add_node("reformulate_question", reformulate_question)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_node("get_final_answer", get_final_answer)
+    builder.add_node("python_interpreter", python_interpreter)
+    builder.add_node("image_interpreter", image_interpreter)
+    builder.add_node("not_handled", not_handled)
+    
+    # Define edges: these determine how the control flow moves
+    builder.add_edge(START, "get_question")
+    builder.add_conditional_edges(
+        "get_question",
+        files_condition,
+        {
+            "python_interpreter": "python_interpreter",
+            "image_interpreter": "image_interpreter",
+            "reformulate_question": "reformulate_question",
+            "not_handled": "not_handled"
+        }
+    )
+    builder.add_edge("python_interpreter", "get_final_answer")
+    builder.add_edge("image_interpreter", "get_final_answer")
+    builder.add_edge("reformulate_question", "assistant")
+    builder.add_conditional_edges(
+        "assistant",
+        tools_condition,
+        {
+            "tools": "tools",
+            "get_final_answer": "get_final_answer"
+        }
+    )
+    builder.add_edge("tools", "assistant")
+    builder.add_edge("get_final_answer", END)
+    builder.add_edge("not_handled", END)
+    mySuperAgent = builder.compile()
+    return mySuperAgent