diff --git a/main.py b/main.py index cd512db..da57a0d 100644 --- a/main.py +++ b/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"]}): - print(partial.get("partial", ""), end="") + # Stream output with separators between steps + 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!") \ No newline at end of file