Improve tracing output of state graph

langgraph/graph/state.py

@@ -3,7 +3,7 @@ from functools import partial
 from inspect import signature
 from typing import Any, Optional, Sequence, Type, Union
 
-from langchain_core.runnables import Runnable, RunnableLambda
+from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.base import RunnableLike
 
 from langgraph.channels.base import BaseChannel, InvalidUpdateError
@@ -16,6 +16,7 @@ from langgraph.constants import TAG_HIDDEN
 from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
+from langgraph.utils import RunnableCallable
 
 logger = logging.getLogger(__name__)
 
@@ -123,7 +124,7 @@ class CompiledStateGraph(CompiledGraph):
     graph: StateGraph
 
     def attach_node(self, key: str, node: Optional[Runnable]) -> None:
-        def _get_state_key(key: str, input: dict) -> Any:
+        def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any:
             if input is None:
                 return SKIP_WRITE
             elif not isinstance(input, dict):
@@ -136,7 +137,9 @@ class CompiledStateGraph(CompiledGraph):
         state_write_entries = [
             ChannelWriteEntry(key, None, skip_none=True)
             if key == "__root__"
-            else ChannelWriteEntry(key, RunnableLambda(partial(_get_state_key, key)))
+            else ChannelWriteEntry(
+                key, RunnableCallable(_get_state_key, key=key, trace=False)
+            )
             for key in state_keys
         ]
         # node that reads current state with (this node's) updates applied

langgraph/prebuilt/agent_executor.py

@@ -3,10 +3,10 @@ from typing import Annotated, Sequence, TypedDict, Union
 
 from langchain_core.agents import AgentAction, AgentFinish
 from langchain_core.messages import BaseMessage
-from langchain_core.runnables import RunnableLambda
 
 from langgraph.graph import END, StateGraph
 from langgraph.prebuilt.tool_executor import ToolExecutor
+from langgraph.utils import RunnableCallable
 
 
 def _get_agent_state(input_schema=None):
@@ -72,35 +72,33 @@ def create_agent_executor(agent_runnable, tools, input_schema=None):
     def execute_tools(data):
         # Get the most recent agent_outcome - this is the key added in the `agent` above
         agent_action = data["agent_outcome"]
-        if isinstance(agent_action, list):
-            output = tool_executor.batch(agent_action, return_exceptions=True)
-            return {
-                "intermediate_steps": [
-                    (action, str(out)) for action, out in zip(agent_action, output)
-                ]
-            }
-        output = tool_executor.invoke(agent_action)
-        return {"intermediate_steps": [(agent_action, str(output))]}
+        if not isinstance(agent_action, list):
+            agent_action = [agent_action]
+        output = tool_executor.batch(agent_action, return_exceptions=True)
+        return {
+            "intermediate_steps": [
+                (action, str(out)) for action, out in zip(agent_action, output)
+            ]
+        }
 
     async def aexecute_tools(data):
         # Get the most recent agent_outcome - this is the key added in the `agent` above
         agent_action = data["agent_outcome"]
-        if isinstance(agent_action, list):
-            output = await tool_executor.abatch(agent_action, return_exceptions=True)
-            return {
-                "intermediate_steps": [
-                    (action, str(out)) for action, out in zip(agent_action, output)
-                ]
-            }
-        output = await tool_executor.ainvoke(agent_action)
-        return {"intermediate_steps": [(agent_action, str(output))]}
+        if not isinstance(agent_action, list):
+            agent_action = [agent_action]
+        output = await tool_executor.abatch(agent_action, return_exceptions=True)
+        return {
+            "intermediate_steps": [
+                (action, str(out)) for action, out in zip(agent_action, output)
+            ]
+        }
 
     # Define a new graph
     workflow = StateGraph(state)
 
     # Define the two nodes we will cycle between
-    workflow.add_node("agent", RunnableLambda(run_agent, arun_agent))
-    workflow.add_node("action", RunnableLambda(execute_tools, aexecute_tools))
+    workflow.add_node("agent", RunnableCallable(run_agent, arun_agent))
+    workflow.add_node("action", RunnableCallable(execute_tools, aexecute_tools))
 
     # Set the entrypoint as `agent`
     # This means that this node is the first one called

langgraph/utils.py

@@ -18,19 +18,28 @@ class RunnableCallable(Runnable):
     def __init__(
         self,
         func: Callable[..., Optional[Runnable]],
-        afunc: Callable[..., Awaitable[Optional[Runnable]]],
-        name: str,
+        afunc: Optional[Callable[..., Awaitable[Optional[Runnable]]]] = None,
+        *,
+        name: Optional[str] = None,
         tags: Optional[list[str]] = None,
         trace: bool = True,
         **kwargs: Any,
     ) -> None:
-        self.name = name
+        self.name = name or func.__name__
         self.func = func
         self.afunc = afunc
         self.config = {"tags": tags} if tags else None
         self.kwargs = kwargs
         self.trace = trace
 
+    def __repr__(self) -> str:
+        repr_args = {
+            k: v
+            for k, v in self.__dict__.items()
+            if k not in {"name", "func", "afunc", "config", "kwargs", "trace"}
+        }
+        return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})"
+
     def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
         if self.trace:
             ret = self._call_with_config(
@@ -43,6 +52,8 @@ class RunnableCallable(Runnable):
         return ret
 
     async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any:
+        if not self.afunc:
+            return self.invoke(input, config)
         if self.trace:
             ret = await self._acall_with_config(
                 self.afunc, input, merge_configs(self.config, config), **self.kwargs