Update app.py

app.py

@@ -31,15 +31,12 @@ together_api_key = os.getenv("TOGETHER_API_KEY")
  
 # LangGraph State
 class State(TypedDict):
-    messages: Annotated[list, add_messages]
-    name: str
-    birthday: str
+    messages: Annotated[list, add_messages]    
     input : str
+    questions : List[str]
+    answers:List[str]
     code : str
     explanation:str
-    execution_result=str
-    
- 
 # LLM
 code_generator = Together(
     model="deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free",  
@@ -51,6 +48,33 @@ code_generator = Together(
 # Memory
 memory = MemorySaver()
 
+def generate_questions(state: State):
+    user_input = state["messages"][-1].content
+    result = subtask_chain.invoke({"user_goal": user_input})
+    text = result["text"]
+    subtasks, questions = parse_subtasks_and_questions(text)
+
+    follow_up = "\n".join(f"Q{i+1}: {q}" for i, q in enumerate(questions))
+    state["messages"].append(AIMessage(content="To proceed, please answer these questions:\n" + follow_up))
+
+    return {
+        "messages": state["messages"],
+        "questions": questions,
+        "answers": [],  # Wait for user
+        "next": "wait_for_answers"
+    }
+if "questions" in state_input and state_input["questions"]:
+    for q in state_input["questions"]:
+        st.text_input(q, key=q)  # Store answers in keys
+
+    submitted = st.form_submit_button("Submit Answers")
+    if submitted:
+        state_input["answers"] = [st.session_state[q] for q in state_input["questions"]]
+        state_input["questions"] = []  # Clear questions
+def handle_answers(state: State):
+    full_input = state["input"] + "\n\n" + "\n".join(state["answers"])
+    return {**state, "input": full_input}
+
 # Define node
 def ai_assistance(state: State):
     result =code_generator.invoke(state["messages"])
@@ -273,21 +297,24 @@ def router(state):
         return "AI_Assistance"
 builder = StateGraph(State)
 
-# Define Nodes
+# Define All Nodes
 builder.add_node("AI_Assistance", ai_assistance)
-builder.add_node("LLM_Agent", agent_node)
+builder.add_node("Generate_Questions", generate_questions)
+builder.add_node("Handle_Answers", handle_answers)
 builder.add_node("Generate_Code", generate_code)
 builder.add_node("Code_Explainer", explain_code)
 
-# Entry Point
+# Set Entry Point
 builder.set_entry_point("AI_Assistance")
 
 # Define Flow
-builder.add_edge("AI_Assistance", "LLM_Agent")
-builder.add_edge("LLM_Agent", "Generate_Code")
+builder.add_edge("AI_Assistance", "Generate_Questions")
+builder.add_edge("Generate_Questions", "Handle_Answers")
+builder.add_edge("Handle_Answers", "Generate_Code")
 builder.add_edge("Generate_Code", "Code_Explainer")
 builder.add_edge("Code_Explainer", END)
 
+
 graph = builder.compile(checkpointer=memory) 
 # Streamlit UI setup
 st.set_page_config(page_title="MitraVerse", layout="wide")