Update src/streamlit_app.py

src/streamlit_app.py

@@ -139,17 +139,27 @@ if "current_chat_id" in st.session_state:
             st.write(user_input)
 
         with st.chat_message("assistant"):
+            # ... (inside with st.chat_message("assistant"))
             response_placeholder = st.empty()
             full_response = ""
 
-            for message, metadata in workflow.stream(
-                {"messages": [system, HumanMessage(user_input)]},
-                config={"configurable": {"thread_id": st.session_state.current_chat_id}},
-                stream_mode="messages",
-            ):
-                if isinstance(message, AIMessage):
-                    full_response += message.content or ""
-                elif isinstance(message, ToolMessage):
-                    render_tool_message(message)
-
-            response_placeholder.markdown(full_response)
+# Use stream_mode='values' for a cleaner loop
+            for chunk in workflow.stream(
+                {"messages": [HumanMessage(user_input)]}, # System message should already be in state
+                config=set_config()):
+    # The output of a streaming node is the key of the node itself
+                     if "llmresponse" in chunk:
+                     ai_message = chunk["llmresponse"]["messages"][-1]
+                     full_response += ai_message.content
+                     response_placeholder.markdown(full_response + "▌") # Add a cursor effect
+    
+    # Handle tool calls if they appear in the stream
+                     elif "tool_node" in chunk:
+        # Since we are streaming, render the tool message right away
+        # And clear the response placeholder for the next text chunk
+                     response_placeholder.empty() 
+                     tool_messages = chunk["tool_node"]["messages"]
+                     for tool_message in tool_messages:
+                        render_tool_message(tool_message)
+
+            response_placeholder.markdown(full_response) # Final update without cursor