From a26eec0fb432410167790f869788bf9214da4af1 Mon Sep 17 00:00:00 2001 From: goulustis Date: Wed, 10 Dec 2025 22:10:42 +0800 Subject: [PATCH] stream toolnode --- lang_agent/graphs/tool_nodes.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/lang_agent/graphs/tool_nodes.py b/lang_agent/graphs/tool_nodes.py index 3657b6d..f2607f8 100644 --- a/lang_agent/graphs/tool_nodes.py +++ b/lang_agent/graphs/tool_nodes.py @@ -232,6 +232,7 @@ def debug_chatty_node(): def debug_tool_node(): + import jax from langchain_core.messages.base import BaseMessageChunk from lang_agent.components.tool_manager import ToolManagerConfig @@ -244,15 +245,13 @@ def debug_tool_node(): memory=mem) query = "use calculator to calculate 33*42" - input = ({"messages":[SystemMessage("you are a kind helper"), - HumanMessage(query)]}, - {"configurable": {"thread_id": '3'}}) - inp = { - "inp":input - } + input = ({"messages":[SystemMessage("you are a kind helper"), HumanMessage(query)]}, + {"configurable": {"thread_id": '3'}}) + graph = tool_node.tool_agent - out = tool_node.invoke(inp) - assert 0 + for chunk in graph.stream(*input, stream_mode="updates"): + el = jax.tree.leaves(chunk)[-1] + print(el.pretty_print()) if __name__ == "__main__":