From 03b7d6fe0da7b49fcecd73d26e49cbfaa77647b0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 30 Mar 2024 22:04:27 -0700 Subject: [PATCH] Remove __end__ node from state graph No longer necessary with stream output modes --- langgraph/graph/graph.py | 48 ++-- langgraph/graph/state.py | 23 +- tests/__snapshots__/test_pregel.ambr | 332 +++++++++++++-------------- tests/test_pregel.py | 38 +-- tests/test_pregel_async.py | 22 +- 5 files changed, 226 insertions(+), 237 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index b21b6f0b3..c07c7402a 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -30,25 +30,32 @@ class Branch(NamedTuple): condition: Runnable[Any, str] ends: Optional[dict[str, str]] - @property - def runnable(self): - return RunnableLambda(self._route, self._aroute, name=self.condition.name) + def run(self, writer: Callable[[str], Optional[Runnable]]) -> None: + return RunnableLambda( + self._route, + self._aroute, + name=self.condition.name, + ).bind(writer=writer) - def _route(self, input: Any) -> Runnable: + def _route( + self, input: Any, *, writer: Callable[[str], Optional[Runnable]] + ) -> Runnable: result = self.condition.invoke(input, {"run_name": "condition"}) if self.ends: destination = self.ends[result] else: destination = result - return Channel.write_to(f"{destination}:inbox" if destination != END else END) + return writer(destination) - async def _aroute(self, input: Any) -> Runnable: + async def _aroute( + self, input: Any, *, writer: Callable[[str], Optional[Runnable]] + ) -> Runnable: result = await self.condition.ainvoke(input, {"run_name": "condition"}) if self.ends: destination = self.ends[result] else: destination = result - return Channel.write_to(f"{destination}:inbox" if destination != END else END) + return writer(destination) class Graph: @@ -104,7 +111,7 @@ class Graph: condition: Union[ Callable[..., str], Callable[..., Awaitable[str]], Runnable[Any, str] ], - conditional_edge_mapping: Optional[Dict[str, str]] = None, + conditional_edge_mapping: Optional[dict[str, str]] = None, ) -> None: if self.compiled: logger.warning( @@ -224,6 +231,9 @@ class Graph: for key in self.nodes } + def branch_writer(dest: str) -> Optional[Runnable]: + return Channel.write_to(f"{dest}:inbox" if dest != END else END) + for key in self.nodes: outgoing = outgoing_edges[key] edges_key = f"{key}:edges" @@ -233,13 +243,12 @@ class Graph: nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing]) if key in self.branches: for branch in self.branches[key]: - nodes[edges_key] |= branch.runnable + nodes[edges_key] |= branch.run(branch_writer) if self.entry_point_branch: - nodes[f"{START}:edges"] = ( - Channel.subscribe_to(START, tags=["langsmith:hidden"]) - | self.entry_point_branch.runnable - ) + nodes[f"{START}:edges"] = Channel.subscribe_to( + START, tags=["langsmith:hidden"] + ) | self.entry_point_branch.run(branch_writer) elif self.entry_point is None: raise ValueError("No entry point set") @@ -296,10 +305,10 @@ class CompiledGraph(Pregel): graph.add_edge(start_nodes[start], end_nodes[end]) for start, branches in self.graph.branches.items(): for i, branch in enumerate(branches): - name = f"{start}_{branch.runnable.name}" + name = f"{start}_{branch.condition.name or 'condition'}" if i > 0: name += f"_{i}" - cond = graph.add_node(branch.runnable, name) + cond = graph.add_node(branch.condition, name) graph.add_edge(start_nodes[start], cond) ends = branch.ends or { **{k: k for k in self.graph.nodes}, @@ -307,14 +316,13 @@ class CompiledGraph(Pregel): } for label, end in ends.items(): graph.add_edge(cond, end_nodes[end], label) - if self.graph.entry_point_branch: + if entry_point_branch := self.graph.entry_point_branch: cond = graph.add_node( - self.graph.entry_point_branch.runnable, f"{START}_condition" + entry_point_branch.condition, + entry_point_branch.condition.name or f"{START}_condition", ) graph.add_edge(start_nodes[START], cond) - ends = self.graph.entry_point_branch.ends or { - k: k for k in self.graph.nodes - } + ends = entry_point_branch.ends or {k: k for k in self.graph.nodes} for label, end in ends.items(): graph.add_edge(cond, end_nodes[end], label) elif self.graph.entry_point: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index b61581736..243bc0e0c 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -109,7 +109,8 @@ class StateGraph(Graph): outgoing_edges = defaultdict(list) for start, end in self.edges: - outgoing_edges[start].append(f"{end}:inbox" if end != END else END) + if end != END: + outgoing_edges[start].append(f"{end}:inbox") for key, starts, end in waiting_edges: for start in starts: outgoing_edges[start].append(key) @@ -143,6 +144,12 @@ class StateGraph(Graph): for key in list(self.nodes) + [START] } + def branch_writer(src: str, dest: str) -> Optional[ChannelWrite]: + if dest != END: + return ChannelWrite( + channels=[ChannelWriteEntry(f"{dest}:inbox", src, False)] + ) + for key in self.nodes: outgoing = outgoing_edges[key] edges_key = f"{key}:edges" @@ -152,14 +159,11 @@ class StateGraph(Graph): ) if outgoing: nodes[edges_key] |= ChannelWrite( - channels=[ - ChannelWriteEntry(dest, None if dest == END else key, True) - for dest in outgoing - ] + channels=[ChannelWriteEntry(dest, key, True) for dest in outgoing] ) if key in self.branches: for branch in self.branches[key]: - nodes[edges_key] |= branch.runnable + nodes[edges_key] |= branch.run(partial(branch_writer, key)) nodes[START] = Channel.subscribe_to( f"{START}:inbox", tags=["langsmith:hidden"] @@ -172,7 +176,9 @@ class StateGraph(Graph): if self.entry_point: nodes[f"{START}:edges"] |= Channel.write_to(f"{self.entry_point}:inbox") elif self.entry_point_branch: - nodes[f"{START}:edges"] |= self.entry_point_branch.runnable + nodes[f"{START}:edges"] |= self.entry_point_branch.run( + partial(branch_writer, START) + ) else: raise ValueError("No entry point set") @@ -184,11 +190,10 @@ class StateGraph(Graph): **node_inboxes, **node_outboxes, **waiting_edge_channels, - END: LastValue(self.schema), }, input_channels=f"{START}:inbox", stream_mode="updates", - output_channels=END, + output_channels=state_keys_read, stream_channels=state_keys_read, checkpointer=checkpointer, interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 7ae35f8cd..baded0215 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -50,7 +50,7 @@ } }, { - "id": "left__route", + "id": "left_condition", "type": "runnable", "data": { "id": [ @@ -59,11 +59,11 @@ "base", "RunnableLambda" ], - "name": "_route" + "name": "RunnableLambda" } }, { - "id": "__start___condition", + "id": "should_start", "type": "runnable", "data": { "id": [ @@ -83,34 +83,34 @@ }, { "source": "left", - "target": "left__route" + "target": "left_condition" }, { - "source": "left__route", + "source": "left_condition", "target": "left", "data": "left" }, { - "source": "left__route", + "source": "left_condition", "target": "right", "data": "right" }, { - "source": "left__route", + "source": "left_condition", "target": "__end__", "data": "__end__" }, { "source": "__start__", - "target": "__start___condition" + "target": "should_start" }, { - "source": "__start___condition", + "source": "should_start", "target": "left", "data": "go-left" }, { - "source": "__start___condition", + "source": "should_start", "target": "right", "data": "go-right" } @@ -120,46 +120,46 @@ # --- # name: test_conditional_entrypoint_graph.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------------+ - | __start___condition | - +---------------------+ - *** *** - * * - ** *** - +------+ * - | left | * - +------+ * - * * - * * - * * - +-------------+ * - | left__route | * - +-------------+** * - * **** * - * **** * - * ** * - * +-------+ - ** | right | - ** +-------+ - ** *** - ** * - * ** - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +--------------+ + | should_start | + +--------------+ + *** ** + * ** + ** ** + +------+ * + | left | * + +------+ * + * * + * * + * * + +----------------+ * + | left_condition | * + +----------------+* * + * ***** * + * *** * + * *** * + ** +-------+ + * | right | + *** +-------+ + * *** + *** * + * ** + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_conditional_entrypoint_graph_state '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}}}}}' # --- # name: test_conditional_entrypoint_graph_state.1 - '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}}}}}' + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}}}' # --- # name: test_conditional_entrypoint_graph_state.2 ''' @@ -194,21 +194,15 @@ "type": "schema", "data": { "title": "LangGraphOutput", - "$ref": "#/definitions/AgentState", - "definitions": { - "AgentState": { - "title": "AgentState", - "type": "object", - "properties": { - "input": { - "title": "Input", - "type": "string" - }, - "output": { - "title": "Output", - "type": "string" - } - } + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "output": { + "title": "Output", + "type": "string" } } } @@ -240,7 +234,7 @@ } }, { - "id": "left__route", + "id": "left_condition", "type": "runnable", "data": { "id": [ @@ -249,11 +243,11 @@ "base", "RunnableLambda" ], - "name": "_route" + "name": "RunnableLambda" } }, { - "id": "__start___condition", + "id": "should_start", "type": "runnable", "data": { "id": [ @@ -273,34 +267,34 @@ }, { "source": "left", - "target": "left__route" + "target": "left_condition" }, { - "source": "left__route", + "source": "left_condition", "target": "left", "data": "left" }, { - "source": "left__route", + "source": "left_condition", "target": "right", "data": "right" }, { - "source": "left__route", + "source": "left_condition", "target": "__end__", "data": "__end__" }, { "source": "__start__", - "target": "__start___condition" + "target": "should_start" }, { - "source": "__start___condition", + "source": "should_start", "target": "left", "data": "go-left" }, { - "source": "__start___condition", + "source": "should_start", "target": "right", "data": "go-right" } @@ -310,39 +304,39 @@ # --- # name: test_conditional_entrypoint_graph_state.3 ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------------+ - | __start___condition | - +---------------------+ - *** *** - * * - ** *** - +------+ * - | left | * - +------+ * - * * - * * - * * - +-------------+ * - | left__route | * - +-------------+** * - * **** * - * **** * - * ** * - * +-------+ - ** | right | - ** +-------+ - ** *** - ** * - * ** - +---------+ - | __end__ | - +---------+ + +-----------+ + | __start__ | + +-----------+ + * + * + * + +--------------+ + | should_start | + +--------------+ + *** ** + * ** + ** ** + +------+ * + | left | * + +------+ * + * * + * * + * * + +----------------+ * + | left_condition | * + +----------------+* * + * ***** * + * *** * + * *** * + ** +-------+ + * | right | + *** +-------+ + * *** + *** * + * ** + +---------+ + | __end__ | + +---------+ ''' # --- # name: test_conditional_graph @@ -754,7 +748,7 @@ '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' # --- # name: test_conditional_graph_state.1 - '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- # name: test_conditional_graph_state.2 ''' @@ -879,7 +873,41 @@ "type": "schema", "data": { "title": "LangGraphOutput", - "$ref": "#/definitions/AgentState", + "type": "object", + "properties": { + "input": { + "title": "Input", + "type": "string" + }, + "agent_outcome": { + "title": "Agent Outcome", + "anyOf": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "$ref": "#/definitions/AgentFinish" + } + ] + }, + "intermediate_steps": { + "title": "Intermediate Steps", + "type": "array", + "items": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": [ + { + "$ref": "#/definitions/AgentAction" + }, + { + "type": "string" + } + ] + } + } + }, "definitions": { "AgentAction": { "title": "AgentAction", @@ -946,44 +974,6 @@ "return_values", "log" ] - }, - "AgentState": { - "title": "AgentState", - "type": "object", - "properties": { - "input": { - "title": "Input", - "type": "string" - }, - "agent_outcome": { - "title": "Agent Outcome", - "anyOf": [ - { - "$ref": "#/definitions/AgentAction" - }, - { - "$ref": "#/definitions/AgentFinish" - } - ] - }, - "intermediate_steps": { - "title": "Intermediate Steps", - "type": "array", - "items": { - "type": "array", - "minItems": 2, - "maxItems": 2, - "items": [ - { - "$ref": "#/definitions/AgentAction" - }, - { - "type": "string" - } - ] - } - } - } } } } @@ -1903,7 +1893,7 @@ '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' # --- # name: test_prebuilt_chat.1 - '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' + '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' # --- # name: test_prebuilt_chat.2 ''' @@ -1988,7 +1978,16 @@ "type": "schema", "data": { "title": "LangGraphOutput", - "$ref": "#/definitions/AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, "definitions": { "BaseMessage": { "title": "BaseMessage", @@ -2037,22 +2036,6 @@ "content", "type" ] - }, - "AgentState": { - "title": "AgentState", - "type": "object", - "properties": { - "messages": { - "title": "Messages", - "type": "array", - "items": { - "$ref": "#/definitions/BaseMessage" - } - } - }, - "required": [ - "messages" - ] } } } @@ -2153,7 +2136,7 @@ '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' # --- # name: test_prebuilt_tool_chat.1 - '{"title": "LangGraphOutput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' + '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' # --- # name: test_prebuilt_tool_chat.2 ''' @@ -2238,7 +2221,16 @@ "type": "schema", "data": { "title": "LangGraphOutput", - "$ref": "#/definitions/AgentState", + "type": "object", + "properties": { + "messages": { + "title": "Messages", + "type": "array", + "items": { + "$ref": "#/definitions/BaseMessage" + } + } + }, "definitions": { "BaseMessage": { "title": "BaseMessage", @@ -2287,22 +2279,6 @@ "content", "type" ] - }, - "AgentState": { - "title": "AgentState", - "type": "object", - "properties": { - "messages": { - "title": "Messages", - "type": "array", - "items": { - "$ref": "#/definitions/BaseMessage" - } - } - }, - "required": [ - "messages" - ] } } } diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 3a26a2c67..d0562a271 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2234,7 +2234,7 @@ def test_message_graph( "__start__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000038", + id="00000000-0000-4000-8000-000000000037", ) ] }, @@ -2251,7 +2251,7 @@ def test_message_graph( "action": FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000051", + id="00000000-0000-4000-8000-000000000050", ) }, { @@ -2267,7 +2267,7 @@ def test_message_graph( "action": FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000064", + id="00000000-0000-4000-8000-000000000063", ) }, {"agent": AIMessage(content="answer", id="ai3")}, @@ -2287,7 +2287,7 @@ def test_message_graph( { "__start__": HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ) }, { @@ -2305,7 +2305,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2329,7 +2329,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2351,7 +2351,7 @@ def test_message_graph( "action": FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ) }, { @@ -2369,7 +2369,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2384,7 +2384,7 @@ def test_message_graph( FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ), AIMessage( content="", @@ -2408,7 +2408,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2423,7 +2423,7 @@ def test_message_graph( FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ), AIMessage(content="answer", id="ai2"), ], @@ -2448,7 +2448,7 @@ def test_message_graph( { "__start__": HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", + id="00000000-0000-4000-8000-000000000096", ) }, { @@ -2466,7 +2466,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", + id="00000000-0000-4000-8000-000000000096", ), AIMessage( content="", @@ -2490,7 +2490,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", + id="00000000-0000-4000-8000-000000000096", ), AIMessage( content="", @@ -2512,7 +2512,7 @@ def test_message_graph( "action": FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000116", + id="00000000-0000-4000-8000-000000000113", ) }, { @@ -2530,7 +2530,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", + id="00000000-0000-4000-8000-000000000096", ), AIMessage( content="", @@ -2545,7 +2545,7 @@ def test_message_graph( FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000116", + id="00000000-0000-4000-8000-000000000113", ), AIMessage( content="", @@ -2569,7 +2569,7 @@ def test_message_graph( values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", + id="00000000-0000-4000-8000-000000000096", ), AIMessage( content="", @@ -2584,7 +2584,7 @@ def test_message_graph( FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000116", + id="00000000-0000-4000-8000-000000000113", ), AIMessage(content="answer", id="ai2"), ], diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index c969f5261..1ba1249d7 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2243,7 +2243,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: "__start__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000038", + id="00000000-0000-4000-8000-000000000037", ) ] }, @@ -2260,7 +2260,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: "action": FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000051", + id="00000000-0000-4000-8000-000000000050", ) }, { @@ -2276,7 +2276,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: "action": FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000064", + id="00000000-0000-4000-8000-000000000063", ) }, {"agent": AIMessage(content="answer", id="ai3")}, @@ -2296,7 +2296,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: { "__start__": HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ) }, { @@ -2314,7 +2314,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2338,7 +2338,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2360,7 +2360,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: "action": FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ) }, { @@ -2378,7 +2378,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2393,7 +2393,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ), AIMessage( content="", @@ -2417,7 +2417,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: values=[ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", + id="00000000-0000-4000-8000-000000000072", ), AIMessage( content="", @@ -2432,7 +2432,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: FunctionMessage( content="result for a different query", name="search_api", - id="00000000-0000-4000-8000-000000000088", + id="00000000-0000-4000-8000-000000000086", ), AIMessage(content="answer", id="ai2"), ],