Upload app.py

app.py

@@ -71,7 +71,8 @@ if "agent" not in st.session_state or st.session_state.get("current_model") != m
     if "memory" not in st.session_state:
         st.session_state.memory = AgentMemory()
         initial_message = BaseAgentOutputSchema(chat_message="Hello! I'm here to help with math problems. What can I assist you with today?")
-        st.session_state.memory.add_message("assistant", initial_message)
+        # Pass the chat_message string instead of the BaseAgentOutputSchema object
+        st.session_state.memory.add_message("assistant", initial_message.chat_message)
         st.session_state.conversation = [("assistant", initial_message.chat_message)]
     st.session_state.agent = BaseAgent(config=BaseAgentConfig(
         client=client,
@@ -103,7 +104,7 @@ if user_input:
     # Add user message to conversation and memory
     st.session_state.conversation.append(("user", user_input))
     input_schema = BaseAgentInputSchema(chat_message=user_input)
-    st.session_state.memory.add_message("user", input_schema)
+    st.session_state.memory.add_message("user", input_schema.chat_message)  # Pass the chat_message string
 
     # Display user message immediately
     with st.chat_message("user"):
@@ -122,7 +123,7 @@ if user_input:
 
             # After streaming completes, add the final response to conversation and memory
             st.session_state.conversation.append(("assistant", current_response))
-            st.session_state.memory.add_message("assistant", BaseAgentOutputSchema(chat_message=current_response))
+            st.session_state.memory.add_message("assistant", current_response)  # Pass the string directly
 
         # Run the async function
         asyncio.run(stream_response())