diff --git a/agent.py b/agent.py
index 6338280..f2afae3 100644
--- a/agent.py
+++ b/agent.py
@@ -61,10 +61,28 @@ def agent(state: State) -> Dict:
     """Ask the LLM with the current conversation and stream the answer."""
     # Build prompt from history
     messages = [HumanMessage(content=state["input"])] + state.get("messages", [])
-    # Stream response
-    for chunk in llm.stream(messages):
-        # Yield each token as a partial AI message
-        yield {"partial": chunk.content}
+    # Stream response using stream_mode to get both messages and updates
+    stream = llm.stream(messages, stream_mode=["messages", "updates"])
+    step = 1
+    def format_chunk_message(chunk):
+        message, meta = chunk
+        nonlocal step
+        if meta.get("langgraph_step") != step:
+            step = meta.get("langgraph_step")
+            print("\n --- --- --- \n", end="")
+        if message.content:
+            print(message.content, end="", flush=True)
+    def format_message(message):
+        if message.content:
+            return message.content
+        return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})"
+    for chunk_type, chunk_data in stream:
+        if chunk_type == "messages":
+            format_chunk_message(chunk_data)
+        elif chunk_type == "updates":
+            if chunk_data.get("model"):
+                last_msg = chunk_data["model"]["messages"][-1]
+                print(format_message(last_msg), end="", flush=True)
     # After streaming, append full answer to history
     final = llm.invoke(messages)
     state["messages"] = state.get("messages", []) + [AIMessage(content=final.content)]