Update app.py

app.py

@@ -21,6 +21,7 @@ from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 from langchain_core.messages import AIMessage
 from typing import List
+from langgraph.graph import StateGraph, END
 # Load environment
 load_dotenv()
 # os.environ["TAVILY_API_KEY"] = os.getenv("TAVILY_API_KEY")
@@ -66,7 +67,10 @@ def generate_questions(state: State):
         "next": "wait_for_answers"
     }
 
-
+def wait_for_answers(state: State):
+    # Just pass through until answers are submitted
+    return state
+    
 def handle_answers(state: State):
     full_input = state["input"] + "\n\n" + "\n".join(state["answers"])
     return {**state, "input": full_input}
@@ -291,26 +295,47 @@ def router(state):
         return "Generate_Code"
     else:
         return "AI_Assistance"
+
+# Define your graph builder with the state schema
 builder = StateGraph(State)
 
-# Define All Nodes
+# Add Nodes
 builder.add_node("LLM_Agent", agent_node)
+builder.add_node("AI_Assistance", ai_assistance)
 builder.add_node("Generate_Questions", generate_questions)
 builder.add_node("Handle_Answers", handle_answers)
+builder.add_node("Wait_For_Answers", wait_for_answers)  # this must be defined
 builder.add_node("Generate_Code", generate_code)
 builder.add_node("Code_Explainer", explain_code)
 
 # Set Entry Point
 builder.set_entry_point("LLM_Agent")
 
+# Define Conditional Function
+def check_if_answered(state: State) -> str:
+    if "questions" in state and state["questions"]:
+        return "answered"
+    else:
+        return "not_answered"
+
 # Define Flow
 builder.add_edge("LLM_Agent", "Generate_Questions")
-builder.add_edge("Generate_Questions", "Handle_Answers")
+
+builder.add_conditional_edges(
+    "Generate_Questions",
+    check_if_answered,
+    {
+        "answered": "Handle_Answers",
+        "not_answered": "Wait_For_Answers"
+    }
+)
+
 builder.add_edge("Handle_Answers", "Generate_Code")
 builder.add_edge("Generate_Code", "Code_Explainer")
 builder.add_edge("Code_Explainer", END)
 
-
+# Optionally: define what happens after waiting (if it's a loop)
+builder.add_edge("Wait_For_Answers", "Generate_Questions")  # retry loop
 
 graph = builder.compile(checkpointer=memory) 
 # Streamlit UI setup
@@ -423,9 +448,9 @@ with st.container():
 # ================================
 # TOOL BUTTONS SECTION
 # ================================
-col1, col2, col3 = st.columns(3)
+# col1, col2, col3 = st.columns(3)
 
-user_prompt = st.session_state.get("latest_code", "") or user_input  # fallback to user_input if needed
+# 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"):
@@ -435,23 +460,23 @@ with st.container():
     #             st.success("✅ Output:")
     #             st.code(result, language="python")
         # else:
-        #     st.warning("Please enter Python code in the input box.")
-
-    if col1.button("🌐 Web Search"):
-        if user_prompt:
-            with st.spinner("Searching the web..."):
-                result = web_search.invoke({"query": user_prompt})
-                st.success("🔎 Search Result:")
-                st.write(result)
-        else:
-            st.warning("Please enter a search query.")
-
-    if col2.button("🧠 Deep Think"):
-        if user_prompt:
-            with st.spinner("Thinking deeply..."):
-                result = deep_think.invoke({"prompt": user_prompt})
-                st.success("🧠 Reasoned Output:")
-                st.write(result)
-        else:
-            st.warning("Please enter a prompt.")
+    #     #     st.warning("Please enter Python code in the input box.")
+
+    # if col1.button("🌐 Web Search"):
+    #     if user_prompt:
+    #         with st.spinner("Searching the web..."):
+    #             result = web_search.invoke({"query": user_prompt})
+    #             st.success("🔎 Search Result:")
+    #             st.write(result)
+    #     else:
+    #         st.warning("Please enter a search query.")
+
+    # if col2.button("🧠 Deep Think"):
+    #     if user_prompt:
+    #         with st.spinner("Thinking deeply..."):
+    #             result = deep_think.invoke({"prompt": user_prompt})
+    #             st.success("🧠 Reasoned Output:")
+    #             st.write(result)
+    #     else:
+    #         st.warning("Please enter a prompt.")