Submit solution for task 699cc158d6d3a5544a3ed35b: update main.py

main.py

@@ -10,8 +10,27 @@ if __name__ == "__main__":
             continue
         if user_input.lower() in {"exit", "quit"}:
             break
-        for partial in graph.stream({"input": user_input, "messages": state["messages"]}):
+        # Stream output with separators between steps
-            print(partial.get("partial", ""), end="")
+        current_step = None
+        for chunk_type, chunk_data in graph.stream({"input": user_input, "messages": state["messages"]}, stream_mode=["messages", "updates"]):
+            if chunk_type == "messages":
+                message, meta = chunk_data
+                step = meta.get("langgraph_step")
+                if current_step is None:
+                    current_step = step
+                elif step != current_step:
+                    print("\n--- --- ---\n", end="")
+                    current_step = step
+                if message.content:
+                    print(message.content, end="")
+            elif chunk_type == "updates":
+                last_message = chunk_data.get("model", {}).get("messages", [-1])[-1]
+                # If the last message is an AIMessage with content, print it
+                try:
+                    if hasattr(last_message, "content") and last_message.content:
+                        print(last_message.content, end="")
+                except Exception:
+                    pass
         print()
         state["messages"] = graph.invoke({"input": user_input, "messages": state["messages"]})["messages"]
-    print("Goodbye!")
+    print("Goodbye!")