From c86ce019b20ee37eea166e2ac8d3db7998ef3837 Mon Sep 17 00:00:00 2001
From: goulustis <cnimaiv@gmail.com>
Date: Wed, 22 Oct 2025 16:29:45 +0800
Subject: [PATCH] update invoke

---
 lang_agent/graphs/routing.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/lang_agent/graphs/routing.py b/lang_agent/graphs/routing.py
index a680da7..56051d3 100644
--- a/lang_agent/graphs/routing.py
+++ b/lang_agent/graphs/routing.py
@@ -59,7 +59,13 @@ class RoutingGraph(GraphBase):
     def invoke(self, *nargs, as_stream:bool=False, **kwargs):
         assert len(kwargs) == 0, "due to inp assumptions"
 
-        state = self.workflow.invoke({"inp": nargs})
+        if as_stream:
+            for step in self.workflow.stream({"inp": nargs}, stream_mode="values", **kwargs):
+                step["messages"][-1].pretty_print()
+            state = step
+        else:
+            state = self.workflow.invoke({"inp": nargs})
+            
         return state["output"]
     
     def _build_modules(self):