From 1891d5afbb32862b3228ecb04f3bd7fa69a213f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?=
 =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= <mirzagitov03@mail.ru>
Date: Wed, 10 Jun 2026 14:19:38 +0000
Subject: [PATCH] =?UTF-8?q?=D0=9F=D1=83=D0=B1=D0=BB=D0=B8=D0=BA=D0=B0?=
 =?UTF-8?q?=D1=86=D0=B8=D1=8F=20=D0=B7=D0=B0=D0=B2=D0=B5=D1=80=D1=88=D0=B5?=
 =?UTF-8?q?=D0=BD=D0=B0:=20update=20agent.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agent.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

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)]