diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index d637beaaf..3413b5a97 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -27,7 +27,7 @@ from langgraph.checkpoint import BaseCheckpointSaver from langgraph.constants import TAG_HIDDEN from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode -from langgraph.pregel.write import ChannelWrite +from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.utils import RunnableCallable logger = logging.getLogger(__name__) @@ -71,7 +71,7 @@ class Branch(NamedTuple): destinations = [self.ends[r] for r in result] else: destinations = result - return writer(destinations) + return writer(destinations) or input async def _aroute( self, @@ -90,7 +90,7 @@ class Branch(NamedTuple): destinations = [self.ends[r] for r in result] else: destinations = result - return writer(destinations) + return writer(destinations) or input class Graph: @@ -278,14 +278,16 @@ class CompiledGraph(Pregel): self.nodes[key] = ( PregelNode(channels=[], triggers=[]) | node - | Channel.write_to(key, tags=[TAG_HIDDEN]) + | ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN]) ) cast(list[str], self.stream_channels).append(key) def attach_edge(self, start: str, end: str) -> None: if end == END: # publish to end channel - self.nodes[start].writers.append(Channel.write_to(END, tags=[TAG_HIDDEN])) + self.nodes[start].writers.append( + ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN]) + ) else: # subscribe to start channel self.nodes[end].triggers.append(start) @@ -296,7 +298,9 @@ class CompiledGraph(Pregel): channels = [ f"branch:{start}:{name}:{end}" if end != END else END for end in ends ] - return Channel.write_to(*channels, tags=[TAG_HIDDEN]) + return ChannelWrite( + [ChannelWriteEntry(ch) for ch in channels], tags=[TAG_HIDDEN] + ) # add hidden start node if start == START and start not in self.nodes: diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 89ef5085e..999017ed3 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1106,7 +1106,7 @@ def _apply_writes( checkpoint["channel_versions"][chan] = max_version + 1 updated_channels.add(chan) else: - logger.warning(f"Skipping write for channel {chan} which has no readers") + logger.warning(f"Skipping write for channel '{chan}' which has no readers") # Channels that weren't updated in this step are notified of a new step for chan in channels: if chan not in updated_channels: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6f4d70128..9d8e0e971 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -947,6 +947,69 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # Check that agent (one of the nodes) has its output streamed to the logs assert "/logs/agent/streamed_output/-" in patch_paths + assert "/logs/agent:2/streamed_output/-" in patch_paths + assert "/logs/agent:3/streamed_output/-" in patch_paths + # Check that agent (one of the nodes) has its final output set in the logs + assert "/logs/agent/final_output" in patch_paths + assert "/logs/agent:2/final_output" in patch_paths + assert "/logs/agent:3/final_output" in patch_paths + assert [ + p["value"] + for log in patches + for p in log.ops + if p["path"] == "/logs/agent/final_output" + or p["path"] == "/logs/agent:2/final_output" + or p["path"] == "/logs/agent:3/final_output" + ] == [ + { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ) + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + }, + { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ), + ( + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ), + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + }, + ] # test state get/update methods with interrupt_after @@ -1534,6 +1597,38 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: }, ] + patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] + patch_paths = {op["path"] for log in patches for op in log.ops} + + # Check that agent (one of the nodes) has its output streamed to the logs + assert "/logs/agent/streamed_output/-" in patch_paths + # Check that agent (one of the ndoes) has its final output set in the logs + assert "/logs/agent/final_output" in patch_paths + assert [ + p["value"] + for log in patches + for p in log.ops + if p["path"] == "/logs/agent/final_output" + or p["path"] == "/logs/agent:2/final_output" + or p["path"] == "/logs/agent:3/final_output" + ] == [ + { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ) + }, + { + "agent_outcome": AgentAction( + tool="search_api", tool_input="another", log="tool:search_api:another" + ) + }, + { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + }, + ] + # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile(