Update new_app.py

new_app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import uuid
 from langgraph.graph import StateGraph, END, START
 from langgraph.checkpoint.memory import MemorySaver
 from typing_extensions import TypedDict
@@ -92,6 +93,9 @@ def handle_answers(state: State) -> State:
         HumanMessage(content=user_msg)
     ])
 
+    if 'messages' not in state:
+        state['messages'] =[]
+        
     state["messages"].append(response.content.strip())
     return state
 
@@ -167,23 +171,29 @@ builder.add_edge("Generate_Code", "Code_Explainer")
 builder.add_edge("Code_Explainer", END)
 
 # graph = builder.compile(checkpointer = MemorySaver())
-graph = StateGraph(State)
+# graph = StateGraph(State)
 
 # -------------------
 # Compile and Run
 # -------------------
 
-memory = MemorySaver()
-graph = builder.compile(checkpointer=memory)
+# memory = MemorySaver()
+# graph = builder.compile(checkpointer=memory)
 
-inputs = {"messages": [], "answers": [], "retry_count": 0, "code": "", "explanation": "", "questions": [], "task_plan" :"","user_input": "I want to create an agent"}
+# inputs = {"messages": [], "answers": [], "retry_count": 0, "code": "", "explanation": "", "questions": [], "task_plan" :"","user_input": "I want to create an agent"}
 
-for step in graph.stream(inputs, configurable={"thread_id": st.session_state.thread_id}):
-    for key, val in step.items():
-        print(f"\n--- {key} ---\n{val}")
+# for step in graph.stream(inputs, configurable={"thread_id": st.session_state.thread_id}):
+#     for key, val in step.items():
+#         print(f"\n--- {key} ---\n{val}")
 
 st.set_page_config(page_title="MitraVerse", layout="wide")
- 
+if "thread_id" not in st.session_state:
+    st.session_state.thread_id = str(uuid.uuid4())
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
+
+memory = MemorySaver()
+graph = builder.compile(checkpointer=memory)
 st.markdown("""
 <style>
 .stChatMessage {
@@ -229,16 +239,16 @@ st.title("🧠MitraVerse")
 # Columns for button layout
 col1, col2, col3 = st.columns(3) 
 # Initialize session
-if "thread_id" not in st.session_state:
-    st.session_state.thread_id = "1"
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = [] 
+# if "thread_id" not in st.session_state:
+#     st.session_state.thread_id = "1"
+# if "chat_history" not in st.session_state:
+#     st.session_state.chat_history = [] 
 # Show chat
 for msg in st.session_state.chat_history:
     role = "user" if isinstance(msg, HumanMessage) else "bot"
     st.markdown(f"<div class='stChatMessage {role}'>{msg.content}</div>",unsafe_allow_html=True)
 
-    user_input = st.chat_input("What would you like to build?")
+user_input = st.chat_input("What would you like to build?")
 if user_input:
     st.session_state.chat_history.append(HumanMessage(content=user_input))
 
@@ -253,7 +263,8 @@ if user_input:
         "user_input": user_input,
     }
 
-    for step in graph.stream(state_input):
+    for step in graph.stream(state_input, configurable={"thread_id": st.session_state.thread_id}):
+
         for _, state in step.items():
             # Display messages from different stages
             if "task_plan" in state and state["task_plan"]: