Update app.py

app.py

@@ -37,6 +37,8 @@ class State(TypedDict):
     answers:List[str]
     code : str
     explanation:str
+    subtasks: List[str]
+    follow_up_questions: List[str]
 # LLM
 code_generator = Together(
     model="deepseek-ai/DeepSeek-R1-Distill-Llama-70B-free",  
@@ -298,23 +300,24 @@ def router(state):
 builder = StateGraph(State)
 
 # 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)
 
 # Set Entry Point
-builder.set_entry_point("AI_Assistance")
+builder.set_entry_point("LLM_Agent")
 
 # Define Flow
-builder.add_edge("AI_Assistance", "Generate_Questions")
+builder.add_edge("LLM_Agent", "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")
@@ -374,36 +377,77 @@ for msg in st.session_state.chat_history:
     st.markdown(f"<div class='stChatMessage {role}'>{msg.content}</div>", unsafe_allow_html=True) 
 
 with st.container():
+    # First Form: Ask Question
     with st.form("chat_form", clear_on_submit=True):
-        st.markdown('<div id="floating-container">', unsafe_allow_html=True)        
-        st.markdown('</div>', unsafe_allow_html=True) 
-
         user_input = st.text_input("Ask me", label_visibility="collapsed", placeholder="Ask me Anything")
         submitted = st.form_submit_button(label="Send")
-        
+
         if submitted and user_input:
             st.session_state.chat_history.append(HumanMessage(content=user_input))
-        
             config = {"configurable": {"thread_id": st.session_state.thread_id}}
             state_input = {
                 "messages": st.session_state.chat_history,
                 "input": user_input,
             }
-        
+
             result = graph.invoke(state_input, config=config)
-            
+            st.session_state.last_result = result  # store temporarily
+
+            if result.get("questions"):
+                st.session_state.pending_questions = result["questions"]
+            else:
+                st.session_state.pending_questions = []
+
             if result.get("code"):
                 st.session_state.latest_code = result["code"]
                 st.session_state.chat_history.append(
                     AIMessage(content="**💻 Generated Code:**\n\n```python\n" + result["code"] + "\n```")
-                )            
+                )
+
             if result.get("explanation"):
                 st.session_state.latest_explanation = result["explanation"]
                 st.session_state.chat_history.append(
                     AIMessage(content="**🔍 Code Explanation:**\n\n```\n" + result["explanation"] + "\n```")
-                )            
-            elif result.get("messages"):
+                )
+
+            if result.get("subtasks"):
+                st.markdown("### 🧩 Subtasks:")
+                for s in result["subtasks"]:
+                    st.markdown(f"- {s}")
+
+            if result.get("messages"):
                 st.session_state.chat_history.append(result["messages"][-1])
+
+    # Second Form: Answer Follow-up Questions
+    if "pending_questions" in st.session_state and st.session_state.pending_questions:
+        with st.form("answer_form"):
+            st.markdown("### 🤖 Please answer the following:")
+            for q in st.session_state.pending_questions:
+                st.text_input(q, key=q)
+            submit_answers = st.form_submit_button("Submit Answers")
+
+            if submit_answers:
+                answers = [st.session_state[q] for q in st.session_state.pending_questions]
+                new_input = {
+                    "messages": st.session_state.chat_history,
+                    "answers": answers,
+                    "questions": [],
+                }
+                result = graph.invoke(new_input, config={"configurable": {"thread_id": st.session_state.thread_id}})
+                st.session_state.pending_questions = []
+
+                if result.get("code"):
+                    st.session_state.latest_code = result["code"]
+                    st.session_state.chat_history.append(
+                        AIMessage(content="**💻 Generated Code:**\n\n```python\n" + result["code"] + "\n```")
+                    )
+
+                if result.get("explanation"):
+                    st.session_state.latest_explanation = result["explanation"]
+                    st.session_state.chat_history.append(
+                        AIMessage(content="**🔍 Code Explanation:**\n\n```\n" + result["explanation"] + "\n```")
+                    )
+              
             st.rerun()
             st.markdown("<script>window.scrollTo(0, document.body.scrollHeight);</script>", unsafe_allow_html=True)