Update app.py

app.py

@@ -49,6 +49,18 @@ memory = MemorySaver()
 def ai_assistance(state: State):
     result =code_generator.invoke(state["messages"])
     return {"messages": state['messages']+[result]}
+
+def agent_node(state: State):
+    # Use your LLM here (e.g., Together, OpenAI, etc.)
+    model = ChatGoogleGenerativeAI(model = "gemini-2.0-flash-001").bind_tools(tools)
+    follow_up_prompt = "Break down this task into subtasks and ask follow-up questions if needed:\n\n"
+    last_user_msg = state["messages"][-1].content
+    full_prompt = follow_up_prompt + last_user_msg
+
+    response = model.invoke(full_prompt)
+
+    return {"messages": state["messages"] + [AIMessage(content=response)]}
+
     
 import time
 def generate_code(state: State):
@@ -184,7 +196,6 @@ def get_all_tools():
         execute_python_code
     ]
 
-
 def router(state):
     user_input = state["input"].lower()
     if "generate" in user_input:
@@ -304,16 +315,16 @@ col1, col2, col3 = st.columns(3)
 user_prompt = st.session_state.get("latest_code", "") or user_input  # fallback to user_input if needed
 
 with st.container():
-    if col1.button("⚙️ Run Python Code"):
-        if user_prompt:
-            with st.spinner("Executing your Python code..."):
-                result = execute_python_code.invoke({"code": user_prompt})
-                st.success("✅ Output:")
-                st.code(result, language="python")
+    # if col1.button("⚙️ Run Python Code"):
+    #     if user_prompt:
+    #         with st.spinner("Executing your Python code..."):
+    #             result = execute_python_code.invoke({"code": user_prompt})
+    #             st.success("✅ Output:")
+    #             st.code(result, language="python")
         else:
             st.warning("Please enter Python code in the input box.")
 
-    if col2.button("🌐 Web Search"):
+    if col1.button("🌐 Web Search"):
         if user_prompt:
             with st.spinner("Searching the web..."):
                 result = web_search.invoke({"query": user_prompt})
@@ -322,7 +333,7 @@ with st.container():
         else:
             st.warning("Please enter a search query.")
 
-    if col3.button("🧠 Deep Think"):
+    if col2.button("🧠 Deep Think"):
         if user_prompt:
             with st.spinner("Thinking deeply..."):
                 result = deep_think.invoke({"prompt": user_prompt})