+
+Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "
+__start__
]):::first + agent(agent) + foo(foo) + __end__([__end__
]):::last + __start__ --> agent; + agent -.-> foo; + foo --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[memory] + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + __end__([__end__
]):::last + __start__ --> router_node; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph\3amodel_node; + normal_llm_node --> __end__; + weather_graph\3aweather_node --> __end__; + subgraph weather_graph + weather_graph\3amodel_node(model_node) + weather_graph\3aweather_node(weather_node__start__
]):::first + uno(uno) + dos(dos) + __end__([__end__
]):::last + __start__ --> uno; + uno -.-> dos; + uno -.-> subgraph\3aone; + dos --> __end__; + subgraph\3a__end__ --> __end__; + subgraph subgraph + subgraph\3aone(one) + subgraph\3atwo(two) + subgraph\3athree(three) + subgraph\3a__end__(__end__
) + subgraph\3aone -.-> subgraph\3athree; + subgraph\3aone -.-> subgraph\3atwo; + subgraph\3athree --> subgraph\3a__end__; + subgraph\3atwo --> subgraph\3a__end__; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_nested_graph + ''' + graph TD; + __start__ --> inner; + inner --> side; + side --> __end__; + + ''' +# --- +# name: test_nested_graph.1 + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + side(side) + __end__([__end__
]):::last + __start__ --> inner\3aup; + inner\3aup --> side; + side --> __end__; + subgraph inner + inner\3aup(up) + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_nested_graph_xray + dict({ + 'edges': list([ + dict({ + 'conditional': True, + 'source': '__start__', + 'target': 'tool_one', + }), + dict({ + 'conditional': True, + 'source': '__start__', + 'target': 'tool_three', + }), + dict({ + 'conditional': True, + 'source': '__start__', + 'target': 'tool_two:__start__', + }), + dict({ + 'source': 'tool_one', + 'target': '__end__', + }), + dict({ + 'source': 'tool_three', + 'target': '__end__', + }), + dict({ + 'source': 'tool_two:__end__', + 'target': '__end__', + }), + dict({ + 'conditional': True, + 'source': 'tool_two:__start__', + 'target': 'tool_two:tool_two_fast', + }), + dict({ + 'conditional': True, + 'source': 'tool_two:__start__', + 'target': 'tool_two:tool_two_slow', + }), + dict({ + 'source': 'tool_two:tool_two_fast', + 'target': 'tool_two:__end__', + }), + dict({ + 'source': 'tool_two:tool_two_slow', + 'target': 'tool_two:__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': '__start__', + }), + 'id': '__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'tool_one', + }), + 'id': 'tool_one', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'tool_three', + }), + 'id': 'tool_three', + 'type': 'runnable', + }), + dict({ + 'id': '__end__', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'tool_two:__start__', + }), + 'id': 'tool_two:__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'tool_two:tool_two_slow', + }), + 'id': 'tool_two:tool_two_slow', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'tool_two:tool_two_fast', + }), + 'id': 'tool_two:tool_two_fast', + 'type': 'runnable', + }), + dict({ + 'id': 'tool_two:__end__', + }), + ]), + }) +# --- +# name: test_nested_graph_xray.1 + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + tool_one(tool_one) + tool_three(tool_three) + __end__([__end__
]):::last + __start__ -.-> tool_one; + __start__ -.-> tool_three; + __start__ -.-> tool_two\3a__start__; + tool_one --> __end__; + tool_three --> __end__; + tool_two\3a__end__ --> __end__; + subgraph tool_two + tool_two\3a__start__(__start__
) + tool_two\3atool_two_slow(tool_two_slow) + tool_two\3atool_two_fast(tool_two_fast) + tool_two\3a__end__(__end__
) + tool_two\3a__start__ -.-> tool_two\3atool_two_fast; + tool_two\3a__start__ -.-> tool_two\3atool_two_slow; + tool_two\3atool_two_fast --> tool_two\3a__end__; + tool_two\3atool_two_slow --> tool_two\3a__end__; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_repeat_condition + ''' + graph TD; + Call\20Tool -.-> Chart\20Generator; + Call\20Tool -.-> Researcher; + Chart\20Generator -. call_tool .-> Call\20Tool; + Chart\20Generator -. continue .-> Researcher; + Chart\20Generator -. end .-> __end__; + Researcher -. call_tool .-> Call\20Tool; + Researcher -. continue .-> Chart\20Generator; + Researcher -. end .-> __end__; + __start__ --> Researcher; + Researcher -. redo .-> Researcher; + + ''' +# --- +# name: test_simple_multi_edge + ''' + graph TD; + __start__ --> up; + side --> down; + up --> down; + up --> other; + up --> side; + down --> __end__; + other --> __end__; + + ''' +# --- +# name: test_state_graph_w_config_inherited_state_keys + '{"properties": {"tools": {"items": {"type": "string"}, "title": "Tools", "type": "array"}}, "title": "Context", "type": "object"}' +# --- +# name: test_state_graph_w_config_inherited_state_keys.1 + '{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"additionalProperties": true, "type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"additionalProperties": true, "title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "title": "Agent Outcome"}, "intermediate_steps": {"items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input", "agent_outcome"], "title": "AgentState", "type": "object"}' +# --- +# name: test_state_graph_w_config_inherited_state_keys.2 + '{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"additionalProperties": true, "type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"additionalProperties": true, "title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "title": "Agent Outcome"}, "intermediate_steps": {"items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input", "agent_outcome"], "title": "AgentState", "type": "object"}' +# --- +# name: test_xray_bool + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + gp_one(gp_one) + __end__([__end__
]):::last + __start__ --> gp_one; + gp_one -. 1 .-> __end__; + gp_one -. 0 .-> gp_two\3a__start__; + gp_two\3a__end__ --> gp_one; + subgraph gp_two + gp_two\3a__start__(__start__
) + gp_two\3ap_one(p_one) + gp_two\3a__end__(__end__
) + gp_two\3a__start__ --> gp_two\3ap_one; + gp_two\3ap_one -. 1 .-> gp_two\3a__end__; + gp_two\3ap_one -. 0 .-> gp_two\3ap_two\3a__start__; + gp_two\3ap_two\3a__end__ --> gp_two\3ap_one; + subgraph p_two + gp_two\3ap_two\3a__start__(__start__
) + gp_two\3ap_two\3ac_one(c_one) + gp_two\3ap_two\3ac_two(c_two) + gp_two\3ap_two\3a__end__(__end__
) + gp_two\3ap_two\3a__start__ --> gp_two\3ap_two\3ac_one; + gp_two\3ap_two\3ac_one -. 1 .-> gp_two\3ap_two\3a__end__; + gp_two\3ap_two\3ac_one -. 0 .-> gp_two\3ap_two\3ac_two; + gp_two\3ap_two\3ac_two --> gp_two\3ap_two\3ac_one; + end + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_xray_issue + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + p_one(p_one) + __end__([__end__
]):::last + __start__ --> p_one; + p_one -. 1 .-> __end__; + p_one -. 0 .-> p_two\3a__start__; + p_two\3a__end__ --> p_one; + subgraph p_two + p_two\3a__start__(__start__
) + p_two\3ac_one(c_one) + p_two\3ac_two(c_two) + p_two\3a__end__(__end__
) + p_two\3a__start__ --> p_two\3ac_one; + p_two\3ac_one -. 1 .-> p_two\3a__end__; + p_two\3ac_one -. 0 .-> p_two\3ac_two; + p_two\3ac_two --> p_two\3ac_one; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_xray_lance + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'ask_question', + }), + dict({ + 'conditional': True, + 'source': 'answer_question', + 'target': '__end__', + }), + dict({ + 'conditional': True, + 'source': 'answer_question', + 'target': 'ask_question', + }), + dict({ + 'source': 'ask_question', + 'target': 'answer_question', + }), + ]), + 'nodes': list([ + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': '__start__', + }), + 'id': '__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'ask_question', + }), + 'id': 'ask_question', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'answer_question', + }), + 'id': 'answer_question', + 'type': 'runnable', + }), + dict({ + 'id': '__end__', + }), + ]), + }) +# --- +# name: test_xray_lance.1 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'generate_analysts', + }), + dict({ + 'source': 'conduct_interview', + 'target': 'generate_sections', + }), + dict({ + 'conditional': True, + 'source': 'generate_analysts', + 'target': 'conduct_interview', + }), + dict({ + 'source': 'generate_sections', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': '__start__', + }), + 'id': '__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'generate_analysts', + }), + 'id': 'generate_analysts', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'graph', + 'state', + 'CompiledStateGraph', + ]), + 'name': 'conduct_interview', + }), + 'id': 'conduct_interview', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'generate_sections', + }), + 'id': 'generate_sections', + 'type': 'runnable', + }), + dict({ + 'id': '__end__', + }), + ]), + }) +# --- +# name: test_xray_lance.2 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'generate_analysts', + }), + dict({ + 'source': 'conduct_interview:__end__', + 'target': 'generate_sections', + }), + dict({ + 'conditional': True, + 'source': 'generate_analysts', + 'target': 'conduct_interview:__start__', + }), + dict({ + 'source': 'generate_sections', + 'target': '__end__', + }), + dict({ + 'source': 'conduct_interview:__start__', + 'target': 'conduct_interview:ask_question', + }), + dict({ + 'conditional': True, + 'source': 'conduct_interview:answer_question', + 'target': 'conduct_interview:__end__', + }), + dict({ + 'conditional': True, + 'source': 'conduct_interview:answer_question', + 'target': 'conduct_interview:ask_question', + }), + dict({ + 'source': 'conduct_interview:ask_question', + 'target': 'conduct_interview:answer_question', + }), + ]), + 'nodes': list([ + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': '__start__', + }), + 'id': '__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'generate_analysts', + }), + 'id': 'generate_analysts', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'generate_sections', + }), + 'id': 'generate_sections', + 'type': 'runnable', + }), + dict({ + 'id': '__end__', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'conduct_interview:__start__', + }), + 'id': 'conduct_interview:__start__', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'conduct_interview:ask_question', + }), + 'id': 'conduct_interview:ask_question', + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + '_internal', + '_runnable', + 'RunnableCallable', + ]), + 'name': 'conduct_interview:answer_question', + }), + 'id': 'conduct_interview:answer_question', + 'type': 'runnable', + }), + dict({ + 'id': 'conduct_interview:__end__', + }), + ]), + }) +# --- diff --git a/langgraph/source/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/langgraph/source/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr new file mode 100644 index 0000000000000000000000000000000000000000..0832ff61edb5e7b9b8096c3f5d565d70125d0cc4 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -0,0 +1,143 @@ +# serializer version: 1 +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + qa --> __end__; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_send_react_interrupt_control[memory] + ''' + --- + config: + flowchart: + curve: linear + --- + graph TD; + __start__([__start__
]):::first + agent(agent) + foo(foo) + __end__([__end__
]):::last + __start__ --> agent; + agent -.-> foo; + foo --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- diff --git a/langgraph/source/libs/langgraph/tests/agents.py b/langgraph/source/libs/langgraph/tests/agents.py new file mode 100644 index 0000000000000000000000000000000000000000..d9e9257d928e9bcb8610c5b3600713c5519b2738 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/agents.py @@ -0,0 +1,30 @@ +from typing import Literal + +from pydantic import BaseModel + + +# define these objects to avoid importing langchain_core.agents +# and therefore avoid relying on core Pydantic version +class AgentAction(BaseModel): + """ + Represents a request to execute an action by an agent. + + The action consists of the name of the tool to execute and the input to pass + to the tool. The log is used to pass along extra information about the action. + """ + + tool: str + tool_input: str | dict + log: str + type: Literal["AgentAction"] = "AgentAction" + + +class AgentFinish(BaseModel): + """Final return value of an ActionAgent. + + Agents return an AgentFinish when they have reached a stopping condition. + """ + + return_values: dict + log: str + type: Literal["AgentFinish"] = "AgentFinish" diff --git a/langgraph/source/libs/langgraph/tests/any_int.py b/langgraph/source/libs/langgraph/tests/any_int.py new file mode 100644 index 0000000000000000000000000000000000000000..2fb2dba5564294a54fca98d244f4418e848520ca --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/any_int.py @@ -0,0 +1,6 @@ +class AnyInt(int): + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, int) diff --git a/langgraph/source/libs/langgraph/tests/any_str.py b/langgraph/source/libs/langgraph/tests/any_str.py new file mode 100644 index 0000000000000000000000000000000000000000..e4175fd8b48b063a5061efb0e738bc59dcd9549d --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/any_str.py @@ -0,0 +1,92 @@ +import re +from collections.abc import Sequence +from typing import Any + +from typing_extensions import Self + + +class AnyObject: + def __eq__(self, value): + return True + + +class FloatBetween(float): + def __new__(cls, min_value: float, max_value: float) -> Self: + return super().__new__(cls, min_value) + + def __init__(self, min_value: float, max_value: float) -> None: + super().__init__() + self.min_value = min_value + self.max_value = max_value + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, float) + and other >= self.min_value + and other <= self.max_value + ) + + def __hash__(self) -> int: + return hash((float(self), self.min_value, self.max_value)) + + +class AnyStr(str): + def __init__(self, prefix: str | re.Pattern = "") -> None: + super().__init__() + self.prefix = prefix + + def __eq__(self, other: object) -> bool: + return isinstance(other, str) and ( + other.startswith(self.prefix) + if isinstance(self.prefix, str) + else self.prefix.match(other) + ) + + def __hash__(self) -> int: + return hash((str(self), self.prefix)) + + +class AnyDict(dict): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, dict) or len(self) != len(other): + return False + for k, v in self.items(): + if kk := next((kk for kk in other if kk == k), None): + if v == other[kk]: + continue + else: + return False + else: + return True + + +class AnyVersion: + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, (str, int, float)) + + def __hash__(self) -> int: + return hash(str(self)) + + +class UnsortedSequence: + def __init__(self, *values: Any) -> None: + self.seq = values + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, Sequence) + and len(self.seq) == len(value) + and all(a in value for a in self.seq) + ) + + def __hash__(self) -> int: + return hash(frozenset(self.seq)) + + def __repr__(self) -> str: + return repr(self.seq) diff --git a/langgraph/source/libs/langgraph/tests/compose-postgres.yml b/langgraph/source/libs/langgraph/tests/compose-postgres.yml new file mode 100644 index 0000000000000000000000000000000000000000..221b35dafe48b4e1fa1cb714f5571dd3d8b6230f --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/compose-postgres.yml @@ -0,0 +1,17 @@ +name: langgraph-tests +services: + postgres-test: + image: postgres:16 + ports: + - "5442:5432" + environment: + POSTGRES_DB: postgres + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + healthcheck: + test: pg_isready -U postgres + start_period: 10s + timeout: 1s + retries: 5 + interval: 60s + start_interval: 1s diff --git a/langgraph/source/libs/langgraph/tests/compose-redis.yml b/langgraph/source/libs/langgraph/tests/compose-redis.yml new file mode 100644 index 0000000000000000000000000000000000000000..be43437095aac34831e9d947c453bfccadb4603d --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/compose-redis.yml @@ -0,0 +1,16 @@ +name: langgraph-tests +services: + redis-test: + image: redis:7-alpine + ports: + - "6379:6379" + command: redis-server --maxmemory 256mb --maxmemory-policy allkeys-lru + healthcheck: + test: redis-cli ping + start_period: 10s + timeout: 1s + retries: 5 + interval: 5s + start_interval: 1s + tmpfs: + - /data # Use tmpfs for faster testing diff --git a/langgraph/source/libs/langgraph/tests/conftest.py b/langgraph/source/libs/langgraph/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..2e2857ae3e667da3f20e9b74c9473094821d9efc --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/conftest.py @@ -0,0 +1,226 @@ +import os +from collections.abc import AsyncIterator, Iterator +from uuid import UUID + +import pytest +import redis +from langgraph.cache.base import BaseCache +from langgraph.cache.memory import InMemoryCache +from langgraph.cache.redis import RedisCache +from langgraph.cache.sqlite import SqliteCache +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.store.base import BaseStore +from pytest_mock import MockerFixture + +from langgraph.types import Durability +from tests.conftest_checkpointer import ( + _checkpointer_memory, + _checkpointer_memory_migrate_sends, + _checkpointer_postgres, + _checkpointer_postgres_aio, + _checkpointer_postgres_aio_pipe, + _checkpointer_postgres_aio_pool, + _checkpointer_postgres_pipe, + _checkpointer_postgres_pool, + _checkpointer_sqlite, + _checkpointer_sqlite_aes, + _checkpointer_sqlite_aio, +) +from tests.conftest_store import ( + _store_memory, + _store_postgres, + _store_postgres_aio, + _store_postgres_aio_pipe, + _store_postgres_aio_pool, + _store_postgres_pipe, + _store_postgres_pool, +) + +NO_DOCKER = os.getenv("NO_DOCKER", "false") == "true" + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +@pytest.fixture() +def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: + side_effect = ( + UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000) + ) + return mocker.patch("uuid.uuid4", side_effect=side_effect) + + +@pytest.fixture(params=["sync", "async", "exit"]) +def durability(request: pytest.FixtureRequest) -> Durability: + return request.param + + +@pytest.fixture( + scope="function", + params=["sqlite", "memory"] if NO_DOCKER else ["sqlite", "memory", "redis"], +) +def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: + if request.param == "sqlite": + yield SqliteCache(path=":memory:") + elif request.param == "memory": + yield InMemoryCache() + elif request.param == "redis": + # Get worker ID for parallel test isolation + worker_id = getattr(request.config, "workerinput", {}).get("workerid", "master") + + redis_client = redis.Redis( + host="localhost", port=6379, db=0, decode_responses=False + ) + # Use worker-specific prefix to avoid cache pollution between parallel tests + cache = RedisCache(redis_client, prefix=f"test:cache:{worker_id}:") + yield cache + + try: + # Only clear keys with our specific prefix + pattern = f"test:cache:{worker_id}:*" + keys = redis_client.keys(pattern) + if keys: + redis_client.delete(*keys) + except Exception: + pass + else: + raise ValueError(f"Unknown cache type: {request.param}") + + +@pytest.fixture( + scope="function", + params=["in_memory"] + if NO_DOCKER + else ["in_memory", "postgres", "postgres_pipe", "postgres_pool"], +) +def sync_store(request: pytest.FixtureRequest) -> Iterator[BaseStore]: + store_name = request.param + if store_name is None: + yield None + elif store_name == "in_memory": + with _store_memory() as store: + yield store + elif store_name == "postgres": + with _store_postgres() as store: + yield store + elif store_name == "postgres_pipe": + with _store_postgres_pipe() as store: + yield store + elif store_name == "postgres_pool": + with _store_postgres_pool() as store: + yield store + else: + raise NotImplementedError(f"Unknown store {store_name}") + + +@pytest.fixture( + scope="function", + params=["in_memory"] + if NO_DOCKER + else ["in_memory", "postgres_aio", "postgres_aio_pipe", "postgres_aio_pool"], +) +async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore]: + store_name = request.param + if store_name is None: + yield None + elif store_name == "in_memory": + with _store_memory() as store: + yield store + elif store_name == "postgres_aio": + async with _store_postgres_aio() as store: + yield store + elif store_name == "postgres_aio_pipe": + async with _store_postgres_aio_pipe() as store: + yield store + elif store_name == "postgres_aio_pool": + async with _store_postgres_aio_pool() as store: + yield store + else: + raise NotImplementedError(f"Unknown store {store_name}") + + +@pytest.fixture( + scope="function", + params=[ + "memory", + "sqlite", + "sqlite_aes", + ] + if NO_DOCKER + else [ + "memory", + "memory_migrate_sends", + "sqlite", + "sqlite_aes", + "postgres", + "postgres_pipe", + "postgres_pool", + ], +) +def sync_checkpointer( + request: pytest.FixtureRequest, +) -> Iterator[BaseCheckpointSaver]: + checkpointer_name = request.param + if checkpointer_name == "memory": + with _checkpointer_memory() as checkpointer: + yield checkpointer + elif checkpointer_name == "memory_migrate_sends": + with _checkpointer_memory_migrate_sends() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite": + with _checkpointer_sqlite() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite_aes": + with _checkpointer_sqlite_aes() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres": + with _checkpointer_postgres() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_pipe": + with _checkpointer_postgres_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_pool": + with _checkpointer_postgres_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") + + +@pytest.fixture( + scope="function", + params=[ + "memory", + "sqlite_aio", + ] + if NO_DOCKER + else [ + "memory", + "sqlite_aio", + "postgres_aio", + "postgres_aio_pipe", + "postgres_aio_pool", + ], +) +async def async_checkpointer( + request: pytest.FixtureRequest, +) -> AsyncIterator[BaseCheckpointSaver]: + checkpointer_name = request.param + if checkpointer_name == "memory": + with _checkpointer_memory() as checkpointer: + yield checkpointer + elif checkpointer_name == "sqlite_aio": + async with _checkpointer_sqlite_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio": + async with _checkpointer_postgres_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pipe": + async with _checkpointer_postgres_aio_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pool": + async with _checkpointer_postgres_aio_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") diff --git a/langgraph/source/libs/langgraph/tests/conftest_checkpointer.py b/langgraph/source/libs/langgraph/tests/conftest_checkpointer.py new file mode 100644 index 0000000000000000000000000000000000000000..99c0686fcb2deb16528f6a00865074b1192f72c1 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/conftest_checkpointer.py @@ -0,0 +1,200 @@ +from contextlib import asynccontextmanager, contextmanager +from uuid import uuid4 + +import pytest +from langgraph.checkpoint.postgres import PostgresSaver +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from langgraph.checkpoint.serde.encrypted import EncryptedSerializer +from langgraph.checkpoint.sqlite import SqliteSaver +from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver +from psycopg import AsyncConnection, Connection +from psycopg_pool import AsyncConnectionPool, ConnectionPool + +pytest.register_assert_rewrite("tests.memory_assert") + +from tests.memory_assert import ( # noqa: E402 + MemorySaverAssertImmutable, + MemorySaverNeedsPendingSendsMigration, +) + +DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" + + +@contextmanager +def _checkpointer_memory(): + yield MemorySaverAssertImmutable() + + +@contextmanager +def _checkpointer_memory_migrate_sends(): + yield MemorySaverNeedsPendingSendsMigration() + + +@contextmanager +def _checkpointer_sqlite(): + with SqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer + + +@contextmanager +def _checkpointer_sqlite_aes(): + with SqliteSaver.from_conn_string(":memory:") as checkpointer: + checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( + key=b"1234567890123456" + ) + yield checkpointer + + +@contextmanager +def _checkpointer_postgres(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _checkpointer_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with PostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + checkpointer.setup() + # setup can't run inside pipeline because of implicit transaction + with checkpointer.conn.pipeline() as pipe: + checkpointer.pipe = pipe + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _checkpointer_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + with ConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = PostgresSaver(pool) + checkpointer.setup() + yield checkpointer + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _checkpointer_sqlite_aio(): + async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_postgres_aio(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncPostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + await checkpointer.setup() + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncPostgresSaver.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as checkpointer: + await checkpointer.setup() + # setup can't run inside pipeline because of implicit transaction + async with checkpointer.conn.pipeline() as pipe: + checkpointer.pipe = pipe + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _checkpointer_postgres_aio_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + # yield checkpointer + async with AsyncConnectionPool( + DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True} + ) as pool: + checkpointer = AsyncPostgresSaver(pool) + await checkpointer.setup() + yield checkpointer + finally: + # drop unique db + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +__all__ = [ + "_checkpointer_memory", + "_checkpointer_memory_migrate_sends", + "_checkpointer_sqlite", + "_checkpointer_sqlite_aes", + "_checkpointer_postgres", + "_checkpointer_postgres_pipe", + "_checkpointer_postgres_pool", + "_checkpointer_sqlite_aio", + "_checkpointer_postgres_aio", + "_checkpointer_postgres_aio_pipe", + "_checkpointer_postgres_aio_pool", +] diff --git a/langgraph/source/libs/langgraph/tests/conftest_store.py b/langgraph/source/libs/langgraph/tests/conftest_store.py new file mode 100644 index 0000000000000000000000000000000000000000..4ae8c2445e161fd29349dbbf59d16cbafc63d88a --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/conftest_store.py @@ -0,0 +1,145 @@ +from contextlib import asynccontextmanager, contextmanager +from uuid import uuid4 + +from langgraph.store.memory import InMemoryStore +from langgraph.store.postgres import AsyncPostgresStore, PostgresStore +from psycopg import AsyncConnection, Connection + +DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" + + +@contextmanager +def _store_memory(): + store = InMemoryStore() + yield store + + +@contextmanager +def _store_postgres(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _store_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() # Run in its own transaction + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _store_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10} + ) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio(): + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pipe(): + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() # Run in its own transaction + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pool(): + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, + pool_config={"max_size": 10}, + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +__all__ = [ + "_store_memory", + "_store_postgres", + "_store_postgres_pipe", + "_store_postgres_pool", + "_store_postgres_aio", + "_store_postgres_aio_pipe", + "_store_postgres_aio_pool", +] diff --git a/langgraph/source/libs/langgraph/tests/example_app/example_graph.py b/langgraph/source/libs/langgraph/tests/example_app/example_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..2a3f2a52c57e64896b1953f66cf28b8736d743a3 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/example_app/example_graph.py @@ -0,0 +1,89 @@ +from typing import Annotated + +from langchain_core.messages import AIMessage, BaseMessage, ToolMessage +from langchain_core.tools import tool +from typing_extensions import TypedDict + +from langgraph.func import entrypoint, task +from langgraph.graph.message import add_messages +from tests.fake_chat import FakeChatModel + + +class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + +@tool +def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + +tools = [search_api] +tools_by_name = {t.name: t for t in tools} + + +def get_model(): + model = FakeChatModel( + messages=[ + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + return model + + +@task +def foo(): + return "foo" + + +@entrypoint() +async def app(state: AgentState) -> AgentState: + model = get_model() + max_steps = 100 + messages = state["messages"][:] + await foo() # Very useful call here ya know. + for _ in range(max_steps): + message = await model.ainvoke(messages) + messages.append(message) + if not message.tool_calls: + break + # Assume it's the search tool + tool_results = await search_api.abatch( + [t["args"]["query"] for t in message.tool_calls] + ) + messages.extend( + [ + ToolMessage(content=tool_res, tool_call_id=tc["id"]) + for tc, tool_res in zip(message.tool_calls, tool_results) + ] + ) + + return entrypoint.final(value=messages[-1], save={"messages": messages}) diff --git a/langgraph/source/libs/langgraph/tests/example_app/langgraph.json b/langgraph/source/libs/langgraph/tests/example_app/langgraph.json new file mode 100644 index 0000000000000000000000000000000000000000..5d0a88e41fa10e8708ec565a74d20de4ef4bc0dd --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/example_app/langgraph.json @@ -0,0 +1,7 @@ +{ + "$schema": "https://langgra.ph/schema.json", + "graphs": { + "app": "tests/example_app/example_graph.py:app" + }, + "dependencies": ["tests/example_app"] +} diff --git a/langgraph/source/libs/langgraph/tests/example_app/requirements.txt b/langgraph/source/libs/langgraph/tests/example_app/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..b3557ec4e5e330784260d9c33fa9e0e48bfbeb04 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/example_app/requirements.txt @@ -0,0 +1,2 @@ +langchain-core +-e . \ No newline at end of file diff --git a/langgraph/source/libs/langgraph/tests/fake_chat.py b/langgraph/source/libs/langgraph/tests/fake_chat.py new file mode 100644 index 0000000000000000000000000000000000000000..827089d4a96613e9d43e42d0ced2dc1fd25f34db --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/fake_chat.py @@ -0,0 +1,157 @@ +import re +from collections.abc import AsyncIterator, Iterator +from typing import Any, cast + +from langchain_core.callbacks import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult + + +class FakeChatModel(GenericFakeChatModel): + messages: list[BaseMessage] + + i: int = 0 + + def bind_tools(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + """Top Level call""" + if self.i >= len(self.messages): + self.i = 0 + message = self.messages[self.i] + self.i += 1 + if isinstance(message, str): + message_ = AIMessage(content=message) + else: + if hasattr(message, "model_copy"): + message_ = message.model_copy() + else: + message_ = message.copy() + generation = ChatGeneration(message=message_) + return ChatResult(generations=[generation]) + + def _stream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> Iterator[ChatGenerationChunk]: + """Stream the output of the model.""" + chat_result = self._generate( + messages, stop=stop, run_manager=run_manager, **kwargs + ) + if not isinstance(chat_result, ChatResult): + raise ValueError( + f"Expected generate to return a ChatResult, " + f"but got {type(chat_result)} instead." + ) + + message = chat_result.generations[0].message + + if not isinstance(message, AIMessage): + raise ValueError( + f"Expected invoke to return an AIMessage, " + f"but got {type(message)} instead." + ) + + content = message.content + + if content: + # Use a regular expression to split on whitespace with a capture group + # so that we can preserve the whitespace in the output. + assert isinstance(content, str) + content_chunks = cast(list[str], re.split(r"(\s)", content)) + + for i, token in enumerate(content_chunks): + if i == len(content_chunks) - 1: + chunk = ChatGenerationChunk( + message=AIMessageChunk( + content=token, id=message.id, chunk_position="last" + ) + ) + else: + chunk = ChatGenerationChunk( + message=AIMessageChunk(content=token, id=message.id) + ) + if run_manager: + run_manager.on_llm_new_token(token, chunk=chunk) + yield chunk + else: + args = message.__dict__ + args.pop("type") + chunk = ChatGenerationChunk( + message=AIMessageChunk(**args, chunk_position="last") + ) + if run_manager: + run_manager.on_llm_new_token("", chunk=chunk) + yield chunk + + async def _astream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> AsyncIterator[ChatGenerationChunk]: + """Stream the output of the model.""" + chat_result = self._generate( + messages, stop=stop, run_manager=run_manager, **kwargs + ) + if not isinstance(chat_result, ChatResult): + raise ValueError( + f"Expected generate to return a ChatResult, " + f"but got {type(chat_result)} instead." + ) + + message = chat_result.generations[0].message + + if not isinstance(message, AIMessage): + raise ValueError( + f"Expected invoke to return an AIMessage, " + f"but got {type(message)} instead." + ) + + content = message.content + + if content: + # Use a regular expression to split on whitespace with a capture group + # so that we can preserve the whitespace in the output. + assert isinstance(content, str) + content_chunks = cast(list[str], re.split(r"(\s)", content)) + + for i, token in enumerate(content_chunks): + if i == len(content_chunks) - 1: + chunk = ChatGenerationChunk( + message=AIMessageChunk( + content=token, id=message.id, chunk_position="last" + ) + ) + else: + chunk = ChatGenerationChunk( + message=AIMessageChunk(content=token, id=message.id) + ) + + if run_manager: + run_manager.on_llm_new_token(token, chunk=chunk) + yield chunk + else: + args = message.__dict__ + args.pop("type") + chunk = ChatGenerationChunk( + message=AIMessageChunk(**args, chunk_position="last") + ) + if run_manager: + await run_manager.on_llm_new_token("", chunk=chunk) + yield chunk diff --git a/langgraph/source/libs/langgraph/tests/fake_tracer.py b/langgraph/source/libs/langgraph/tests/fake_tracer.py new file mode 100644 index 0000000000000000000000000000000000000000..37133980f1f4646459f9111fe82ee0a3fa0fa079 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/fake_tracer.py @@ -0,0 +1,91 @@ +from typing import Any +from uuid import UUID + +from langchain_core.messages.base import BaseMessage +from langchain_core.outputs.chat_generation import ChatGeneration +from langchain_core.outputs.llm_result import LLMResult +from langchain_core.tracers import BaseTracer, Run + + +class FakeTracer(BaseTracer): + """Fake tracer that records LangChain execution. + It replaces run ids with deterministic UUIDs for snapshotting.""" + + def __init__(self) -> None: + """Initialize the tracer.""" + super().__init__() + self.runs: list[Run] = [] + self.uuids_map: dict[UUID, UUID] = {} + self.uuids_generator = ( + UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000) + ) + + def _replace_uuid(self, uuid: UUID) -> UUID: + if uuid not in self.uuids_map: + self.uuids_map[uuid] = next(self.uuids_generator) + return self.uuids_map[uuid] + + def _replace_message_id(self, maybe_message: Any) -> Any: + if isinstance(maybe_message, BaseMessage): + maybe_message.id = str(next(self.uuids_generator)) + if isinstance(maybe_message, ChatGeneration): + maybe_message.message.id = str(next(self.uuids_generator)) + if isinstance(maybe_message, LLMResult): + for i, gen_list in enumerate(maybe_message.generations): + for j, gen in enumerate(gen_list): + maybe_message.generations[i][j] = self._replace_message_id(gen) + if isinstance(maybe_message, dict): + for k, v in maybe_message.items(): + maybe_message[k] = self._replace_message_id(v) + if isinstance(maybe_message, list): + for i, v in enumerate(maybe_message): + maybe_message[i] = self._replace_message_id(v) + + return maybe_message + + def _copy_run(self, run: Run) -> Run: + if run.dotted_order: + levels = run.dotted_order.split(".") + processed_levels = [] + for level in levels: + timestamp, run_id = level.split("Z") + new_run_id = self._replace_uuid(UUID(run_id)) + processed_level = f"{timestamp}Z{new_run_id}" + processed_levels.append(processed_level) + new_dotted_order = ".".join(processed_levels) + else: + new_dotted_order = None + return run.copy( + update={ + "id": self._replace_uuid(run.id), + "parent_run_id": ( + self.uuids_map[run.parent_run_id] if run.parent_run_id else None + ), + "child_runs": [self._copy_run(child) for child in run.child_runs], + "trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None, + "dotted_order": new_dotted_order, + "inputs": self._replace_message_id(run.inputs), + "outputs": self._replace_message_id(run.outputs), + } + ) + + def _persist_run(self, run: Run) -> None: + """Persist a run.""" + + self.runs.append(self._copy_run(run)) + + def flattened_runs(self) -> list[Run]: + q = [] + self.runs + result = [] + while q: + parent = q.pop() + result.append(parent) + if parent.child_runs: + q.extend(parent.child_runs) + return result + + @property + def run_ids(self) -> list[UUID | None]: + runs = self.flattened_runs() + uuids_map = {v: k for k, v in self.uuids_map.items()} + return [uuids_map.get(r.id) for r in runs] diff --git a/langgraph/source/libs/langgraph/tests/memory_assert.py b/langgraph/source/libs/langgraph/tests/memory_assert.py new file mode 100644 index 0000000000000000000000000000000000000000..d9ca4904c8d549a4f79f762f2edff679839936ac --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/memory_assert.py @@ -0,0 +1,101 @@ +import os +import tempfile +from collections import defaultdict +from functools import partial +from typing import Any + +from langchain_core.runnables import RunnableConfig +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + ChannelVersions, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, + SerializerProtocol, +) +from langgraph.checkpoint.memory import InMemorySaver, PersistentDict + +from langgraph.constants import TASKS + + +class NoopSerializer(SerializerProtocol): + def loads_typed(self, data: tuple[str, bytes]) -> Any: + return data[1] + + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + return "type", obj + + +class MemorySaverNeedsPendingSendsMigration(BaseCheckpointSaver): + def __init__(self) -> None: + self.saver = InMemorySaver() + + def __getattribute__(self, name): + if name in ("saver", "__class__", "get_tuple"): + return object.__getattribute__(self, name) + return getattr(self.saver, name) + + def get_tuple(self, config): + if tup := self.saver.get_tuple(config): + if tup.checkpoint["v"] == 4 and tup.checkpoint["channel_values"].get(TASKS): + tup.checkpoint["v"] = 3 + tup.checkpoint["pending_sends"] = tup.checkpoint["channel_values"].pop( + TASKS + ) + tup.checkpoint["channel_versions"].pop(TASKS) + for seen in tup.checkpoint["versions_seen"].values(): + seen.pop(TASKS, None) + return tup + + +class MemorySaverAssertImmutable(InMemorySaver): + storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]] + + def __init__( + self, + *, + serde: SerializerProtocol | None = None, + put_sleep: float | None = None, + ) -> None: + _, filename = tempfile.mkstemp() + super().__init__( + serde=serde, factory=partial(PersistentDict, filename=filename) + ) + self.storage_for_copies = defaultdict(lambda: defaultdict(dict)) + self.put_sleep = put_sleep + self.stack.callback(os.remove, filename) + + def put( + self, + config: dict, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> None: + if self.put_sleep: + import time + + time.sleep(self.put_sleep) + # assert checkpoint hasn't been modified since last written + thread_id = config["configurable"]["thread_id"] + checkpoint_ns = config["configurable"]["checkpoint_ns"] + if saved := super().get(config): + assert ( + self.serde.loads_typed( + self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]] + ) + == saved + ), config["configurable"]["checkpoint_ns"] + self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = ( + self.serde.dumps_typed(checkpoint) + ) + # call super to write checkpoint + return super().put(config, checkpoint, metadata, new_versions) + + +class MemorySaverNoPending(InMemorySaver): + def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + result = super().get_tuple(config) + if result: + return CheckpointTuple(result.config, result.checkpoint, result.metadata) + return result diff --git a/langgraph/source/libs/langgraph/tests/messages.py b/langgraph/source/libs/langgraph/tests/messages.py new file mode 100644 index 0000000000000000000000000000000000000000..ecc657a3622b048fc63b2cee0c3fd5612efc84bf --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/messages.py @@ -0,0 +1,50 @@ +"""Redefined messages as a work-around for pydantic issue with AnyStr. + +The code below creates version of pydantic models +that will work in unit tests with AnyStr as id field +Please note that the `id` field is assigned AFTER the model is created +to workaround an issue with pydantic ignoring the __eq__ method on +subclassed strings. +""" + +from typing import Any + +from langchain_core.documents import Document +from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage + +from tests.any_str import AnyStr + + +def _AnyIdDocument(**kwargs: Any) -> Document: + """Create a document with an id field.""" + message = Document(**kwargs) + message.id = AnyStr() + return message + + +def _AnyIdAIMessage(**kwargs: Any) -> AIMessage: + """Create ai message with an any id field.""" + message = AIMessage(**kwargs) + message.id = AnyStr() + return message + + +def _AnyIdAIMessageChunk(**kwargs: Any) -> AIMessageChunk: + """Create ai message with an any id field.""" + message = AIMessageChunk(**kwargs) + message.id = AnyStr() + return message + + +def _AnyIdHumanMessage(**kwargs: Any) -> HumanMessage: + """Create a human message with an any id field.""" + message = HumanMessage(**kwargs) + message.id = AnyStr() + return message + + +def _AnyIdToolMessage(**kwargs: Any) -> ToolMessage: + """Create a tool message with an any id field.""" + message = ToolMessage(**kwargs) + message.id = AnyStr() + return message diff --git a/langgraph/source/libs/langgraph/tests/test_algo.py b/langgraph/source/libs/langgraph/tests/test_algo.py new file mode 100644 index 0000000000000000000000000000000000000000..0bf988173212f86b4961f7a7df97a7ee3924650e --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_algo.py @@ -0,0 +1,67 @@ +from langgraph._internal._constants import PULL, PUSH +from langgraph.pregel._algo import prepare_next_tasks, task_path_str +from langgraph.pregel._checkpoint import channels_from_checkpoint, empty_checkpoint + + +def test_prepare_next_tasks() -> None: + config = {} + processes = {} + checkpoint = empty_checkpoint() + channels, managed = channels_from_checkpoint({}, checkpoint) + + assert ( + prepare_next_tasks( + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + -1, + for_execution=False, + ) + == {} + ) + assert ( + prepare_next_tasks( + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + -1, + for_execution=True, + checkpointer=None, + store=None, + manager=None, + ) + == {} + ) + + # TODO: add more tests + + +def test_tuple_str() -> None: + push_path_a = (PUSH, 2) + pull_path_a = (PULL, "abc") + push_path_b = (PUSH, push_path_a, 1) + push_path_c = (PUSH, push_path_b, 3) + + assert task_path_str(push_path_a) == f"~{PUSH}, 0000000002" + assert task_path_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001" + assert ( + task_path_str(push_path_c) + == f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003" + ) + assert task_path_str(pull_path_a) == f"~{PULL}, abc" + + path_list = [push_path_b, push_path_a, pull_path_a, push_path_c] + assert sorted(map(task_path_str, path_list)) == [ + f"~{PULL}, abc", + f"~{PUSH}, 0000000002", + f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001", + f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003", + ] diff --git a/langgraph/source/libs/langgraph/tests/test_channels.py b/langgraph/source/libs/langgraph/tests/test_channels.py new file mode 100644 index 0000000000000000000000000000000000000000..37a5eb432bdc8d7ae289c6f9b742545be96613f3 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_channels.py @@ -0,0 +1,119 @@ +import operator +from collections.abc import Sequence + +import pytest + +from langgraph._internal._typing import MISSING +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.last_value import LastValue +from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.errors import EmptyChannelError, InvalidUpdateError + +pytestmark = pytest.mark.anyio + + +def test_last_value() -> None: + channel = LastValue(int).from_checkpoint(MISSING) + assert channel.ValueType is int + assert channel.UpdateType is int + + with pytest.raises(EmptyChannelError): + channel.get() + with pytest.raises(InvalidUpdateError): + channel.update([5, 6]) + + channel.update([3]) + assert channel.get() == 3 + channel.update([4]) + assert channel.get() == 4 + checkpoint = channel.checkpoint() + channel = LastValue(int).from_checkpoint(checkpoint) + assert channel.get() == 4 + + +def test_topic() -> None: + channel = Topic(str).from_checkpoint(MISSING) + assert channel.ValueType == Sequence[str] + assert channel.UpdateType == str | list[str] + + assert channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + assert channel.update([["c", "d"], "d"]) + assert channel.get() == ["c", "d", "d"] + assert channel.update([]) + with pytest.raises(EmptyChannelError): + channel.get() + assert not channel.update([]), "channel already empty" + assert channel.update(["e"]) + assert channel.get() == ["e"] + checkpoint = channel.checkpoint() + channel = Topic(str).from_checkpoint(checkpoint) + assert channel.get() == ["e"] + channel_copy = Topic(str).from_checkpoint(checkpoint) + channel_copy.update(["f"]) + assert channel_copy.get() == ["f"] + assert channel.get() == ["e"] + + +def test_topic_accumulate() -> None: + channel = Topic(str, accumulate=True).from_checkpoint(MISSING) + assert channel.ValueType == Sequence[str] + assert channel.UpdateType == str | list[str] + + assert channel.update(["a", "b"]) + assert channel.get() == ["a", "b"] + assert channel.update(["b", ["c", "d"], "d"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + assert not channel.update([]) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + checkpoint = channel.checkpoint() + channel = Topic(str, accumulate=True).from_checkpoint(checkpoint) + assert channel.get() == ["a", "b", "b", "c", "d", "d"] + assert channel.update(["e"]) + assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"] + + +def test_binop() -> None: + channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(MISSING) + assert channel.ValueType is int + assert channel.UpdateType is int + + assert channel.get() == 0 + + channel.update([1, 2, 3]) + assert channel.get() == 6 + channel.update([4]) + assert channel.get() == 10 + checkpoint = channel.checkpoint() + channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(checkpoint) + assert channel.get() == 10 + + +def test_untracked_value() -> None: + channel = UntrackedValue(dict).from_checkpoint(MISSING) + assert channel.ValueType is dict + assert channel.UpdateType is dict + + # UntrackedValue should start empty + with pytest.raises(EmptyChannelError): + channel.get() + + # Should be able to update with a value + test_data = {"session": "test", "temp": "dir"} + channel.update([test_data]) + assert channel.get() == test_data + + # Update with new value + new_data = {"session": "updated", "temp": "newdir"} + channel.update([new_data]) + assert channel.get() == new_data + + # On checkpoint, UntrackedValue should return MISSING + checkpoint = channel.checkpoint() + assert checkpoint is MISSING + + # Creating from checkpoint with MISSING should start empty + new_channel = UntrackedValue(dict).from_checkpoint(checkpoint) + with pytest.raises(EmptyChannelError): + new_channel.get() diff --git a/langgraph/source/libs/langgraph/tests/test_checkpoint_migration.py b/langgraph/source/libs/langgraph/tests/test_checkpoint_migration.py new file mode 100644 index 0000000000000000000000000000000000000000..c53c5c72f4ebda6974ec37f690d1aa365c5cc804 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_checkpoint_migration.py @@ -0,0 +1,1727 @@ +import operator +import sys +import time +from collections import defaultdict +from typing import Annotated, Literal + +import pytest +from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple +from typing_extensions import TypedDict + +from langgraph._internal._config import patch_configurable +from langgraph.graph.state import StateGraph +from langgraph.pregel._checkpoint import copy_checkpoint +from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, interrupt +from tests.any_int import AnyInt +from tests.any_str import AnyDict, AnyObject, AnyStr + +pytestmark = pytest.mark.anyio + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +def get_expected_history(*, exc_task_results: int = 0) -> list[StateSnapshot]: + return [ + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + next=("qa",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="qa", + path=("__pregel_pull", "qa"), + error=None, + interrupts=() + if exc_task_results + else ( + Interrupt( + value="", + id=AnyStr(), + ), + ), + state=None, + result=None + if exc_task_results + else {"answer": "doc1,doc2,doc3,doc4"}, + ), + ), + interrupts=() + if exc_task_results + else ( + Interrupt( + value="", + id=AnyStr(), + ), + ), + ), + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc3", "doc4"], + }, + next=("retriever_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="retriever_one", + path=("__pregel_pull", "retriever_one"), + error=None, + interrupts=(), + state=None, + result=None if exc_task_results else {"docs": ["doc1", "doc2"]}, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={"query": "query: what is weather in sf", "docs": []}, + next=("analyzer_one", "retriever_two"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="analyzer_one", + path=("__pregel_pull", "analyzer_one"), + error=None, + interrupts=(), + state=None, + result=None + if exc_task_results + else {"query": "analyzed: query: what is weather in sf"}, + ), + PregelTask( + id=AnyStr(), + name="retriever_two", + path=("__pregel_pull", "retriever_two"), + error=None, + interrupts=(), + state=None, + result=None + if exc_task_results >= 2 + else {"docs": ["doc3", "doc4"]}, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={"query": "what is weather in sf", "docs": []}, + next=("rewrite_query",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="rewrite_query", + path=("__pregel_pull", "rewrite_query"), + error=None, + interrupts=(), + state=None, + result=None + if exc_task_results + else {"query": "query: what is weather in sf"}, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={"docs": []}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"query": "what is weather in sf"}, + ), + ), + interrupts=(), + ), + ] + + +SAVED_CHECKPOINTS = { + "3": [ + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2149-6faa-8004-9d848038f10a", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.237381+00:00", + "id": "1f00fd5f-2149-6faa-8004-9d848038f10a", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "branch:to:qa": "00000000000000000000000000000006.0.6375257096095945", + "answer": "00000000000000000000000000000006.0.9100669543952636", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + "retriever_one": { + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264" + }, + "__interrupt__": { + "query": "00000000000000000000000000000004.0.18727156933289513", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962", + }, + "qa": { + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, + "updated_channels": None, + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + } + }, + pending_writes=[], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.233695+00:00", + "id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + "retriever_one": { + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "branch:to:qa": None, + }, + "updated_channels": None, + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-213c-6940-8002-a28f475a6478", + } + }, + pending_writes=[ + ( + "2430f303-da9f-2e3e-738c-2e8ea28e8973", + "__interrupt__", + [ + Interrupt( + value="", + resumable=True, # type: ignore[arg-type] + ns=["qa:2430f303-da9f-2e3e-738c-2e8ea28e8973"], # type: ignore[arg-type] + ) + ], + ), + ("00000000-0000-0000-0000-000000000000", "__resume__", ""), + ("2430f303-da9f-2e3e-738c-2e8ea28e8973", "__resume__", [""]), + ( + "2430f303-da9f-2e3e-738c-2e8ea28e8973", + "answer", + "doc1,doc2,doc3,doc4", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-213c-6940-8002-a28f475a6478", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.231890+00:00", + "id": "1f00fd5f-213c-6940-8002-a28f475a6478", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264", + "docs": "00000000000000000000000000000004.0.972701399851098", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "branch:to:retriever_one": None, + "docs": ["doc3", "doc4"], + }, + "updated_channels": None, + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + } + }, + pending_writes=[ + ("a5602426-85f2-1fe4-c9e4-bd0127e8e53e", "docs", ["doc1", "doc2"]), + ("a5602426-85f2-1fe4-c9e4-bd0127e8e53e", "branch:to:qa", None), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.125661+00:00", + "id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000003.0.04057405566428263", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241", + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + }, + "channel_values": { + "query": "query: what is weather in sf", + "branch:to:analyzer_one": None, + "branch:to:retriever_two": None, + }, + "updated_channels": None, + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + } + }, + pending_writes=[ + ( + "4e7cb70b-7e0f-52d0-d8aa-5439bd3f84de", + "query", + "analyzed: query: what is weather in sf", + ), + ( + "4e7cb70b-7e0f-52d0-d8aa-5439bd3f84de", + "branch:to:retriever_one", + None, + ), + ("abcbc448-cfba-ac2b-2e39-346808f20add", "docs", ["doc3", "doc4"]), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.125200+00:00", + "id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000002.0.3399249312096154", + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + }, + "channel_values": { + "query": "what is weather in sf", + "branch:to:rewrite_query": None, + }, + "updated_channels": None, + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + } + }, + pending_writes=[ + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "query", + "query: what is weather in sf", + ), + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "branch:to:analyzer_one", + None, + ), + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "branch:to:retriever_two", + None, + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.124678+00:00", + "id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + "channel_versions": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "versions_seen": {"__input__": {}}, + "channel_values": {"__start__": {"query": "what is weather in sf"}}, + "updated_channels": None, + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + parent_config=None, + pending_writes=[ + ( + "a9e2a749-9870-1952-0a6c-b23b6729ffda", + "query", + "what is weather in sf", + ), + ( + "a9e2a749-9870-1952-0a6c-b23b6729ffda", + "branch:to:rewrite_query", + None, + ), + ], + ), + ], + "2-start:*": [ + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515f-6b88-8004-6fffb69dd465", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.825576+00:00", + "id": "1f00fe48-515f-6b88-8004-6fffb69dd465", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "retriever_one": "00000000000000000000000000000006.0.9577453764095437", + "answer": "00000000000000000000000000000006.0.27361287406148327", + "qa": "00000000000000000000000000000006.0.24260043089701677", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268" + }, + "__interrupt__": { + "query": "00000000000000000000000000000004.0.05732679770452498", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "retriever_one": "00000000000000000000000000000005.0.0523757506060204", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + }, + "qa": { + "retriever_one": "00000000000000000000000000000005.0.0523757506060204" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + "qa": "qa", + }, + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + } + }, + pending_writes=[], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.824251+00:00", + "id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "retriever_one": "00000000000000000000000000000005.0.0523757506060204", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "retriever_one": "retriever_one", + }, + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + } + }, + pending_writes=[ + ( + "4ee8637e-0a95-285e-75bc-4da721c0beab", + "__interrupt__", + [ + Interrupt( + value="", + resumable=True, # type: ignore[arg-type] + ns=["qa:4ee8637e-0a95-285e-75bc-4da721c0beab"], # type: ignore[arg-type] + ) + ], + ), + ("00000000-0000-0000-0000-000000000000", "__resume__", ""), + ("4ee8637e-0a95-285e-75bc-4da721c0beab", "__resume__", [""]), + ( + "4ee8637e-0a95-285e-75bc-4da721c0beab", + "answer", + "doc1,doc2,doc3,doc4", + ), + ("4ee8637e-0a95-285e-75bc-4da721c0beab", "qa", "qa"), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.823978+00:00", + "id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268", + "docs": "00000000000000000000000000000004.0.49012772235571145", + "retriever_two": "00000000000000000000000000000004.0.9223450775254257", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "analyzer_one": "analyzer_one", + "docs": ["doc3", "doc4"], + "retriever_two": "retriever_two", + }, + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + } + }, + pending_writes=[ + ("16295c56-f44e-31fa-8fad-fff3f9022629", "docs", ["doc1", "doc2"]), + ( + "16295c56-f44e-31fa-8fad-fff3f9022629", + "retriever_one", + "retriever_one", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.718258+00:00", + "id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000003.0.10748450241039154", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955", + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + }, + "channel_values": { + "query": "query: what is weather in sf", + "rewrite_query": "rewrite_query", + "branch:to:retriever_two": "rewrite_query", + }, + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5058-6a46-8000-086bffc73797", + } + }, + pending_writes=[ + ( + "baecc0e3-ea00-0e00-9436-e33cd2527faf", + "query", + "analyzed: query: what is weather in sf", + ), + ( + "baecc0e3-ea00-0e00-9436-e33cd2527faf", + "analyzer_one", + "analyzer_one", + ), + ("96b7bfe4-269f-092c-e685-14dba6a27271", "docs", ["doc3", "doc4"]), + ( + "96b7bfe4-269f-092c-e685-14dba6a27271", + "retriever_two", + "retriever_two", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5058-6a46-8000-086bffc73797", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.717827+00:00", + "id": "1f00fe48-5058-6a46-8000-086bffc73797", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000002.0.706632616485588", + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + }, + "channel_values": { + "query": "what is weather in sf", + "start:rewrite_query": "__start__", + }, + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + } + }, + pending_writes=[ + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "query", + "query: what is weather in sf", + ), + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "rewrite_query", + "rewrite_query", + ), + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "branch:to:retriever_two", + "rewrite_query", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.717221+00:00", + "id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + "channel_versions": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "versions_seen": {"__input__": {}}, + "channel_values": {"__start__": {"query": "what is weather in sf"}}, + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + parent_config=None, + pending_writes=[ + ( + "891e8564-d78f-7fb2-f15d-bce2a0ddf1c6", + "query", + "what is weather in sf", + ), + ( + "891e8564-d78f-7fb2-f15d-bce2a0ddf1c6", + "start:rewrite_query", + "__start__", + ), + ], + ), + ], + "2-quadratic": [ + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-546a-64f0-8004-d2d4e06b6fb6", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.630535+00:00", + "id": "1f00ffbe-546a-64f0-8004-d2d4e06b6fb6", + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "qa": "qa", + }, + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.141080617837112", + "query": "00000000000000000000000000000004.0.9004905790383284", + "start:rewrite_query": "00000000000000000000000000000003.0.013109117892399547", + "rewrite_query": "00000000000000000000000000000004.0.1679336326485974", + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000004.0.7474512867042074", + "analyzer_one": "00000000000000000000000000000005.0.5817293698381076", + "docs": "00000000000000000000000000000005.0.9650795030435029", + "retriever_two": "00000000000000000000000000000005.0.77101858493518", + "retriever_one": "00000000000000000000000000000006.0.4984661612084784", + "answer": "00000000000000000000000000000006.0.6244466008661432", + "qa": "00000000000000000000000000000006.0.06630110662217248", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.32578323811902354" + }, + "retriever_two": { + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.2684613370070208" + }, + "__interrupt__": { + "query": "00000000000000000000000000000004.0.9004905790383284", + "docs": "00000000000000000000000000000005.0.9650795030435029", + "__start__": "00000000000000000000000000000002.0.141080617837112", + "rewrite_query": "00000000000000000000000000000004.0.1679336326485974", + "analyzer_one": "00000000000000000000000000000005.0.5817293698381076", + "retriever_one": "00000000000000000000000000000005.0.222301724202566", + "retriever_two": "00000000000000000000000000000005.0.77101858493518", + "start:rewrite_query": "00000000000000000000000000000003.0.013109117892399547", + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000004.0.7474512867042074", + }, + "qa": { + "retriever_one": "00000000000000000000000000000005.0.222301724202566" + }, + }, + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5466-61ac-8003-7ec684cd12cc", + } + }, + pending_writes=[], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5466-61ac-8003-7ec684cd12cc", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.628817+00:00", + "id": "1f00ffbe-5466-61ac-8003-7ec684cd12cc", + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "retriever_one": "retriever_one", + }, + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.141080617837112", + "query": "00000000000000000000000000000004.0.9004905790383284", + "start:rewrite_query": "00000000000000000000000000000003.0.013109117892399547", + "rewrite_query": "00000000000000000000000000000004.0.1679336326485974", + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000004.0.7474512867042074", + "analyzer_one": "00000000000000000000000000000005.0.5817293698381076", + "docs": "00000000000000000000000000000005.0.9650795030435029", + "retriever_two": "00000000000000000000000000000005.0.77101858493518", + "retriever_one": "00000000000000000000000000000005.0.222301724202566", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.32578323811902354" + }, + "retriever_two": { + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.2684613370070208" + }, + }, + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5465-6248-8002-ee1d8bdbbee5", + } + }, + pending_writes=[ + ( + "369e94b1-77d1-d67a-ab59-23d1ba20ee73", + "__interrupt__", + [ + Interrupt( + value="", + resumable=True, + ns=["qa:369e94b1-77d1-d67a-ab59-23d1ba20ee73"], # type: ignore[arg-type] + ) + ], + ), + ("00000000-0000-0000-0000-000000000000", "__resume__", ""), + ("369e94b1-77d1-d67a-ab59-23d1ba20ee73", "__resume__", [""]), + ( + "369e94b1-77d1-d67a-ab59-23d1ba20ee73", + "answer", + "doc1,doc2,doc3,doc4", + ), + ("369e94b1-77d1-d67a-ab59-23d1ba20ee73", "qa", "qa"), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5465-6248-8002-ee1d8bdbbee5", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.628408+00:00", + "id": "1f00ffbe-5465-6248-8002-ee1d8bdbbee5", + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc3", "doc4"], + "analyzer_one": "analyzer_one", + "retriever_two": "retriever_two", + }, + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.141080617837112", + "query": "00000000000000000000000000000004.0.9004905790383284", + "start:rewrite_query": "00000000000000000000000000000003.0.013109117892399547", + "rewrite_query": "00000000000000000000000000000004.0.1679336326485974", + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000004.0.7474512867042074", + "analyzer_one": "00000000000000000000000000000004.0.2684613370070208", + "docs": "00000000000000000000000000000004.0.37458911821520957", + "retriever_two": "00000000000000000000000000000004.0.7340649978617967", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.32578323811902354" + }, + "retriever_two": { + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405" + }, + }, + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-536a-6dca-8001-c7e021a73244", + } + }, + pending_writes=[ + ("601f2099-cb23-41f9-ae64-9b4e4a6b675e", "docs", ["doc1", "doc2"]), + ( + "601f2099-cb23-41f9-ae64-9b4e4a6b675e", + "retriever_one", + "retriever_one", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-536a-6dca-8001-c7e021a73244", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.525915+00:00", + "id": "1f00ffbe-536a-6dca-8001-c7e021a73244", + "channel_values": { + "query": "query: what is weather in sf", + "rewrite_query": "rewrite_query", + "branch:rewrite_query:rewrite_query_then:retriever_two": "rewrite_query", + }, + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.141080617837112", + "query": "00000000000000000000000000000003.0.8982471206042032", + "start:rewrite_query": "00000000000000000000000000000003.0.013109117892399547", + "rewrite_query": "00000000000000000000000000000003.0.32578323811902354", + "branch:rewrite_query:rewrite_query_then:retriever_two": "00000000000000000000000000000003.0.8992241767805405", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445" + }, + }, + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5369-6de4-8000-70bd3810f0ca", + } + }, + pending_writes=[ + ( + "88757475-bc8c-934c-7d18-921cb2d94864", + "query", + "analyzed: query: what is weather in sf", + ), + ( + "88757475-bc8c-934c-7d18-921cb2d94864", + "analyzer_one", + "analyzer_one", + ), + ("53c60600-588b-e49e-a9d2-bbcbf30a7497", "docs", ["doc3", "doc4"]), + ( + "53c60600-588b-e49e-a9d2-bbcbf30a7497", + "retriever_two", + "retriever_two", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5369-6de4-8000-70bd3810f0ca", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.525508+00:00", + "id": "1f00ffbe-5369-6de4-8000-70bd3810f0ca", + "channel_values": { + "query": "what is weather in sf", + "start:rewrite_query": "__start__", + }, + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.141080617837112", + "query": "00000000000000000000000000000002.0.06948551802156189", + "start:rewrite_query": "00000000000000000000000000000002.0.32002588286540445", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + }, + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5368-6c0a-bfff-ac22bea4c512", + } + }, + pending_writes=[ + ( + "3bb470dd-9bf8-6216-b5fb-e50162991da1", + "query", + "query: what is weather in sf", + ), + ( + "3bb470dd-9bf8-6216-b5fb-e50162991da1", + "rewrite_query", + "rewrite_query", + ), + ( + "3bb470dd-9bf8-6216-b5fb-e50162991da1", + "branch:rewrite_query:rewrite_query_then:retriever_two", + "rewrite_query", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00ffbe-5368-6c0a-bfff-ac22bea4c512", + } + }, + checkpoint={ + "v": 1, + "ts": "2025-04-02T19:51:40.525051+00:00", + "id": "1f00ffbe-5368-6c0a-bfff-ac22bea4c512", + "channel_values": {"__start__": {"query": "what is weather in sf"}}, + "channel_versions": { + "__start__": "00000000000000000000000000000001.0.6759219622820284" + }, + "versions_seen": {"__input__": {}}, + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + parent_config=None, + pending_writes=[ + ( + "aa8c5e8a-da6f-ccb1-f8a9-3b145cdfe7a4", + "query", + "what is weather in sf", + ), + ( + "aa8c5e8a-da6f-ccb1-f8a9-3b145cdfe7a4", + "start:rewrite_query", + "__start__", + ), + ], + ), + ], +} + + +def make_state_graph() -> StateGraph: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + interrupt("") + return {"answer": ",".join(data["docs"])} + + def rewrite_query_then(data: State) -> Literal["retriever_two"]: + return "retriever_two" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges("rewrite_query", rewrite_query_then) + workflow.add_edge("retriever_one", "qa") + workflow.set_finish_point("qa") + return workflow + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("source,target", [("2-start:*", "3"), ("2-quadratic", "3")]) +def test_migrate_checkpoints(source: str, target: str) -> None: + # Check that the migration function works as expected + builder = make_state_graph() + graph = builder.compile() + + source_checkpoints = list(reversed(SAVED_CHECKPOINTS[source])) + target_checkpoints = list(reversed(SAVED_CHECKPOINTS[target])) + assert len(source_checkpoints) == len(target_checkpoints) + for idx, (source_checkpoint, target_checkpoint) in enumerate( + zip(source_checkpoints, target_checkpoints) + ): + # copy the checkpoint to avoid modifying the original + migrated = copy_checkpoint(source_checkpoint.checkpoint) + # migrate the checkpoint + graph._migrate_checkpoint(migrated) + # replace values that don't need to match exactly + migrated["id"] = AnyStr() + migrated["ts"] = AnyStr() + migrated["v"] = AnyInt() + for k in migrated["channel_values"]: + migrated["channel_values"][k] = AnyObject() + for v in migrated["channel_versions"]: + migrated["channel_versions"][v] = AnyStr( + migrated["channel_versions"][v].split(".")[0] + ) + for c in migrated["versions_seen"]: + for v in migrated["versions_seen"][c]: + migrated["versions_seen"][c][v] = AnyStr( + migrated["versions_seen"][c][v].split(".")[0] + ) + # check that the migrated checkpoint matches the target checkpoint + assert migrated == target_checkpoint.checkpoint, ( + f"Checkpoint mismatch at index {idx}" + ) + + +@NEEDS_CONTEXTVARS +def test_latest_checkpoint_state_graph( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + builder = make_state_graph() + app = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert [ + *app.stream({"query": "what is weather in sf"}, config, durability="async") + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + { + "__interrupt__": ( + Interrupt( + value="", + id=AnyStr(), + ), + ) + }, + ] + + assert [*app.stream(Command(resume=""), config, durability="async")] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + # check history with current checkpoints matches expected history + history = [*app.get_state_history(config)] + expected_history = get_expected_history() + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + +@NEEDS_CONTEXTVARS +async def test_latest_checkpoint_state_graph_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: + builder = make_state_graph() + app = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app.astream( + {"query": "what is weather in sf"}, config, durability="async" + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + { + "__interrupt__": ( + Interrupt( + value="", + id=AnyStr(), + ), + ) + }, + ] + + assert [ + c async for c in app.astream(Command(resume=""), config, durability="async") + ] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + # check history with current checkpoints matches expected history + history = [c async for c in app.aget_state_history(config)] + expected_history = get_expected_history() + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"]) +def test_saved_checkpoint_state_graph( + sync_checkpointer: BaseCheckpointSaver, + checkpoint_version: str, +) -> None: + builder = make_state_graph() + app = builder.compile(checkpointer=sync_checkpointer) + + thread1 = "1" + config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} + + # save checkpoints + parent_id: str | None = None + for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): + grouped_writes = defaultdict(list) + for write in checkpoint.pending_writes: + grouped_writes[write[0]].append(write[1:]) + for tid, group in grouped_writes.items(): + sync_checkpointer.put_writes(checkpoint.config, group, tid) + sync_checkpointer.put( + patch_configurable(config, {"checkpoint_id": parent_id}), + checkpoint.checkpoint, + checkpoint.metadata, + checkpoint.checkpoint["channel_versions"], + ) + parent_id = checkpoint.checkpoint["id"] + + # load history + history = [*app.get_state_history(config)] + # check history with saved checkpoints matches expected history + exc_task_results: int = 0 + if checkpoint_version == "2-start:*": + exc_task_results = 1 + elif checkpoint_version == "2-quadratic": + exc_task_results = 2 + expected_history = get_expected_history(exc_task_results=exc_task_results) + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + # resume from 2nd to latest checkpoint + assert [*app.stream(Command(resume=""), history[1].config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + # new checkpoint should match the latest checkpoint in history + latest_state = app.get_state(config) + assert ( + StateSnapshot( + values=latest_state.values, + next=latest_state.next, + config=patch_configurable(latest_state.config, {"checkpoint_id": AnyStr()}), + metadata=AnyDict(latest_state.metadata), + created_at=AnyStr(), + parent_config=latest_state.parent_config, + tasks=latest_state.tasks, + interrupts=latest_state.interrupts, + ) + == history[0] + ) + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*", "2-quadratic"]) +async def test_saved_checkpoint_state_graph_async( + async_checkpointer: BaseCheckpointSaver, + checkpoint_version: str, +) -> None: + builder = make_state_graph() + app = builder.compile(checkpointer=async_checkpointer) + + thread1 = "1" + config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} + + # save checkpoints + parent_id: str | None = None + for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): + grouped_writes = defaultdict(list) + for write in checkpoint.pending_writes: + grouped_writes[write[0]].append(write[1:]) + for tid, group in grouped_writes.items(): + await async_checkpointer.aput_writes(checkpoint.config, group, tid) + await async_checkpointer.aput( + patch_configurable(config, {"checkpoint_id": parent_id}), + checkpoint.checkpoint, + checkpoint.metadata, + checkpoint.checkpoint["channel_versions"], + ) + parent_id = checkpoint.checkpoint["id"] + + # load history + history = [c async for c in app.aget_state_history(config)] + # check history with saved checkpoints matches expected history + exc_task_results: int = 0 + if checkpoint_version == "2-start:*": + exc_task_results = 1 + elif checkpoint_version == "2-quadratic": + exc_task_results = 2 + expected_history = get_expected_history(exc_task_results=exc_task_results) + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + # resume from 2nd to latest checkpoint + assert [c async for c in app.astream(Command(resume=""), history[1].config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + # new checkpoint should match the latest checkpoint in history + latest_state = await app.aget_state(config) + assert ( + StateSnapshot( + values=latest_state.values, + next=latest_state.next, + config=patch_configurable(latest_state.config, {"checkpoint_id": AnyStr()}), + metadata=AnyDict(latest_state.metadata), + created_at=AnyStr(), + parent_config=latest_state.parent_config, + tasks=latest_state.tasks, + interrupts=latest_state.interrupts, + ) + == history[0] + ) diff --git a/langgraph/source/libs/langgraph/tests/test_config_async.py b/langgraph/source/libs/langgraph/tests/test_config_async.py new file mode 100644 index 0000000000000000000000000000000000000000..a65a5f2af7c90c425aad9064318f5b56055ae94e --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_config_async.py @@ -0,0 +1,19 @@ +import pytest +from langchain_core.callbacks import AsyncCallbackManager + +from langgraph._internal._config import get_async_callback_manager_for_config + +pytestmark = pytest.mark.anyio + + +def test_new_async_manager_includes_tags() -> None: + config = {"callbacks": None} + manager = get_async_callback_manager_for_config(config, tags=["x", "y"]) + assert isinstance(manager, AsyncCallbackManager) + assert manager.inheritable_tags == ["x", "y"] + + +def test_new_async_manager_merges_tags_with_config() -> None: + config = {"callbacks": None, "tags": ["a"]} + manager = get_async_callback_manager_for_config(config, tags=["b"]) + assert manager.inheritable_tags == ["a", "b"] diff --git a/langgraph/source/libs/langgraph/tests/test_deprecation.py b/langgraph/source/libs/langgraph/tests/test_deprecation.py new file mode 100644 index 0000000000000000000000000000000000000000..919a7ca2703a342d1e1e72889430fc7c2cb4b83e --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_deprecation.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import warnings +from typing import Any, Optional + +import pytest +from langchain_core.runnables import RunnableConfig +from pytest_mock import MockerFixture +from typing_extensions import NotRequired, TypedDict + +from langgraph.channels.last_value import LastValue +from langgraph.errors import NodeInterrupt +from langgraph.func import entrypoint, task +from langgraph.graph import StateGraph +from langgraph.graph.message import MessageGraph +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.types import Interrupt, RetryPolicy +from langgraph.warnings import LangGraphDeprecatedSinceV05, LangGraphDeprecatedSinceV10 + + +class PlainState(TypedDict): ... + + +def test_add_node_retry_arg() -> None: + builder = StateGraph(PlainState) + + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.", + ): + builder.add_node("test_node", lambda state: state, retry=RetryPolicy()) # type: ignore[arg-type] + + +def test_task_retry_arg() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.", + ): + + @task(retry=RetryPolicy()) # type: ignore[arg-type] + def my_task(state: PlainState) -> PlainState: + return state + + +def test_entrypoint_retry_arg() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`retry` is deprecated and will be removed. Please use `retry_policy` instead.", + ): + + @entrypoint(retry=RetryPolicy()) # type: ignore[arg-type] + def my_entrypoint(state: PlainState) -> PlainState: + return state + + +def test_state_graph_input_schema() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`input` is deprecated and will be removed. Please use `input_schema` instead.", + ): + StateGraph(PlainState, input=PlainState) # type: ignore[arg-type] + + +def test_state_graph_output_schema() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`output` is deprecated and will be removed. Please use `output_schema` instead.", + ): + StateGraph(PlainState, output=PlainState) # type: ignore[arg-type] + + +def test_add_node_input_schema() -> None: + builder = StateGraph(PlainState) + + with pytest.warns( + LangGraphDeprecatedSinceV05, + match="`input` is deprecated and will be removed. Please use `input_schema` instead.", + ): + builder.add_node("test_node", lambda state: state, input=PlainState) # type: ignore[arg-type] + + +def test_constants_deprecation() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="Importing Send from langgraph.constants is deprecated. Please use 'from langgraph.types import Send' instead.", + ): + from langgraph.constants import Send # noqa: F401 + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="Importing Interrupt from langgraph.constants is deprecated. Please use 'from langgraph.types import Interrupt' instead.", + ): + from langgraph.constants import Interrupt # noqa: F401 + + +def test_pregel_types_deprecation() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="Importing from langgraph.pregel.types is deprecated. Please use 'from langgraph.types import ...' instead.", + ): + from langgraph.pregel.types import StateSnapshot # noqa: F401 + + +def test_config_schema_deprecation() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_schema` is deprecated and will be removed. Please use `context_schema` instead.", + ): + builder = StateGraph(PlainState, config_schema=PlainState) + assert builder.context_schema == PlainState + + builder.add_node("test_node", lambda state: state) + builder.set_entry_point("test_node") + graph = builder.compile() + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_schema` is deprecated. Use `get_context_jsonschema` for the relevant schema instead.", + ): + assert graph.config_schema() is not None + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`get_config_jsonschema` is deprecated. Use `get_context_jsonschema` instead.", + ): + graph.get_config_jsonschema() + + +def test_config_schema_deprecation_on_entrypoint() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_schema` is deprecated and will be removed. Please use `context_schema` instead.", + ): + + @entrypoint(config_schema=PlainState) # type: ignore[arg-type] + def my_entrypoint(state: PlainState) -> PlainState: + return state + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_schema` is deprecated. Use `get_context_jsonschema` for the relevant schema instead.", + ): + assert my_entrypoint.context_schema == PlainState + assert my_entrypoint.config_schema() is not None + + +@pytest.mark.filterwarnings("ignore:`config_type` is deprecated") +def test_config_type_deprecation_pregel(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_type` is deprecated and will be removed. Please use `context_schema` instead.", + ): + instance = Pregel( + nodes={ + "one": chain, + }, + channels={ + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + config_type=PlainState, + ) + assert instance.context_schema == PlainState + + +def test_interrupt_attributes_deprecation() -> None: + interrupt = Interrupt(value="question", id="abc") + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`interrupt_id` is deprecated. Use `id` instead.", + ): + interrupt.interrupt_id + + +def test_node_interrupt_deprecation() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="NodeInterrupt is deprecated. Please use `langgraph.types.interrupt` instead.", + ): + NodeInterrupt(value="test") + + +def test_deprecated_import() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="Importing PREVIOUS from langgraph.constants is deprecated. This constant is now private and should not be used directly.", + ): + from langgraph.constants import PREVIOUS # noqa: F401 + + +@pytest.mark.filterwarnings( + "ignore:`durability` has no effect when no checkpointer is present" +) +def test_checkpoint_during_deprecation_state_graph() -> None: + class CheckDurability(TypedDict): + durability: NotRequired[str] + + def plain_node(state: CheckDurability, config: RunnableConfig) -> CheckDurability: + return {"durability": config["configurable"]["__pregel_durability"]} + + builder = StateGraph(CheckDurability) + builder.add_node("plain_node", plain_node) + builder.set_entry_point("plain_node") + graph = builder.compile() + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", + ): + result = graph.invoke({}, checkpoint_during=True) + assert result["durability"] == "async" + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", + ): + result = graph.invoke({}, checkpoint_during=False) + assert result["durability"] == "exit" + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", + ): + for chunk in graph.stream({}, checkpoint_during=True): # type: ignore[arg-type] + assert chunk["plain_node"]["durability"] == "async" + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.", + ): + for chunk in graph.stream({}, checkpoint_during=False): # type: ignore[arg-type] + assert chunk["plain_node"]["durability"] == "exit" + + +def test_config_parameter_incorrect_typing() -> None: + """Test that a warning is raised when config parameter is typed incorrectly.""" + builder = StateGraph(PlainState) + + # Test sync function with config: dict + with pytest.warns( + UserWarning, + match="The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig | None', not '.*dict.*'. ", + ): + + def sync_node_with_dict_config(state: PlainState, config: dict) -> PlainState: + return state + + builder.add_node(sync_node_with_dict_config) + + # Test async function with config: dict + with pytest.warns( + UserWarning, + match="The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig | None', not '.*dict.*'. ", + ): + + async def async_node_with_dict_config( + state: PlainState, config: dict + ) -> PlainState: + return state + + builder.add_node(async_node_with_dict_config) + + # Test with other incorrect types + with pytest.warns( + UserWarning, + match="The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig | None', not '.*Any.*'. ", + ): + + def sync_node_with_any_config(state: PlainState, config: Any) -> PlainState: + return state + + builder.add_node(sync_node_with_any_config) + + with pytest.warns( + UserWarning, + match="The 'config' parameter should be typed as 'RunnableConfig' or 'RunnableConfig | None', not '.*Any.*'. ", + ): + + async def async_node_with_any_config( + state: PlainState, config: Any + ) -> PlainState: + return state + + builder.add_node(async_node_with_any_config) + + with warnings.catch_warnings(record=True) as w: + + def node_with_correct_config( + state: PlainState, config: RunnableConfig + ) -> PlainState: + return state + + builder.add_node(node_with_correct_config) + + def node_with_optional_config( + state: PlainState, + config: Optional[RunnableConfig], # noqa: UP045 + ) -> PlainState: + return state + + builder.add_node(node_with_optional_config) + + def node_with_untyped_config(state: PlainState, config) -> PlainState: + return state + + builder.add_node(node_with_untyped_config) + + async def async_node_with_correct_config( + state: PlainState, config: RunnableConfig + ) -> PlainState: + return state + + builder.add_node(async_node_with_correct_config) + + async def async_node_with_optional_config( + state: PlainState, + config: Optional[RunnableConfig], # noqa: UP045 + ) -> PlainState: + return state + + builder.add_node(async_node_with_optional_config) + + async def async_node_with_untyped_config( + state: PlainState, config + ) -> PlainState: + return state + + builder.add_node(async_node_with_untyped_config) + assert len(w) == 0 + + +def test_message_graph_deprecation() -> None: + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="MessageGraph is deprecated in LangGraph v1.0.0, to be removed in v2.0.0. Please use StateGraph with a `messages` key instead.", + ): + MessageGraph() diff --git a/langgraph/source/libs/langgraph/tests/test_interrupt_migration.py b/langgraph/source/libs/langgraph/tests/test_interrupt_migration.py new file mode 100644 index 0000000000000000000000000000000000000000..53f2269f220f71dae5f5eaa8488795f9faf39c93 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_interrupt_migration.py @@ -0,0 +1,50 @@ +import warnings + +import pytest +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + +from langgraph.types import Interrupt +from langgraph.warnings import LangGraphDeprecatedSinceV10 + + +@pytest.mark.filterwarnings("ignore:LangGraphDeprecatedSinceV10") +def test_interrupt_legacy_ns() -> None: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=LangGraphDeprecatedSinceV10) + + old_interrupt = Interrupt( + value="abc", resumable=True, when="during", ns=["a:b", "c:d"] + ) + + new_interrupt = Interrupt.from_ns(value="abc", ns="a:b|c:d") + assert new_interrupt.value == old_interrupt.value + assert new_interrupt.id == old_interrupt.id + + +serializer = JsonPlusSerializer(allowed_json_modules=True) + + +def test_serialization_roundtrip() -> None: + """Test that the legacy interrupt (pre v1) can be reserialized as the modern interrupt without id corruption.""" + + # generated with: + # JsonPlusSerializer().dumps_typed(Interrupt(value="legacy_test", ns=["legacy_test"], resumable=True, when="during")) + legacy_interrupt_bytes = b'{"lc": 2, "type": "constructor", "id": ["langgraph", "types", "Interrupt"], "kwargs": {"value": "legacy_test", "resumable": true, "ns": ["legacy_test"], "when": "during"}}' + legacy_interrupt_id = "f1fa625689ec006a5b32b76863e22a6c" + + interrupt = serializer.loads_typed(("json", legacy_interrupt_bytes)) + assert interrupt.id == legacy_interrupt_id + assert interrupt.value == "legacy_test" + + +def test_serialization_roundtrip_complex_ns() -> None: + """Test that the legacy interrupt (pre v1), with a more complex ns can be reserialized as the modern interrupt without id corruption.""" + + # generated with: + # JsonPlusSerializer().dumps_typed(Interrupt(value="legacy_test", ns=["legacy:test", "with:complex", "name:space"], resumable=True, when="during")) + legacy_interrupt_bytes = b'{"lc": 2, "type": "constructor", "id": ["langgraph", "types", "Interrupt"], "kwargs": {"value": "legacy_test", "resumable": true, "ns": ["legacy:test", "with:complex", "name:space"], "when": "during"}}' + legacy_interrupt_id = "e69356a9ee3630ee7f4f597f2693000c" + + interrupt = serializer.loads_typed(("json", legacy_interrupt_bytes)) + assert interrupt.id == legacy_interrupt_id + assert interrupt.value == "legacy_test" diff --git a/langgraph/source/libs/langgraph/tests/test_interruption.py b/langgraph/source/libs/langgraph/tests/test_interruption.py new file mode 100644 index 0000000000000000000000000000000000000000..a484d74e5adc4c6391835134a23ec8acbfab9326 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_interruption.py @@ -0,0 +1,92 @@ +import pytest +from langgraph.checkpoint.base import BaseCheckpointSaver +from typing_extensions import TypedDict + +from langgraph.graph import END, START, StateGraph +from langgraph.types import Durability + +pytestmark = pytest.mark.anyio + + +def test_interruption_without_state_updates( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Test interruption without state updates. This test confirms that + interrupting doesn't require a state key having been updated in the prev step""" + + class State(TypedDict): + input: str + + def noop(_state): + pass + + builder = StateGraph(State) + builder.add_node("step_1", noop) + builder.add_node("step_2", noop) + builder.add_node("step_3", noop) + builder.add_edge(START, "step_1") + builder.add_edge("step_1", "step_2") + builder.add_edge("step_2", "step_3") + builder.add_edge("step_3", END) + + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_after="*") + + initial_input = {"input": "hello world"} + thread = {"configurable": {"thread_id": "1"}} + + graph.invoke(initial_input, thread, durability=durability) + assert graph.get_state(thread).next == ("step_2",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (3 if durability != "exit" else 1) + + graph.invoke(None, thread, durability=durability) + assert graph.get_state(thread).next == ("step_3",) + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (4 if durability != "exit" else 2) + + graph.invoke(None, thread, durability=durability) + assert graph.get_state(thread).next == () + n_checkpoints = len([c for c in graph.get_state_history(thread)]) + assert n_checkpoints == (5 if durability != "exit" else 3) + + +async def test_interruption_without_state_updates_async( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Test interruption without state updates. This test confirms that + interrupting doesn't require a state key having been updated in the prev step""" + + class State(TypedDict): + input: str + + async def noop(_state): + pass + + builder = StateGraph(State) + builder.add_node("step_1", noop) + builder.add_node("step_2", noop) + builder.add_node("step_3", noop) + builder.add_edge(START, "step_1") + builder.add_edge("step_1", "step_2") + builder.add_edge("step_2", "step_3") + builder.add_edge("step_3", END) + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_after="*") + + initial_input = {"input": "hello world"} + thread = {"configurable": {"thread_id": "1"}} + + await graph.ainvoke(initial_input, thread, durability=durability) + assert (await graph.aget_state(thread)).next == ("step_2",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (3 if durability != "exit" else 1) + + await graph.ainvoke(None, thread, durability=durability) + assert (await graph.aget_state(thread)).next == ("step_3",) + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (4 if durability != "exit" else 2) + + await graph.ainvoke(None, thread, durability=durability) + assert (await graph.aget_state(thread)).next == () + n_checkpoints = len([c async for c in graph.aget_state_history(thread)]) + assert n_checkpoints == (5 if durability != "exit" else 3) diff --git a/langgraph/source/libs/langgraph/tests/test_large_cases.py b/langgraph/source/libs/langgraph/tests/test_large_cases.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb805c2ca3da7f533be186296430d4b70c5078a --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_large_cases.py @@ -0,0 +1,6970 @@ +import json +import operator +import re +import time +from dataclasses import replace +from typing import Annotated, Any, Literal, cast + +import pytest +from langchain_core.messages import AIMessage, AnyMessage, ToolCall +from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick +from langchain_core.tools import tool +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode +from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion +from typing_extensions import TypedDict + +from langgraph._internal._constants import PULL, PUSH +from langgraph.channels.last_value import LastValue +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from langgraph.graph.message import MessagesState, add_messages +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.types import ( + Command, + Durability, + Interrupt, + PregelTask, + RetryPolicy, + Send, + StateSnapshot, + StreamWriter, + interrupt, +) +from tests.agents import AgentAction, AgentFinish +from tests.any_int import AnyInt +from tests.any_str import AnyDict, AnyStr, UnsortedSequence +from tests.fake_chat import FakeChatModel +from tests.fake_tracer import FakeTracer +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + + +def test_invoke_two_processes_in_out_interrupt( + sync_checkpointer: BaseCheckpointSaver, + mocker: MockerFixture, +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=sync_checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert app.invoke(2, thread1, durability="async") is None + + # inbox == 3 + checkpoint = sync_checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert app.invoke(None, thread1, durability="async") == 4 + + # start execution again, stop at inbox + assert app.invoke(20, thread1, durability="async") is None + + # inbox == 21 + checkpoint = sync_checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert app.invoke(3, thread1, durability="async") is None + assert app.invoke(None, thread1, durability="async") == 5 + + # start execution again, stopping at inbox + assert app.invoke(20, thread2, durability="async") is None + + # inbox == 21 + snapshot = app.get_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + app.update_state(thread2, 25, as_node="one") + assert app.invoke(None, thread2) == 26 + + # no pending tasks + snapshot = app.get_state(thread2) + assert snapshot.next == () + + # list history + history = [c for c in app.get_state_history(thread1)] + assert len(history) == 8 + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 6, + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 4, + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 2, + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=history[7].config, + interrupts=(), + ), + StateSnapshot( + values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + # re-running from any previous checkpoint should re-run nodes + assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == [] + assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [ + {"two": {"output": 5}}, + ] + assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [ + {"one": {"inbox": 4}}, + {"__interrupt__": ()}, + ] + + +def test_fork_always_re_runs_nodes( + sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + add_one = mocker.Mock(side_effect=lambda _: 1) + + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) + graph = builder.compile(checkpointer=sync_checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + + # start execution, stop at inbox + assert [ + *graph.stream(1, thread1, stream_mode=["values", "updates"], durability="async") + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] + + # list history + history = [c for c in graph.get_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values=5, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values=4, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values=3, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values=2, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values=1, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values=0, + tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + # forking from any previous checkpoint should re-run nodes + assert [ + c for c in graph.stream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c for c in graph.stream(None, history[1].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + ] + assert [ + c for c in graph.stream(None, history[2].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] + + +def test_conditional_state_graph( + snapshot: SnapshotAssertion, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.language_models.fake import FakeStreamingListLLM + from langchain_core.prompts import PromptTemplate + from langchain_core.tools import tool + + class AgentState(TypedDict, total=False): + input: Annotated[str, UntrackedValue] + agent_outcome: AgentAction | AgentFinish | None + intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] + + class ToolState(TypedDict, total=False): + agent_outcome: AgentAction | AgentFinish + + # Assemble the tools + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + # Construct the agent + prompt = PromptTemplate.from_template("Hello!") + + llm = FakeStreamingListLLM( + responses=[ + "tool:search_api:query", + "tool:search_api:another", + "finish:answer", + ] + ) + + def agent_parser(input: str) -> dict[str, AgentAction | AgentFinish]: + if input.startswith("finish"): + _, answer = input.split(":") + return { + "agent_outcome": AgentFinish( + return_values={"answer": answer}, log=input + ) + } + else: + _, tool_name, tool_input = input.split(":") + return { + "agent_outcome": AgentAction( + tool=tool_name, tool_input=tool_input, log=input + ) + } + + agent = prompt | llm | agent_parser + + # Define tool execution logic + def execute_tools(data: ToolState) -> dict: + # check session in data + assert "input" not in data + assert "intermediate_steps" not in data + # execute the tool + agent_action: AgentAction = data.pop("agent_outcome") + observation = {t.name: t for t in tools}[agent_action.tool].invoke( + agent_action.tool_input + ) + return {"intermediate_steps": [[agent_action, observation]]} + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # check session in data + # Logic to decide whether to continue in the loop or exit + if isinstance(data["agent_outcome"], AgentFinish): + return "exit" + else: + return "continue" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("agent", agent) + workflow.add_node("tools", execute_tools, input_schema=ToolState) + + workflow.set_entry_point("agent") + + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + + workflow.add_edge("tools", "agent") + + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"input": "what is weather in sf"}) == { + "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" + ), + } + + assert [*app.stream({"input": "what is weather in sf"})] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + } + }, + { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] + + # test state get/update methods with interrupt_after + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + debug=True, + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c + for c in app_w_interrupt.stream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # test w interrupt before all + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before="*", + debug=True, + ) + config = {"configurable": {"thread_id": "3"}} + llm.i = 0 # reset the llm + + assert [ + c + for c in app_w_interrupt.stream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + # test w interrupt after all + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after="*", + ) + config = {"configurable": {"thread_id": "4"}} + llm.i = 0 # reset the llm + + assert [ + c + for c in app_w_interrupt.stream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "4", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + +def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: + from langchain_core.messages import AIMessage, HumanMessage + from langchain_core.tools import tool + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + model = FakeChatModel( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + app = create_react_agent(model, tools) + + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + id=AnyStr(), + ), + _AnyIdAIMessage(content="answer"), + ] + } + + events = [ + c + for c in app.stream( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="messages", + ) + ] + + assert events[:3] == [ + ( + _AnyIdAIMessageChunk( + content="", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + tool_call_chunks=[ + { + "name": "search_api", + "args": '{"query": "query"}', + "id": "tool_call123", + "index": None, + "type": "tool_call_chunk", + } + ], + chunk_position="last", + ), + { + "langgraph_step": 1, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ( + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + { + "langgraph_step": 2, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdAIMessageChunk( + content="", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another"}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one"}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + tool_call_chunks=[ + { + "name": "search_api", + "args": '{"query": "another"}', + "id": "tool_call234", + "index": None, + "type": "tool_call_chunk", + }, + { + "name": "search_api", + "args": '{"query": "a third one"}', + "id": "tool_call567", + "index": None, + "type": "tool_call_chunk", + }, + ], + chunk_position="last", + ), + { + "langgraph_step": 3, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ] + + assert events[3:5] == UnsortedSequence( + ( + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ) + assert events[5:] == [ + ( + _AnyIdAIMessageChunk( + content="answer", + chunk_position="last", + ), + { + "langgraph_step": 5, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ] + + assert app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]}, + {"recursion_limit": 2}, + debug=True, + ) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + _AnyIdAIMessage(content="Sorry, need more steps to process this request."), + ] + } + + model.i = 0 # reset the model + + invoke_updates_events = app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="updates", + ) + + stream_updates_events = [ + *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) + ] + + for output in (invoke_updates_events, stream_updates_events): + assert output[:3] == [ + { + "agent": { + "messages": [ + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + ] + } + }, + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + } + }, + { + "agent": { + "messages": [ + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + ] + } + }, + ] + assert output[3:5] == UnsortedSequence( + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + ] + } + }, + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + ] + } + }, + ) + assert output[5:] == [ + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} + ] + + +def test_state_graph_packets( + sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) + from langchain_core.tools import tool + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + time.sleep(0.1) + return f"result for {query}" + + tools = [search_api] + tools_by_name = {t.name: t for t in tools} + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + + def agent(data: AgentState) -> AgentState: + return { + "messages": model.invoke(data["messages"]), + "something_extra": "hi there", + } + + # Define decision-making logic + def should_continue(data: dict) -> str: + assert data["something_extra"] == "hi there", ( + "nodes can pass extra data to their cond edges, which isn't saved in state" + ) + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Send("tools", tool_call) for tool_call in tool_calls] + else: + return END + + def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState: + time.sleep(input["args"].get("idx", 0) / 10) + output = tools_by_name[input["name"]].invoke(input["args"], config) + return { + "messages": ToolMessage( + content=output, name=input["name"], tool_call_id=input["id"] + ) + } + + # Define a new graph + workflow = StateGraph(AgentState) + + # Define the two nodes we will cycle between + workflow.add_node("agent", agent) + workflow.add_node("tools", tools_node) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges("agent", should_continue) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert app.invoke({"messages": HumanMessage(content="what is weather in sf")}) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + AIMessage(content="answer", id="ai3"), + ] + } + + assert [ + c + for c in app.stream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + }, + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + } + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ) + }, + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + }, + }, + {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, + ] + + # interrupt after agent + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + {"messages": HumanMessage(content="what is weather in sf")}, + config, + durability="exit", + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + app_w_interrupt.update_state( + config, {"messages": last_message, "something_extra": "hi there"} + ) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + }, + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 + + assert [ + c + for c in app_w_interrupt.stream( + {"messages": HumanMessage(content="what is weather in sf")}, + config, + durability="exit", + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + app_w_interrupt.update_state( + config, {"messages": last_message, "something_extra": "hi there"} + ) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + }, + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + ), + interrupts=(), + ) + + +def test_message_graph( + snapshot: SnapshotAssertion, + deterministic_uuids: MockerFixture, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from copy import deepcopy + + from langchain_core.callbacks import CallbackManagerForLLMRun + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, BaseMessage, HumanMessage + from langchain_core.outputs import ChatGeneration, ChatResult + from langchain_core.tools import tool + + class FakeFunctionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + response = deepcopy(self.responses[self.i]) + if self.i < len(self.responses) - 1: + self.i += 1 + else: + self.i = 0 + generation = ChatGeneration(message=response) + return ChatResult(generations=[generation]) + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + model = FakeFunctionChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # Define the function that determines whether to continue or not + def should_continue(messages): + last_message = messages[-1] + # If there is no function call, then we finish + if not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + # Define a new graph + workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges( + # First, we define the start node. We use `agent`. + # This means these are the edges taken after the `agent` node is called. + "agent", + # Next, we pass in the function that will determine which node is called next. + should_continue, + # Finally we pass in a mapping. + # The keys are strings, and the values are other nodes. + # END is a special node marking that the graph should finish. + # What will happen is we will call `should_continue`, and then the output of that + # will be matched against the keys in this mapping. + # Based on which one it matches, that node will then be called. + { + # If `tools`, then we call the tool node. + "continue": "tools", + # Otherwise we finish. + "end": END, + }, + ) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke([HumanMessage(content="what is weather in sf")]) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + ("human", "what is weather in sf"), config, durability="exit" + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + next_config = app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=next_config, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), # replace existing message + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 # reset the llm + + assert [ + c + for c in app_w_interrupt.stream( + "what is weather in sf", config, durability="exit" + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # add an extra message as if it came from "tools" node + app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") + + # extra message is coerced BaseMessage and appended + # now the next node is "agent" per the graph edges + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ], + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + +def test_root_graph( + deterministic_uuids: MockerFixture, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from copy import deepcopy + + from langchain_core.callbacks import CallbackManagerForLLMRun + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolMessage, + ) + from langchain_core.outputs import ChatGeneration, ChatResult + from langchain_core.tools import tool + + class FakeFunctionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + response = deepcopy(self.responses[self.i]) + if self.i < len(self.responses) - 1: + self.i += 1 + else: + self.i = 0 + generation = ChatGeneration(message=response) + return ChatResult(generations=[generation]) + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + model = FakeFunctionChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # Define the function that determines whether to continue or not + def should_continue(messages): + last_message = messages[-1] + # If there is no function call, then we finish + if not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + class State(TypedDict): + __root__: Annotated[list[BaseMessage], add_messages] + + # Define a new graph + workflow = StateGraph(State) + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges( + # First, we define the start node. We use `agent`. + # This means these are the edges taken after the `agent` node is called. + "agent", + # Next, we pass in the function that will determine which node is called next. + should_continue, + # Finally we pass in a mapping. + # The keys are strings, and the values are other nodes. + # END is a special node marking that the graph should finish. + # What will happen is we will call `should_continue`, and then the output of that + # will be matched against the keys in this mapping. + # Based on which one it matches, that node will then be called. + { + # If `tools`, then we call the tool node. + "continue": "tools", + # Otherwise we finish. + "end": END, + }, + ) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert app.invoke(HumanMessage(content="what is weather in sf")) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + ToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + id="00000000-0000-4000-8000-000000000004", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + ToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + id="00000000-0000-4000-8000-000000000005", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + ("human", "what is weather in sf"), config, durability="exit" + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + next_config = app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=next_config, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), # replace existing message + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 # reset the llm + + assert [ + c + for c in app_w_interrupt.stream( + "what is weather in sf", config, durability="exit" + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = app_w_interrupt.get_state(config).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + app_w_interrupt.update_state(config, last_message) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + app_w_interrupt.update_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # add an extra message as if it came from "tools" node + app_w_interrupt.update_state(config, ("ai", "an extra message"), as_node="tools") + + # extra message is coerced BaseMessage and appended + # now the next node is "agent" per the graph edges + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ], + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + }, + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + ), + interrupts=(), + ) + + # create new graph with one more state key, reuse previous thread history + + def simple_add(left, right): + if not isinstance(right, list): + right = [right] + return left + right + + class MoreState(TypedDict): + __root__: Annotated[list[BaseMessage], simple_add] + something_else: str + + # Define a new graph + new_workflow = StateGraph(MoreState) + new_workflow.add_node( + "agent", RunnableMap(__root__=RunnablePick("__root__") | model) + ) + new_workflow.add_node( + "tools", RunnableMap(__root__=RunnablePick("__root__") | ToolNode(tools)) + ) + new_workflow.set_entry_point("agent") + new_workflow.add_conditional_edges( + "agent", + RunnablePick("__root__") | should_continue, + { + # If `tools`, then we call the tool node. + "continue": "tools", + # Otherwise we finish. + "end": END, + }, + ) + new_workflow.add_edge("tools", "agent") + new_app = new_workflow.compile(checkpointer=sync_checkpointer) + model.i = 0 # reset the llm + + # previous state is converted to new schema + assert new_app.get_state(config) == StateSnapshot( + values={ + "__root__": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + _AnyIdAIMessage(content="an extra message"), + ] + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 6, + }, + parent_config=(list(new_app.checkpointer.list(config, limit=2))[-1].config), + interrupts=(), + ) + + # new input is merged to old state + assert new_app.invoke( + { + "__root__": [HumanMessage(content="what is weather in la")], + "something_else": "value", + }, + config, + interrupt_before=["agent"], + ) == { + "__root__": [ + HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000008", + ), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "a different query"}, + "id": "tool_call123", + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + AIMessage( + content="an extra message", id="00000000-0000-4000-8000-000000000010" + ), + HumanMessage(content="what is weather in la"), + ], + "something_else": "value", + } + + +def test_in_one_fan_out_out_one_graph_state() -> None: + def sorted_add(x: list[str], y: list[str]) -> list[str]: + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def retriever_one(data: State) -> State: + # timer ensures stream output order is stable + # also, it confirms that the update order is not dependent on finishing order + # instead being defined by the order of the nodes/edges in the graph definition + # ie. stable between invocations + time.sleep(0.1) + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert [*app.stream({"query": "what is weather in sf"}, stream_mode="values")] == [ + {"query": "what is weather in sf", "docs": []}, + {"query": "query: what is weather in sf", "docs": []}, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, + ] + + assert [ + *app.stream( + {"query": "what is weather in sf"}, + stream_mode=["values", "updates", "debug"], + ) + ] == [ + ("values", {"query": "what is weather in sf", "docs": []}), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "input": {"query": "what is weather in sf", "docs": []}, + "triggers": ("branch:to:rewrite_query",), + }, + }, + ), + ("updates", {"rewrite_query": {"query": "query: what is weather in sf"}}), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "result": { + "query": "query: what is weather in sf", + }, + "error": None, + "interrupts": [], + }, + }, + ), + ("values", {"query": "query: what is weather in sf", "docs": []}), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_one",), + }, + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_two",), + }, + }, + ), + ( + "updates", + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + ), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "result": { + "docs": ["doc3", "doc4"], + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "updates", + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "result": { + "docs": ["doc1", "doc2"], + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "values", + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "qa", + "input": { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:qa",), + }, + }, + ), + ("updates", {"qa": {"answer": "doc1,doc2,doc3,doc4"}}), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "qa", + "result": { + "answer": "doc1,doc2,doc3,doc4", + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "values", + { + "query": "query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + ), + ] + + +def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_node_count = 0 + + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) == { + "my_key": "value ⛰️", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + interrupts=(), + ) + + +def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def tool_one(s: State) -> State: + return {"my_key": " one"} + + tool_two_node_count = 0 + + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + time.sleep(0.1) + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + def start(state: State) -> list[Send | str]: + return ["tool_two", Send("tool_one", state)] + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy()) + tool_two_graph.add_node("tool_one", tool_one) + tool_two_graph.set_conditional_entry_point(start) + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value one", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value one"} + + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good one", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "tool_one": {"my_key": " one"}, + }, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_one": {"my_key": " one"}, "__metadata__": {"cached": True}}, + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) == { + "my_key": "value ⛰️ one", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + id=AnyStr(), + name="tool_one", + path=("__pregel_push", 0, False), + result={"my_key": " one"}, + ), + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + + # clear the interrupt and next tasks + tool_two.update_state(thread1, None, as_node=END) + + # interrupt and unresolved tasks are cleared, finished tasks are kept + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config), + interrupts=(), + ) + + +def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None: + class SubgraphState(TypedDict): + my_key: str + market: str + + tool_two_node_count = 0 + + def tool_two_node(s: SubgraphState) -> SubgraphState: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + subgraph = StateGraph(SubgraphState) + subgraph.add_node("do", tool_two_node, retry_policy=RetryPolicy()) + subgraph.add_edge(START, "do") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", subgraph.compile()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + "__interrupt__": [ + Interrupt( + value="Just because...", + id=AnyStr(), + ) + ], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer", "market": "DE"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) == { + "my_key": "value ⛰️", + "market": "DE", + "__interrupt__": [ + Interrupt( + value="Just because...", + id=AnyStr(), + ) + ], + } + + assert [ + c.metadata + for c in tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + ) + ] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=( + list( + tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 + ) + )[-1].config + ), + interrupts=(), + ) + + +def test_send_dedupe_on_resume( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + interrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + time.sleep(0) + # sleep makes it more likely to trigger edge case where 1st task + # finishes before 2nd is registered in futures dict + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Command): + return replace(state, update=update) + else: + return update + + def send_for_fun(state): + return [ + Send("2", Command(goto=Send("2", 3))), + Send("2", Command(goto=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + graph = builder.compile(checkpointer=sync_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke(["0"], thread1, durability=durability) == { + "__interrupt__": [ + Interrupt( + value="Bahh", + id=AnyStr(), + ), + ], + } + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # check state + state = graph.get_state(thread1) + assert state.next == ("flaky",) + # check history + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (4 if durability != "exit" else 1) + + # resume execution + assert graph.invoke(None, thread1, durability=durability) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check state + state = graph.get_state(thread1) + assert state.next == () + # check history + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == (6 if durability != "exit" else 2) + expected_history = [ + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + ], + next=("3",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + ], + next=("2", "flaky", "3"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=("__pregel_push", 1, False), + error=None, + interrupts=(Interrupt(value="Bahh", id=AnyStr()),), + state=None, + result=["flaky|4"] if durability != "exit" else None, + ), + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + ), + interrupts=(Interrupt(value="Bahh", id=AnyStr()),), + ), + StateSnapshot( + values=["0", "1"], + next=("2", "2", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 1, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=["0"], + next=("1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + ), + ] + if durability != "exit": + assert history == expected_history + else: + assert history[0] == expected_history[0]._replace( + parent_config=history[1].config + ) + assert history[1] == expected_history[2]._replace(parent_config=None) + + +def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: + class InnerState(TypedDict): + my_key: str + my_other_key: str + + def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + other_parent_key: str + + def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node( + "inner", + inner.compile(interrupt_before=["inner_2"]), + ) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + app.invoke({"my_key": "my value"}, config, durability="exit") + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + expected = StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}}, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + assert app.get_state(config) == expected + assert list(app.get_state_history(config)) == [expected] + # now, get_state with subgraphs state + assert app.get_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state=StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + "inner_2", + (PULL, "inner_2"), + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": { + "": AnyStr(), + }, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + + # get_state_history for a subgraph returns its checkpoints + child_history = [*app.get_state_history(app.get_state(config).tasks[0].state)] + expected_child_history = [ + StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), + ), + ] + assert child_history == expected_child_history + + # resume + app.invoke(None, config, durability="exit") + # test state w/ nested subgraph state (after resuming from interrupt) + assert app.get_state(config) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + # test full history at the end + actual_history = list(app.get_state_history(config)) + expected_history = [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} + }, + result=None, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert app.get_state(actual_snapshot.config) == expected_snapshot + + +def test_doubly_nested_graph_state( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=sync_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert [ + c + for c in app.stream( + {"my_key": "my value"}, config, subgraphs=True, durability="exit" + ) + ] == [ + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] + # get state without subgraphs + outer_state = app.get_state(config) + assert outer_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + child_state = app.get_state(outer_state.tasks[0].state) + assert child_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + grandchild_state = app.get_state(child_state.tasks[0].state) + assert grandchild_state == StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + # get state with subgraphs + assert app.get_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state=StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state=StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + # # resume + assert [c for c in app.stream(None, config, subgraphs=True, durability="exit")] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ] + # get state with and without subgraphs + assert ( + app.get_state(config) + == app.get_state(config, subgraphs=True) + == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + ) + + # get outer graph history + outer_history = list(app.get_state_history(config)) + assert outer_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + result=None, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + # get child graph history + child_history = list(app.get_state_history(outer_history[1].tasks[0].state)) + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="child_1", + path=(PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + } + }, + result=None, + ), + ), + interrupts=(), + ), + ] + # get grandchild graph history + grandchild_history = list(app.get_state_history(child_history[0].tasks[0].state)) + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result=None, + ), + ), + interrupts=(), + ), + ] + + +def test_send_to_nested_graphs(sync_checkpointer: BaseCheckpointSaver) -> None: + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(JokeState, output_schema=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile(interrupt_before=["generate"]), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + tracer = FakeTracer() + + # invoke and pause at nested interrupt + assert graph.invoke( + {"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]} + ) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + assert len(tracer.runs) == 1, "Should produce exactly 1 root run" + + # check state + outer_state = graph.get_state(config) + + # update state of dogs joke graph + graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) + + # continue past interrupt + assert sorted( + graph.stream(None, config=config), + key=lambda d: d["generate_joke"]["jokes"][0], + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] + + +def test_send_react_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + def agent(state): + return {"messages": ai_message} + + def route(state): + if isinstance(state["messages"][-1], AIMessage): + return [ + Send(call["name"], call) for call in state["messages"][-1].tool_calls + ] + + foo_called = 0 + + def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + builder.add_conditional_edges("agent", route) + graph = builder.compile() + + assert graph.invoke({"messages": [HumanMessage("hello")]}) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert graph.invoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # remove the tool call, clearing the pending task + graph.update_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert graph.get_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ), + ] + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + tasks=(), + ) + + # tool call not executed + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + # interrupt-update-resume flow, creating new Send in update call + foo_called = 0 + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "3"}} + assert graph.invoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # replace the tool call, should clear previous send, create new one + graph.update_state( + thread1, + { + "messages": AIMessage( + "", + id=ai_message.id, + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ) + }, + ) + + # prev tool call no longer in pending tasks, new tool call is + assert graph.get_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # prev tool call not executed, new tool call is + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"), + ] + } + assert foo_called == 1 + + +def test_send_react_interrupt_control( + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + def agent(state) -> Command[Literal["foo"]]: + return Command( + update={"messages": ai_message}, + goto=[Send(call["name"], call) for call in ai_message.tool_calls], + ) + + foo_called = 0 + + def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + graph = builder.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert graph.get_graph().draw_mermaid() == snapshot + + assert graph.invoke({"messages": [HumanMessage("hello")]}) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert graph.invoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + interrupts=(), + ) + + # remove the tool call, clearing the pending task + graph.update_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert graph.get_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ), + ] + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + tasks=(), + ) + + # tool call not executed + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + # interrupt-update-resume flow, creating new Send in update call + + # TODO add here test with invoke(Command()) + + +def test_weather_subgraph( + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion +) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + + # setup subgraph + + @tool + def get_weather(city: str): + """Get the weather for a specific city""" + return f"I'ts sunny in {city}!" + + weather_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ) + ] + ) + + class SubGraphState(MessagesState): + city: str + + def model_node(state: SubGraphState, writer: StreamWriter): + writer(" very") + result = weather_model.invoke(state["messages"]) + return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} + + def weather_node(state: SubGraphState, writer: StreamWriter): + writer(" good") + result = get_weather.invoke({"city": state["city"]}) + return {"messages": [{"role": "assistant", "content": result}]} + + subgraph = StateGraph(SubGraphState) + subgraph.add_node(model_node) + subgraph.add_node(weather_node) + subgraph.add_edge(START, "model_node") + subgraph.add_edge("model_node", "weather_node") + subgraph.add_edge("weather_node", END) + subgraph = subgraph.compile(interrupt_before=["weather_node"]) + + # setup main graph + + class RouterState(MessagesState): + route: Literal["weather", "other"] + + router_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ) + ] + ) + + def router_node(state: RouterState, writer: StreamWriter): + writer("I'm") + system_message = "Classify the incoming query as either about weather or not." + messages = [{"role": "system", "content": system_message}] + state["messages"] + route = router_model.invoke(messages) + return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]} + + def normal_llm_node(state: RouterState): + return {"messages": [AIMessage("Hello!")]} + + def route_after_prediction(state: RouterState): + if state["route"] == "weather": + return "weather_graph" + else: + return "normal_llm_node" + + def weather_graph(state: RouterState): + return subgraph.invoke(state) + + graph = StateGraph(RouterState) + graph.add_node(router_node) + graph.add_node(normal_llm_node) + graph.add_node("weather_graph", weather_graph) + graph.add_edge(START, "router_node") + graph.add_conditional_edges( + "router_node", + route_after_prediction, + path_map=["weather_graph", "normal_llm_node"], + ) + graph.add_edge("normal_llm_node", END) + graph.add_edge("weather_graph", END) + graph = graph.compile(checkpointer=sync_checkpointer) + + if isinstance(sync_checkpointer, InMemorySaver): + assert graph.get_graph(xray=1).draw_mermaid() == snapshot + + config = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + + # run with custom output + assert [ + c for c in graph.stream(inputs, thread2, stream_mode="custom", subgraphs=True) + ] == [ + ((), "I'm"), + ((AnyStr("weather_graph:"),), " very"), + ] + assert [ + c for c in graph.stream(None, thread2, stream_mode="custom", subgraphs=True) + ] == [ + ((AnyStr("weather_graph:"),), " good"), + ] + + # run until interrupt + assert [ + c + for c in graph.stream( + inputs, + config=config, + stream_mode="updates", + subgraphs=True, + durability="exit", + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] + + # check current state + state = graph.get_state(config) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + interrupts=(), + ) + + # update + graph.update_state(state.tasks[0].state, {"city": "la"}) + + # run after update + assert [ + c + for c in graph.stream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (AnyStr("weather_graph:"),), + { + "weather_node": { + "messages": [{"role": "assistant", "content": "I'ts sunny in la!"}] + } + }, + ), + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="I'ts sunny in la!"), + ] + } + }, + ), + ] + + # try updating acting as weather node + config = {"configurable": {"thread_id": "14"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + assert [ + c + for c in graph.stream( + inputs, + config=config, + stream_mode="updates", + subgraphs=True, + durability="exit", + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] + state = graph.get_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf") + ], + "city": "San Francisco", + }, + next=("weather_node",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_node", + path=(PULL, "weather_node"), + ), + ), + interrupts=(), + ), + ), + ), + interrupts=(), + ) + graph.update_state( + state.tasks[0].state.config, + {"messages": [{"role": "assistant", "content": "rainy"}]}, + as_node="weather_node", + ) + state = graph.get_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="rainy"), + ], + "city": "San Francisco", + }, + next=(), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + } + ), + interrupts=(), + tasks=(), + ), + ), + ), + interrupts=(), + ) + assert [ + c + for c in graph.stream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="rainy"), + ] + } + }, + ), + ] + + # run with custom output, without subgraph streaming, should omit subgraph chunks + assert [ + c + for c in graph.stream( + inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom" + ) + ] == [ + "I'm", + ] + + # run with messages output, with subgraph streaming, should inc subgraph messages + assert [ + c + for c in graph.stream( + inputs, + {"configurable": {"thread_id": "4"}}, + stream_mode="messages", + subgraphs=True, + ) + ] == [ + ( + (), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ( + (AnyStr("weather_graph:"),), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "model_node", + "langgraph_triggers": ("branch:to:model_node",), + "langgraph_path": ("__pregel_pull", "model_node"), + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_ns": AnyStr("weather_graph:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ] + + # run with messages output, without subgraph streaming, should exc subgraph messages + assert [ + c + for c in graph.stream( + inputs, + {"configurable": {"thread_id": "5"}}, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "5", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ] + + +def test_subgraph_to_end_does_not_warn() -> None: + """Regression test for https://github.com/langchain-ai/langgraph/issues/5572.""" + + class State(TypedDict): + x: str + + def update_x(state: State): + return Command(goto=END, update={"x": state["x"] + "!"}) + + # Subgraph + subgraph_builder = StateGraph(State) + subgraph_builder.add_node("update_x", update_x) + subgraph_builder.add_edge(START, "update_x") + subgraph_builder.add_edge("update_x", END) + subgraph = subgraph_builder.compile() + + # Parent graph + builder = StateGraph(State) + builder.add_node("subgraph_node", subgraph) + builder.add_edge(START, "subgraph_node") + builder.add_edge("subgraph_node", END) + graph = builder.compile() + + response = graph.invoke({"x": "hello"}) + assert response == {"x": "hello!"} diff --git a/langgraph/source/libs/langgraph/tests/test_large_cases_async.py b/langgraph/source/libs/langgraph/tests/test_large_cases_async.py new file mode 100644 index 0000000000000000000000000000000000000000..b23bc2387d080539eec596f1e6956f0c5de25834 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_large_cases_async.py @@ -0,0 +1,4040 @@ +import asyncio +import operator +import re +import sys +from typing import ( + Annotated, + Literal, + cast, +) + +import pytest +from langchain_core.messages import AnyMessage, ToolCall +from langchain_core.runnables import RunnableConfig, RunnablePick +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode +from pytest_mock import MockerFixture +from typing_extensions import TypedDict + +from langgraph._internal._constants import PULL, PUSH +from langgraph.channels.last_value import LastValue +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.constants import END, START +from langgraph.graph.message import add_messages +from langgraph.graph.state import StateGraph +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter +from tests.any_int import AnyInt +from tests.any_str import AnyDict, AnyStr, UnsortedSequence +from tests.fake_chat import FakeChatModel +from tests.fake_tracer import FakeTracer +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + +pytestmark = pytest.mark.anyio + + +async def test_invoke_two_processes_in_out_interrupt( + async_checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert await app.ainvoke(2, thread1, durability="async") is None + + # inbox == 3 + checkpoint = await async_checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, thread1, durability="async") == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1, durability="async") is None + + # inbox == 21 + checkpoint = await async_checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, thread1, durability="async") is None + assert await app.ainvoke(None, thread1, durability="async") == 5 + + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2, durability="async") is None + + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 + + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () + + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert len(history) == 8 + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 6, + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 4, + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 2, + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=history[7].config, + interrupts=(), + ), + StateSnapshot( + values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + # forking from any previous checkpoint should re-run nodes + assert [ + c async for c in app.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"two": {"output": 5}}, + ] + assert [ + c async for c in app.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"one": {"inbox": 4}}, + {"__interrupt__": ()}, + ] + + +async def test_fork_always_re_runs_nodes( + async_checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + add_one = mocker.Mock(side_effect=lambda _: 1) + + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) + graph = builder.compile(checkpointer=async_checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + + # start execution, stop at inbox + assert [ + c + async for c in graph.astream( + 1, thread1, stream_mode=["values", "updates"], durability="async" + ) + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] + + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values=5, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values=4, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values=3, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values=2, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values=1, + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values=0, + tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + # forking from any previous checkpoint should re-run nodes + assert [ + c async for c in graph.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in graph.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + ] + assert [ + c async for c in graph.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] + + +async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.agents import AgentAction, AgentFinish + from langchain_core.language_models.fake import FakeStreamingListLLM + from langchain_core.prompts import PromptTemplate + from langchain_core.tools import tool + + class AgentState(TypedDict): + input: Annotated[str, UntrackedValue] + agent_outcome: AgentAction | AgentFinish | None + intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] + + # Assemble the tools + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + # Construct the agent + prompt = PromptTemplate.from_template("Hello!") + + llm = FakeStreamingListLLM( + responses=[ + "tool:search_api:query", + "tool:search_api:another", + "finish:answer", + ] + ) + + def agent_parser(input: str) -> dict[str, AgentAction | AgentFinish]: + if input.startswith("finish"): + _, answer = input.split(":") + return { + "agent_outcome": AgentFinish( + return_values={"answer": answer}, log=input + ) + } + else: + _, tool_name, tool_input = input.split(":") + return { + "agent_outcome": AgentAction( + tool=tool_name, tool_input=tool_input, log=input + ) + } + + agent = prompt | llm | agent_parser + + # Define tool execution logic + def execute_tools(data: AgentState) -> dict: + # execute the tool + agent_action: AgentAction = data.pop("agent_outcome") + observation = {t.name: t for t in tools}[agent_action.tool].invoke( + agent_action.tool_input + ) + return {"intermediate_steps": [[agent_action, observation]]} + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if isinstance(data["agent_outcome"], AgentFinish): + return "exit" + else: + return "continue" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("agent", agent) + workflow.add_node("tools", execute_tools) + + workflow.set_entry_point("agent") + + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + + workflow.add_edge("tools", "agent") + + app = workflow.compile() + + assert await app.ainvoke({"input": "what is weather in sf"}) == { + "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" + ), + } + + assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + } + }, + { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] + + 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 nodes) 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( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config, durability="exit" + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + +async def test_prebuilt_tool_chat() -> None: + from langchain_core.messages import AIMessage, HumanMessage + from langchain_core.tools import tool + + model = FakeChatModel( + messages=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + app = create_react_agent(model, tools) + + assert await app.ainvoke( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + id=AnyStr(), + ), + _AnyIdAIMessage(content="answer"), + ] + } + + events = [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="messages", + ) + ] + + assert events[:3] == [ + ( + _AnyIdAIMessageChunk( + content="", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + tool_call_chunks=[ + { + "name": "search_api", + "args": '{"query": "query"}', + "id": "tool_call123", + "index": None, + "type": "tool_call_chunk", + } + ], + chunk_position="last", + ), + { + "langgraph_step": 1, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ( + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + { + "langgraph_step": 2, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdAIMessageChunk( + content="", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another"}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one"}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + tool_call_chunks=[ + { + "name": "search_api", + "args": '{"query": "another"}', + "id": "tool_call234", + "index": None, + "type": "tool_call_chunk", + }, + { + "name": "search_api", + "args": '{"query": "a third one"}', + "id": "tool_call567", + "index": None, + "type": "tool_call_chunk", + }, + ], + chunk_position="last", + ), + { + "langgraph_step": 3, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ] + + assert events[3:5] == UnsortedSequence( + ( + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": (PUSH,), + "langgraph_path": (PUSH, AnyInt(), False), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ) + assert events[5:] == [ + ( + _AnyIdAIMessageChunk( + content="answer", + chunk_position="last", + ), + { + "langgraph_step": 5, + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ] + + stream_updates_events = [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] + assert stream_updates_events[:3] == [ + { + "agent": { + "messages": [ + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + ] + } + }, + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + } + }, + { + "agent": { + "messages": [ + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + ] + } + }, + ] + assert stream_updates_events[3:5] == UnsortedSequence( + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + ] + } + }, + { + "tools": { + "messages": [ + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + ] + } + }, + ) + assert stream_updates_events[5:] == [ + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} + ] + + +async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolMessage, + ) + from langchain_core.tools import tool + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + tools_by_name = {t.name: t for t in tools} + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Send("tools", tool_call) for tool_call in tool_calls] + else: + return END + + async def tools_node(input: ToolCall, config: RunnableConfig) -> AgentState: + await asyncio.sleep(input["args"].get("idx", 0) / 10) + output = await tools_by_name[input["name"]].ainvoke(input["args"], config) + return { + "messages": ToolMessage( + content=output, name=input["name"], tool_call_id=input["id"] + ) + } + + # Define a new graph + workflow = StateGraph(AgentState) + + # Define the two nodes we will cycle between + workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) + workflow.add_node("tools", tools_node) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges("agent", should_continue) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert await app.ainvoke( + {"messages": HumanMessage(content="what is weather in sf")} + ) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + AIMessage(content="answer", id="ai3"), + ] + } + + assert [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + }, + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + } + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ) + }, + }, + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a third one", + name="search_api", + tool_call_id="tool_call567", + ), + }, + }, + {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, + ] + + # interrupt after agent + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, + config, + durability="exit", + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, + config, + durability="exit", + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + +async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, HumanMessage + from langchain_core.tools import tool + + class FakeFunctionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + model = FakeFunctionChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + AIMessage(content="answer", id="ai3"), + ] + ) + + # Define the function that determines whether to continue or not + def should_continue(messages): + last_message = messages[-1] + # If there is no function call, then we finish + if not last_message.tool_calls: + return "end" + # Otherwise if there is, we continue + else: + return "continue" + + # Define a new graph + workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] + + # Define the two nodes we will cycle between + workflow.add_node("agent", model) + workflow.add_node("tools", ToolNode(tools)) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges( + # First, we define the start node. We use `agent`. + # This means these are the edges taken after the `agent` node is called. + "agent", + # Next, we pass in the function that will determine which node is called next. + should_continue, + # Finally we pass in a mapping. + # The keys are strings, and the values are other nodes. + # END is a special node marking that the graph should finish. + # What will happen is we will call `should_continue`, and then the output of that + # will be matched against the keys in this mapping. + # Based on which one it matches, that node will then be called. + { + # If `tools`, then we call the tool node. + "continue": "tools", + # Otherwise we finish. + "end": END, + }, + ) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert await app.ainvoke([HumanMessage(content="what is weather in sf")]) == [ + _AnyIdHumanMessage( + content="what is weather in sf", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", # respects ids passed in + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ), + AIMessage(content="answer", id="ai3"), + ] + + assert [ + c async for c in app.astream([HumanMessage(content="what is weather in sf")]) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + { + "tools": [ + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call456", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), + config, + durability="exit", + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + parent_config=None, + interrupts=(), + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": [ + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + +async def test_in_one_fan_out_out_one_graph_state() -> None: + def sorted_add(x: list[str], y: list[str]) -> list[str]: + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], operator.add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def retriever_one(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert [ + c + async for c in app.astream( + {"query": "what is weather in sf"}, stream_mode="values" + ) + ] == [ + {"query": "what is weather in sf", "docs": []}, + {"query": "query: what is weather in sf", "docs": []}, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, + ] + + assert [ + c + async for c in app.astream( + {"query": "what is weather in sf"}, + stream_mode=["values", "updates", "debug"], + ) + ] == [ + ("values", {"query": "what is weather in sf", "docs": []}), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "input": {"query": "what is weather in sf", "docs": []}, + "triggers": ("branch:to:rewrite_query",), + }, + }, + ), + ("updates", {"rewrite_query": {"query": "query: what is weather in sf"}}), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "result": { + "query": "query: what is weather in sf", + }, + "error": None, + "interrupts": [], + }, + }, + ), + ("values", {"query": "query: what is weather in sf", "docs": []}), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_one",), + }, + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_two",), + }, + }, + ), + ( + "updates", + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + ), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "result": { + "docs": ["doc3", "doc4"], + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "updates", + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "result": { + "docs": ["doc1", "doc2"], + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "values", + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "qa", + "input": { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:qa",), + }, + }, + ), + ("updates", {"qa": {"answer": "doc1,doc2,doc3,doc4"}}), + ( + "debug", + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "qa", + "result": { + "answer": "doc1,doc2,doc3,doc4", + }, + "error": None, + "interrupts": [], + }, + }, + ), + ( + "values", + { + "query": "query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + ), + ] + + +async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> None: + class InnerState(TypedDict): + my_key: str + my_other_key: str + + def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + other_parent_key: str + + def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node( + "inner", + inner.compile(interrupt_before=["inner_2"]), + ) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, durability="exit") + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + expected = StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}}, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + assert await app.aget_state(config) == expected + # now, get_state with subgraphs state + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state=StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + "inner_2", + (PULL, "inner_2"), + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": { + "": AnyStr(), + }, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + # get_state_history returns outer graph checkpoints + assert [c async for c in app.aget_state_history(config)] == [expected] + + # get_state_history for a subgraph returns its checkpoints + child_history = [ + c + async for c in app.aget_state_history( + (await app.aget_state(config)).tasks[0].state + ) + ] + expected_child_history = [ + StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), + interrupts=(), + ), + ] + + assert child_history == expected_child_history + + # resume + await app.ainvoke(None, config, durability="exit") + # test state w/ nested subgraph state (after resuming from interrupt) + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + # test full history at the end + actual_history = [c async for c in app.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + result=None, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert await app.aget_state(actual_snapshot.config) == expected_snapshot + + +async def test_doubly_nested_graph_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=async_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, subgraphs=True, durability="exit" + ) + ] == [ + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] + # get state without subgraphs + outer_state = await app.aget_state(config) + assert outer_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + child_state = await app.aget_state(outer_state.tasks[0].state) + assert child_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + grandchild_state = await app.aget_state(child_state.tasks[0].state) + assert grandchild_state == StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + # get state with subgraphs + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state=StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state=StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "step": 0, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ) + # resume + assert [ + c async for c in app.astream(None, config, subgraphs=True, durability="exit") + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ( + (AnyStr("child:"),), + {"child_1": {"my_key": "hi my value here and there"}}, + ), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ] + # get state with and without subgraphs + assert ( + await app.aget_state(config) + == await app.aget_state(config, subgraphs=True) + == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + ) + + # get outer graph history + outer_history = [c async for c in app.aget_state_history(config)] + assert outer_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + # get child graph history + child_history = [ + c async for c in app.aget_state_history(outer_history[1].tasks[0].state) + ] + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="child_1", + path=(PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + } + }, + result=None, + ), + ), + interrupts=(), + ), + ] + # get grandchild graph history + grandchild_history = [ + c async for c in app.aget_state_history(child_history[0].tasks[0].state) + ] + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here"}, + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result=None, + ), + ), + interrupts=(), + ), + ] + + +async def test_send_to_nested_graphs(async_checkpointer: BaseCheckpointSaver) -> None: + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + async def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + async def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(JokeState, output_schema=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile(interrupt_before=["generate"]), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + tracer = FakeTracer() + + # invoke and pause at nested interrupt + assert await graph.ainvoke( + {"subjects": ["cats", "dogs"]}, + config={**config, "callbacks": [tracer]}, + ) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + assert len(tracer.runs) == 1, "Should produce exactly 1 root run" + + # check state + outer_state = await graph.aget_state(config) + + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + ) + + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + } + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +async def test_weather_subgraph( + async_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, ToolCall + from langchain_core.tools import tool + + from langgraph.graph import MessagesState + + # setup subgraph + + @tool + def get_weather(city: str): + """Get the weather for a specific city""" + return f"I'ts sunny in {city}!" + + weather_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ) + ] + ) + + class SubGraphState(MessagesState): + city: str + + def model_node(state: SubGraphState, writer: StreamWriter): + writer(" very") + result = weather_model.invoke(state["messages"]) + return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} + + def weather_node(state: SubGraphState, writer: StreamWriter): + writer(" good") + result = get_weather.invoke({"city": state["city"]}) + return {"messages": [{"role": "assistant", "content": result}]} + + subgraph = StateGraph(SubGraphState) + subgraph.add_node(model_node) + subgraph.add_node(weather_node) + subgraph.add_edge(START, "model_node") + subgraph.add_edge("model_node", "weather_node") + subgraph.add_edge("weather_node", END) + subgraph = subgraph.compile(interrupt_before=["weather_node"]) + + # setup main graph + + class RouterState(MessagesState): + route: Literal["weather", "other"] + + router_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ) + ] + ) + + def router_node(state: RouterState, writer: StreamWriter): + writer("I'm") + system_message = "Classify the incoming query as either about weather or not." + messages = [{"role": "system", "content": system_message}] + state["messages"] + route = router_model.invoke(messages) + return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]} + + def normal_llm_node(state: RouterState): + return {"messages": [AIMessage("Hello!")]} + + def route_after_prediction(state: RouterState): + if state["route"] == "weather": + return "weather_graph" + else: + return "normal_llm_node" + + def weather_graph(state: RouterState): + # this tests that all async checkpointers tested also implement sync methods + # as the subgraph called with sync invoke will use sync checkpointer methods + return subgraph.invoke(state) + + graph = StateGraph(RouterState) + graph.add_node(router_node) + graph.add_node(normal_llm_node) + graph.add_node("weather_graph", weather_graph) + graph.add_edge(START, "router_node") + graph.add_conditional_edges( + "router_node", + route_after_prediction, + path_map=["weather_graph", "normal_llm_node"], + ) + graph.add_edge("normal_llm_node", END) + graph.add_edge("weather_graph", END) + + def get_first_in_list(): + return [*graph.get_state_history(config, limit=1)][0] + + graph = graph.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + + # run with custom output + assert [ + c + async for c in graph.astream( + inputs, thread2, stream_mode="custom", subgraphs=True + ) + ] == [ + ((), "I'm"), + ((AnyStr("weather_graph:"),), " very"), + ] + assert [ + c + async for c in graph.astream( + None, thread2, stream_mode="custom", subgraphs=True + ) + ] == [ + ((AnyStr("weather_graph:"),), " good"), + ] + + # run until interrupt + assert [ + c + async for c in graph.astream( + inputs, + config=config, + stream_mode="updates", + subgraphs=True, + durability="exit", + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] + + # check current state + state = await graph.aget_state(config) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + interrupts=(), + ) + # confirm that list() delegates to alist() correctly + assert await asyncio.to_thread(get_first_in_list) == state + + # update + await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) + + # run after update + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (AnyStr("weather_graph:"),), + { + "weather_node": { + "messages": [{"role": "assistant", "content": "I'ts sunny in la!"}] + } + }, + ), + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="I'ts sunny in la!"), + ] + } + }, + ), + ] + + # try updating acting as weather node + config = {"configurable": {"thread_id": "14"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + assert [ + c + async for c in graph.astream( + inputs, + config=config, + stream_mode="updates", + subgraphs=True, + durability="exit", + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf") + ], + "city": "San Francisco", + }, + next=("weather_node",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_node", + path=(PULL, "weather_node"), + ), + ), + ), + ), + ), + interrupts=(), + ) + await graph.aupdate_state( + state.tasks[0].state.config, + {"messages": [{"role": "assistant", "content": "rainy"}]}, + as_node="weather_node", + ) + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="rainy"), + ], + "city": "San Francisco", + }, + next=(), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + } + ), + tasks=(), + interrupts=(), + ), + ), + ), + ) + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="rainy"), + ] + } + }, + ), + ] + + # run with custom output, without subgraph streaming, should omit subgraph chunks + assert [ + c + async for c in graph.astream( + inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom" + ) + ] == [ + "I'm", + ] + + # run with messages output, with subgraph streaming, should inc subgraph messages + assert [ + c + async for c in graph.astream( + inputs, + {"configurable": {"thread_id": "4"}}, + stream_mode="messages", + subgraphs=True, + ) + ] == [ + ( + (), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ( + (AnyStr("weather_graph:"),), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "model_node", + "langgraph_triggers": ("branch:to:model_node",), + "langgraph_path": ("__pregel_pull", "model_node"), + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_ns": AnyStr("weather_graph:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ] + + # run with messages output, without subgraph streaming, should exc subgraph messages + assert [ + c + async for c in graph.astream( + inputs, + {"configurable": {"thread_id": "5"}}, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "5", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ] diff --git a/langgraph/source/libs/langgraph/tests/test_managed_values.py b/langgraph/source/libs/langgraph/tests/test_managed_values.py new file mode 100644 index 0000000000000000000000000000000000000000..0202a1400114cdc4c2e708f1ef7b4e71ccbe9f8c --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_managed_values.py @@ -0,0 +1,27 @@ +from typing_extensions import NotRequired, Required, TypedDict + +from langgraph.graph import StateGraph +from langgraph.managed import RemainingSteps + + +class StatePlain(TypedDict): + remaining_steps: RemainingSteps + + +class StateNotRequired(TypedDict): + remaining_steps: NotRequired[RemainingSteps] + + +class StateRequired(TypedDict): + remaining_steps: Required[RemainingSteps] + + +def test_managed_values_recognized() -> None: + graph = StateGraph(StatePlain) + assert "remaining_steps" in graph.managed + + graph = StateGraph(StateNotRequired) + assert "remaining_steps" in graph.managed + + graph = StateGraph(StateRequired) + assert "remaining_steps" in graph.managed diff --git a/langgraph/source/libs/langgraph/tests/test_messages_state.py b/langgraph/source/libs/langgraph/tests/test_messages_state.py new file mode 100644 index 0000000000000000000000000000000000000000..f41552b16160448cab771cfd898f87f6028788b9 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_messages_state.py @@ -0,0 +1,369 @@ +from typing import Annotated +from uuid import UUID + +import langchain_core +import pytest +from langchain_core.messages import ( + AIMessage, + AnyMessage, + HumanMessage, + RemoveMessage, + SystemMessage, + ToolMessage, +) +from pydantic import BaseModel +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import add_messages +from langgraph.graph.message import REMOVE_ALL_MESSAGES, MessagesState, push_message +from langgraph.graph.state import StateGraph +from tests.messages import _AnyIdHumanMessage + +_, CORE_MINOR, CORE_PATCH = ( + int("".join(c for c in v if c.isdigit())) + for v in langchain_core.__version__.split(".") +) + + +def test_add_single_message(): + left = [HumanMessage(content="Hello", id="1")] + right = AIMessage(content="Hi there!", id="2") + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + assert result == expected_result + + +def test_add_multiple_messages(): + left = [HumanMessage(content="Hello", id="1")] + right = [ + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + ] + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + ] + assert result == expected_result + + +def test_update_existing_message(): + left = [HumanMessage(content="Hello", id="1")] + right = HumanMessage(content="Hello again", id="1") + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello again", id="1")] + assert result == expected_result + + +def test_missing_ids(): + left = [HumanMessage(content="Hello")] + right = [AIMessage(content="Hi there!")] + result = add_messages(left, right) + assert len(result) == 2 + assert all(isinstance(m.id, str) and UUID(m.id, version=4) for m in result) + + +def test_duplicates_in_input(): + left = [] + right = [ + AIMessage(id="1", content="Hi there!"), + AIMessage(id="1", content="Hi there again!"), + ] + result = add_messages(left, right) + assert len(result) == 1 + assert result[0].id == "1" + assert result[0].content == "Hi there again!" + + +def test_duplicates_in_input_with_remove(): + left = [AIMessage(id="1", content="Hello!")] + right = [ + RemoveMessage(id="1"), + AIMessage(id="1", content="Hi there!"), + AIMessage(id="1", content="Hi there again!"), + ] + result = add_messages(left, right) + assert len(result) == 1 + assert result[0].id == "1" + assert result[0].content == "Hi there again!" + + +def test_remove_message(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = RemoveMessage(id="2") + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello", id="1")] + assert result == expected_result + + +def test_duplicate_remove_message(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [RemoveMessage(id="2"), RemoveMessage(id="2")] + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello", id="1")] + assert result == expected_result + + +def test_remove_nonexistent_message(): + left = [HumanMessage(content="Hello", id="1")] + right = RemoveMessage(id="2") + with pytest.raises( + ValueError, match="Attempting to delete a message with an ID that doesn't exist" + ): + add_messages(left, right) + + +def test_mixed_operations(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [ + HumanMessage(content="Updated hello", id="1"), + RemoveMessage(id="2"), + SystemMessage(content="New message", id="3"), + ] + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Updated hello", id="1"), + SystemMessage(content="New message", id="3"), + ] + assert result == expected_result + + +def test_empty_inputs(): + assert add_messages([], []) == [] + assert add_messages([], [HumanMessage(content="Hello", id="1")]) == [ + HumanMessage(content="Hello", id="1") + ] + assert add_messages([HumanMessage(content="Hello", id="1")], []) == [ + HumanMessage(content="Hello", id="1") + ] + + +def test_non_list_inputs(): + left = HumanMessage(content="Hello", id="1") + right = AIMessage(content="Hi there!", id="2") + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + assert result == expected_result + + +def test_delete_all(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [ + RemoveMessage(id="1"), + RemoveMessage(id="2"), + ] + result = add_messages(left, right) + expected_result = [] + assert result == expected_result + + +class MessagesStatePydantic(BaseModel): + messages: Annotated[list[AnyMessage], add_messages] + + +MESSAGES_STATE_SCHEMAS = [MessagesState, MessagesStatePydantic] + + +@pytest.mark.parametrize("state_schema", MESSAGES_STATE_SCHEMAS) +def test_messages_state(state_schema): + def foo(state): + return {"messages": [HumanMessage("foo")]} + + graph = StateGraph(state_schema) + graph.add_edge(START, "foo") + graph.add_edge("foo", END) + graph.add_node(foo) + + app = graph.compile() + + assert app.invoke({"messages": [("user", "meow")]}) == { + "messages": [ + _AnyIdHumanMessage(content="meow"), + _AnyIdHumanMessage(content="foo"), + ] + } + + +@pytest.mark.skipif( + condition=not ((CORE_MINOR == 3 and CORE_PATCH >= 11) or CORE_MINOR > 3), + reason="Requires langchain_core>=0.3.11.", +) +def test_messages_state_format_openai(): + class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages(format="langchain-openai")] + + def foo(state): + messages = [ + HumanMessage( + content=[ + { + "type": "text", + "text": "Here's an image:", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "1234", + }, + }, + ] + ), + AIMessage( + content=[ + { + "type": "tool_use", + "name": "foo", + "input": {"bar": "baz"}, + "id": "1", + } + ] + ), + HumanMessage( + content=[ + { + "type": "tool_result", + "tool_use_id": "1", + "is_error": False, + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "1234", + }, + }, + ], + } + ] + ), + ] + return {"messages": messages} + + expected = [ + HumanMessage(content="meow"), + HumanMessage( + content=[ + {"type": "text", "text": "Here's an image:"}, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,1234"}, + }, + ], + ), + AIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "type": "tool_calls", + "args": {"bar": "baz"}, + "id": "1", + } + ], + ), + ToolMessage( + content=[ + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,1234"}, + } + ], + tool_call_id="1", + ), + ] + + graph = StateGraph(State) + graph.add_edge(START, "foo") + graph.add_edge("foo", END) + graph.add_node(foo) + + app = graph.compile() + + result = app.invoke({"messages": [("user", "meow")]}) + for m in result["messages"]: + m.id = None + assert result == {"messages": expected} + + +def test_remove_all_messages(): + # simple removal + left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")] + right = [RemoveMessage(id=REMOVE_ALL_MESSAGES)] + result = add_messages(left, right) + assert result == [] + + # removal and update (i.e., overwriting) + left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")] + right = [ + RemoveMessage(id=REMOVE_ALL_MESSAGES), + HumanMessage(content="Updated hello"), + ] + result = add_messages(left, right) + assert result == [_AnyIdHumanMessage(content="Updated hello")] + + # test removing preceding messages in the right list + left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")] + right = [ + HumanMessage(content="Updated hello"), + RemoveMessage(id=REMOVE_ALL_MESSAGES), + HumanMessage(content="Updated hi there"), + ] + result = add_messages(left, right) + assert result == [ + _AnyIdHumanMessage(content="Updated hi there"), + ] + + +def test_push_messages_in_graph(): + class MessagesState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + def chat(_: MessagesState) -> MessagesState: + with pytest.raises(ValueError, match="Message ID is required"): + push_message(AIMessage(content="No ID")) + + push_message(AIMessage(content="First", id="1")) + push_message(HumanMessage(content="Second", id="2")) + push_message(AIMessage(content="Third", id="3")) + + builder = StateGraph(MessagesState) + builder.add_node(chat) + builder.add_edge(START, "chat") + + graph = builder.compile() + + messages, values = [], None + for event, chunk in graph.stream( + {"messages": []}, stream_mode=["messages", "values"] + ): + if event == "values": + values = chunk + elif event == "messages": + message, _ = chunk + messages.append(message) + + assert values["messages"] == messages diff --git a/langgraph/source/libs/langgraph/tests/test_pregel.py b/langgraph/source/libs/langgraph/tests/test_pregel.py new file mode 100644 index 0000000000000000000000000000000000000000..4ce5ca24da3d11b9571ddfc71a97437928398547 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_pregel.py @@ -0,0 +1,9061 @@ +import enum +import functools +import gc +import json +import logging +import operator +import threading +import time +import uuid +from collections import Counter, deque +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from random import randrange +from typing import Annotated, Any, Literal, get_type_hints + +import pytest +from langchain_core.language_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage +from langchain_core.runnables import ( + RunnableConfig, + RunnableLambda, + RunnablePassthrough, +) +from langchain_core.runnables.graph import Edge +from langgraph.cache.base import BaseCache +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, +) +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.prebuilt.tool_node import ToolNode +from langgraph.store.base import BaseStore +from langsmith import traceable +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion +from typing_extensions import NotRequired, TypedDict + +from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.ephemeral_value import EphemeralValue +from langgraph.channels.last_value import LastValue +from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.config import get_stream_writer +from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentCommand +from langgraph.func import entrypoint, task +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import MessagesState, add_messages +from langgraph.pregel import ( + NodeBuilder, + Pregel, +) +from langgraph.pregel._loop import SyncPregelLoop +from langgraph.pregel._runner import PregelRunner +from langgraph.types import ( + CachePolicy, + Command, + Durability, + Interrupt, + Overwrite, + PregelTask, + RetryPolicy, + Send, + StateSnapshot, + StateUpdate, + StreamWriter, + interrupt, +) +from tests.agents import AgentAction, AgentFinish +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + +pytestmark = pytest.mark.anyio + +logger = logging.getLogger(__name__) + + +def test_graph_validation() -> None: + class State(TypedDict): + hello: str + + graph = StateGraph(State) + graph.add_node("start", lambda x: x) + graph.add_edge("__start__", "start") + graph.add_edge("unknown", "start") + graph.add_edge("start", "__end__") + with pytest.raises(ValueError, match="Found edge starting at unknown node "): + graph.compile() + + def bad_reducer(a): ... + + class BadReducerState(TypedDict): + hello: Annotated[str, bad_reducer] + + with pytest.raises(ValueError, match="Invalid reducer"): + StateGraph(BadReducerState) + + def node_b(state: State) -> State: + return {"hello": "world"} + + builder = StateGraph(State) + builder.add_node("a", node_b) + builder.add_node("b", node_b) + builder.add_node("c", node_b) + builder.set_entry_point("a") + builder.add_edge("a", "b") + builder.add_edge("a", "c") + graph = builder.compile() + + with pytest.raises(InvalidUpdateError, match="At key 'hello'"): + graph.invoke({"hello": "there"}) + + +def test_invalid_checkpointer_type() -> None: + class State(TypedDict): + foo: str + + builder = StateGraph(State) + builder.add_node("start", lambda state: state) + builder.set_entry_point("start") + builder.set_finish_point("start") + + class NotACheckpointer: + pass + + with pytest.raises(TypeError, match="Invalid checkpointer provided"): + builder.compile(checkpointer=NotACheckpointer()) + + +def test_graph_validation_with_command() -> None: + class State(TypedDict): + foo: str + bar: str + + def node_a(state: State): + return Command(goto="b", update={"foo": "bar"}) + + def node_b(state: State): + return Command(goto=END, update={"bar": "baz"}) + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.add_node("b", node_b) + builder.add_edge(START, "a") + graph = builder.compile() + assert graph.invoke({"foo": ""}) == {"foo": "bar", "bar": "baz"} + + +def test_checkpoint_errors() -> None: + class FaultyGetCheckpointer(InMemorySaver): + def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + raise ValueError("Faulty get_tuple") + + class FaultyPutCheckpointer(InMemorySaver): + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: dict[str, str | int | float] | None = None, + ) -> RunnableConfig: + raise ValueError("Faulty put") + + class FaultyPutWritesCheckpointer(InMemorySaver): + def put_writes( + self, config: RunnableConfig, writes: list[tuple[str, Any]], task_id: str + ) -> RunnableConfig: + raise ValueError("Faulty put_writes") + + class FaultyVersionCheckpointer(InMemorySaver): + def get_next_version(self, current: int | None, channel: None) -> int: + raise ValueError("Faulty get_next_version") + + def logic(inp: str) -> str: + return "" + + builder = StateGraph(Annotated[str, operator.add]) + builder.add_node("agent", logic) + builder.add_edge(START, "agent") + + graph = builder.compile(checkpointer=FaultyGetCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_tuple"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + graph = builder.compile(checkpointer=FaultyPutCheckpointer()) + with pytest.raises(ValueError, match="Faulty put"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_next_version"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + # add parallel node + builder.add_node("parallel", logic) + builder.add_edge(START, "parallel") + graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer()) + with pytest.raises(ValueError, match="Faulty put_writes"): + graph.invoke( + "", {"configurable": {"thread_id": "thread-1"}}, durability="async" + ) + + +def test_context_json_schema() -> None: + """Test that config json schema is generated properly.""" + chain = NodeBuilder().subscribe_only("input").write_to("output") + + @dataclass + class Foo: + x: int + y: str = field(default="foo") + + app = Pregel( + nodes={ + "one": chain, + }, + channels={ + "ephemeral": EphemeralValue(Any), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels=["input", "ephemeral"], + output_channels="output", + context_schema=Foo, + ) + + assert app.get_context_jsonschema() == { + "properties": { + "x": { + "title": "X", + "type": "integer", + }, + "y": { + "default": "foo", + "title": "Y", + "type": "string", + }, + }, + "required": [ + "x", + ], + "title": "Foo", + "type": "object", + } + + +def test_node_schemas_custom_output() -> None: + class State(TypedDict): + hello: str + bye: str + messages: Annotated[list[str], add_messages] + + class Output(TypedDict): + messages: list[str] + + class StateForA(TypedDict): + hello: str + messages: Annotated[list[str], add_messages] + + def node_a(state: StateForA) -> State: + assert state == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + class StateForB(TypedDict): + bye: str + now: int + + def node_b(state: StateForB): + assert state == { + "bye": "world", + } + return { + "now": 123, + "hello": "again", + } + + class StateForC(TypedDict): + hello: str + now: int + + def node_c(state: StateForC) -> StateForC: + assert state == { + "hello": "again", + "now": 123, + } + + builder = StateGraph(State, output_schema=Output) + builder.add_node("a", node_a) + builder.add_node("b", node_b) + builder.add_node("c", node_c) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", "c") + graph = builder.compile() + + assert graph.invoke({"hello": "there", "bye": "world", "messages": "hello"}) == { + "messages": [_AnyIdHumanMessage(content="hello")], + } + + builder = StateGraph(State, output_schema=Output) + builder.add_node("a", node_a) + builder.add_node("b", node_b) + builder.add_node("c", node_c) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", "c") + graph = builder.compile() + + assert graph.invoke( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) == { + "messages": [_AnyIdHumanMessage(content="hello")], + } + + assert [ + c + for c in graph.stream( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) + ] == [ + {"a": None}, + {"b": {"hello": "again", "now": 123}}, + {"c": None}, + ] + + +def test_reducer_before_first_node() -> None: + class State(TypedDict): + hello: str + messages: Annotated[list[str], add_messages] + + def node_a(state: State) -> State: + assert state == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.set_entry_point("a") + builder.set_finish_point("a") + graph = builder.compile() + assert graph.invoke({"hello": "there", "messages": "hello"}) == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + class State(TypedDict): + hello: str + messages: Annotated[list[str], add_messages] + + def node_a(state: State) -> State: + assert state == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.set_entry_point("a") + builder.set_finish_point("a") + graph = builder.compile() + assert graph.invoke({"hello": "there", "messages": "hello"}) == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + class State(TypedDict): + hello: str + messages: Annotated[Sequence[str], add_messages] + + def node_a(state: State) -> State: + assert state == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.set_entry_point("a") + builder.set_finish_point("a") + graph = builder.compile() + assert graph.invoke({"hello": "there", "messages": "hello"}) == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + +def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={ + "one": chain, + }, + channels={ + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "integer", + } + assert app.get_context_jsonschema() is None + + assert app.invoke(2) == 3 + assert app.invoke(2, output_keys=["output"]) == {"output": 3} + assert repr(app), "does not raise recursion error" + + +def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = ( + NodeBuilder() + .subscribe_only("input") + .do(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) + + app = Pregel( + nodes={"one": chain}, + channels={ + "input": LastValue(int), + "output": LastValue(int), + "fixed": LastValue(int), + "output_plus_one": LastValue(int), + }, + output_channels=["output", "fixed", "output_plus_one"], + input_channels="input", + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None}, + "fixed": {"title": "Fixed", "type": "integer", "default": None}, + "output_plus_one": { + "title": "Output Plus One", + "type": "integer", + "default": None, + }, + }, + } + assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} + + +def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels=["output"], + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } + assert app.invoke(2) == {"output": 3} + + +def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels=["input"], + output_channels=["output"], + ) + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer", "default": None}}, + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } + assert app.invoke({"input": 2}) == {"output": 3} + + +def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + assert app.invoke(2) == 4 + + with pytest.raises(GraphRecursionError): + app.invoke(2, {"recursion_limit": 1}, debug=1) + + +def test_run_from_checkpoint_id_retains_previous_writes( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + graph = builder.compile(checkpointer=sync_checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = graph.invoke({"myval": 1}, thread1, durability="async") + assert result["myval"] == 4 + history = [c for c in graph.get_state_history(thread1)] + + assert len(history) == 4 + assert history[-1].values == {"myval": 0} + assert history[0].values == {"myval": 4, "otherval": False} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = graph.invoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + for c in graph.get_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + +def test_batch_two_processes_in_out() -> None: + def add_one_with_delay(inp: int) -> int: + time.sleep(inp / 10) + return inp + 1 + + one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one") + two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] + assert app.batch([3, 2, 1, 3, 5], output_keys=["output"]) == [ + {"output": 5}, + {"output": 4}, + {"output": 3}, + {"output": 5}, + {"output": 7}, + ] + + +def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: + test_size = 100 + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} + for i in range(test_size - 2): + nodes[str(i)] = ( + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) + ) + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") + + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + for _ in range(10): + assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size + + with ThreadPoolExecutor() as executor: + assert [ + *executor.map(app.invoke, [2] * 10, [{"recursion_limit": test_size}] * 10) + ] == [2 + test_size] * 10 + + +def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: + test_size = 100 + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} + for i in range(test_size - 2): + nodes[str(i)] = ( + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) + ) + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") + + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + for _ in range(3): + assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ + 2 + test_size, + 1 + test_size, + 3 + test_size, + 4 + test_size, + 5 + test_size, + ] + + with ThreadPoolExecutor() as executor: + assert [ + *executor.map( + app.batch, [[2, 1, 3, 4, 5]] * 3, [{"recursion_limit": test_size}] * 3 + ) + ] == [ + [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size] + ] * 3 + + +def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={"output": LastValue(int), "input": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + with pytest.raises(InvalidUpdateError): + # LastValue channels can only be updated once per iteration + app.invoke(2) + + class State(TypedDict): + hello: str + + def my_node(input: State) -> State: + return {"hello": "world"} + + builder = StateGraph(State) + builder.add_node("one", my_node) + builder.add_node("two", my_node) + builder.set_conditional_entry_point(lambda _: ["one", "two"]) + + graph = builder.compile() + with pytest.raises(InvalidUpdateError, match="At key 'hello'"): + graph.invoke({"hello": "there"}, debug=True) + + +def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "output": Topic(int), + }, + input_channels="input", + output_channels="output", + ) + + # An Inbox channel accumulates updates into a sequence + assert app.invoke(2) == [3, 3] + + +def test_invoke_checkpoint_two( + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + errored_once = False + + def raise_if_above_10(input: int) -> int: + nonlocal errored_once + if input > 4: + if errored_once: + pass + else: + errored_once = True + raise ConnectionError("I will be retried") + if input > 10: + raise ValueError("Input is too large") + return input + + one = ( + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) + ) + + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=sync_checkpointer, + retry_policy=RetryPolicy(), + ) + + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 2 + # total is now 2, so output is 2+3=5 + assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" + checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) + assert checkpoint_tup is not None + assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated, error is recorded + checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) + assert checkpoint_tup is not None + assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 + assert checkpoint_tup.pending_writes == [ + (AnyStr(), ERROR, "ValueError('Input is too large')") + ] + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "2"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 5 + + +def test_pending_writes_resume( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + value: Annotated[int, operator.add] + + class AwhileMaker: + def __init__(self, sleep: float, rtn: dict | Exception) -> None: + self.sleep = sleep + self.rtn = rtn + self.reset() + + def __call__(self, input: State) -> Any: + self.calls += 1 + time.sleep(self.sleep) + if isinstance(self.rtn, Exception): + raise self.rtn + else: + return self.rtn + + def reset(self): + self.calls = 0 + + one = AwhileMaker(0.1, {"value": 2}) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) + builder = StateGraph(State) + builder.add_node("one", one) + builder.add_node( + "two", + two, + retry_policy=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False), + ) + builder.add_edge(START, "one") + builder.add_edge(START, "two") + graph = builder.compile(checkpointer=sync_checkpointer) + + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ConnectionError, match="I'm not good"): + graph.invoke({"value": 1}, thread1, durability=durability) + + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 2 # two attempts + + # latest checkpoint should be before nodes "one", "two" + # but we should have applied the write from "one" + state = graph.get_state(thread1) + assert state is not None + assert state.values == {"value": 3} + assert state.next == ("two",) + assert state.tasks == ( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), + PregelTask(AnyStr(), "two", (PULL, "two"), 'ConnectionError("I\'m not good")'), + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + } + # get_state with checkpoint_id should not apply any pending writes + state = graph.get_state(state.config) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + # should contain pending write of "one" + checkpoint = sync_checkpointer.get_tuple(thread1) + assert checkpoint is not None + # should contain error from "two" + expected_writes = [ + (AnyStr(), "value", 2), + (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), + ] + assert len(checkpoint.pending_writes) == 2 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] + + # resume execution + with pytest.raises(ConnectionError, match="I'm not good"): + graph.invoke(None, thread1, durability=durability) + + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 4 # two attempts before + two attempts now + + # confirm no new checkpoints saved + state_two = graph.get_state(thread1) + assert state_two.metadata == state.metadata + + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert graph.invoke(None, thread1, durability=durability) == {"value": 6} + + # check all final checkpoints + checkpoints = [c for c in sync_checkpointer.list(thread1)] + # we should have 3 + assert len(checkpoints) == (3 if durability != "exit" else 2) + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": { + "one": { + "branch:to:one": AnyVersion(), + }, + "two": { + "branch:to:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": {"value": 6}, + "updated_channels": ["value"], + }, + metadata={ + "parents": {}, + "step": 1, + "source": "loop", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "branch:to:one": None, + "branch:to:two": None, + }, + "updated_channels": ["branch:to:one", "branch:to:two", "value"], + }, + metadata={ + "parents": {}, + "step": 0, + "source": "loop", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": ( + checkpoints[2].config["configurable"]["checkpoint_id"] + ), + } + } + if durability != "exit" + else None, + pending_writes=( + UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + (AnyStr(), "value", 3), + ) + if durability != "exit" + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved + ) + ), + ) + if durability == "exit": + return + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, + "updated_channels": ["__start__"], + }, + metadata={ + "parents": {}, + "step": -1, + "source": "input", + }, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "branch:to:one", None), + (AnyStr(), "branch:to:two", None), + ), + ) + + +def test_cond_edge_after_send() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return [self.name] + + def send_for_fun(state): + return [Send("2", state), Send("2", state)] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"] + + +def test_concurrent_emit_sends() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + + def send_for_fun(state): + return [Send("2", 1), Send("2", 2), "3.1"] + + def send_for_profit(state): + return [Send("2", 3), Send("2", 4)] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("1.1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_edge(START, "1.1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("1.1", send_for_profit) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert graph.invoke(["0"]) == [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + + +def test_send_sequences() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Command): + return [state, Command(update=update)] + else: + return update + + def send_for_fun(state): + return [ + Send("2", Command(goto=Send("2", 3))), + Send("2", Command(goto=Send("2", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert graph.invoke(["0"]) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + + +def test_imp_task( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + mapper_calls = 0 + + class Context(TypedDict): + model: str + + @task() + def mapper(input: int) -> str: + nonlocal mapper_calls + mapper_calls += 1 + time.sleep(input / 100) + return str(input) * 2 + + @entrypoint(checkpointer=sync_checkpointer, context_schema=Context) + def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = [f.result() for f in futures] + answer = interrupt("question") + return [m + answer for m in mapped] + + assert graph.get_input_jsonschema() == { + "type": "array", + "items": {"type": "integer"}, + "title": "LangGraphInput", + } + assert graph.get_output_jsonschema() == { + "type": "array", + "items": {"type": "string"}, + "title": "LangGraphOutput", + } + assert graph.get_context_jsonschema() == { + "properties": {"model": {"title": "Model", "type": "string"}}, + "required": ["model"], + "title": "Context", + "type": "object", + } + + thread1 = {"configurable": {"thread_id": "1"}} + result = [*graph.stream([0, 1], thread1, durability=durability)] + # mapper tasks run concurrently so output order is non-deterministic + assert sorted(result[:-1], key=lambda d: str(d)) == [ + {"mapper": "00"}, + {"mapper": "11"}, + ] + assert result[-1] == { + "__interrupt__": ( + Interrupt( + value="question", + id=AnyStr(), + ), + ) + } + assert mapper_calls == 2 + + assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [ + "00answer", + "11answer", + ] + assert mapper_calls == 2 + + +def test_imp_nested( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + def mynode(input: list[str]) -> list[str]: + return [it + "a" for it in input] + + builder = StateGraph(list[str]) + builder.add_node(mynode) + builder.add_edge(START, "mynode") + add_a = builder.compile() + + @task + def submapper(input: int) -> str: + time.sleep(input / 100) + return str(input) + + @task() + def mapper(input: int) -> str: + sub = submapper(input) + time.sleep(input / 100) + return sub.result() * 2 + + @entrypoint(checkpointer=sync_checkpointer) + def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = [f.result() for f in futures] + answer = interrupt("question") + final = [m + answer for m in mapped] + return add_a.invoke(final) + + assert graph.get_input_jsonschema() == { + "type": "array", + "items": {"type": "integer"}, + "title": "LangGraphInput", + } + assert graph.get_output_jsonschema() == { + "type": "array", + "items": {"type": "string"}, + "title": "LangGraphOutput", + } + + thread1 = {"configurable": {"thread_id": "1"}} + assert [*graph.stream([0, 1], thread1, durability=durability)] == [ + {"submapper": "0"}, + {"mapper": "00"}, + {"submapper": "1"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + id=AnyStr(), + ), + ) + }, + ] + + assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [ + "00answera", + "11answera", + ] + + +def test_imp_stream_order( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + @task() + def foo(state: dict) -> tuple: + return state["a"] + "foo", "bar" + + @task + def bar(a: str, b: str, c: str | None = None) -> dict: + return {"a": a + b, "c": (c or "") + "bark"} + + @task + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @entrypoint(checkpointer=sync_checkpointer) + def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(*fut_foo.result()) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c for c in graph.stream({"a": "0"}, thread1, durability=durability)] == [ + { + "foo": ( + "0foo", + "bar", + ) + }, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"} + + +def test_invoke_checkpoint_three( + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver +) -> None: + adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + + def raise_if_above_10(input: int) -> int: + if input > 10: + raise ValueError("Input is too large") + return input + + one = ( + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(adder) + .write_to("output", "total") + .do(raise_if_above_10) + ) + + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=sync_checkpointer, + ) + + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, thread_1, durability="async") == 2 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert state.next == () + assert ( + state.config["configurable"]["checkpoint_id"] + == sync_checkpointer.get(thread_1)["id"] + ) + # total is now 2, so output is 2+3=5 + assert app.invoke(3, thread_1, durability="async") == 5 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == sync_checkpointer.get(thread_1)["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, thread_1, durability="async") + # checkpoint is updated with new input + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert app.invoke(2, thread_1, durability="async") == 9 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () + + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, thread_2) == 5 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == (), "checkpoint of other thread not touched" + state = app.get_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () + + assert len(list(app.get_state_history(thread_1, limit=1))) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c for c in app.get_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = list( + app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config) + ) + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert ( + sync_checkpointer.get(thread_1_history[0].config)["id"] + == thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + assert ( + sync_checkpointer.get(thread_1_history[1].config)["id"] + == thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + + thread_1_next_config = app.update_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # update makes new checkpoint child of the previous one + assert ( + app.get_state(thread_1_next_config).parent_config == thread_1_history[1].config + ) + # 1 more checkpoint in history + assert len(list(app.get_state_history(thread_1))) == 8 + assert Counter(c.metadata["source"] for c in app.get_state_history(thread_1)) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert app.get_state(thread_1) == app.get_state(thread_1_next_config) + + +def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_three = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_four = ( + NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output") + ) + + app = Pregel( + nodes={ + "one": one, + "chain_three": chain_three, + "chain_four": chain_four, + }, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + # Then invoke app + # We get a single array result as chain_four waits for all publishers to finish + # before operating on all elements published to topic_two as an array + for _ in range(100): + assert app.invoke(2) == [13, 13] + + with ThreadPoolExecutor() as executor: + assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100 + + +def test_invoke_join_then_call_other_pregel( + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) + + inner_app = Pregel( + nodes={ + "one": NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + }, + channels={ + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + one = NodeBuilder().subscribe_only("input").do(add_10_each).write_to("inbox_one") + two = ( + NodeBuilder() + .subscribe_only("inbox_one") + .do(inner_app.map()) + .write_to("outbox_one") + ) + chain_three = NodeBuilder().subscribe_only("outbox_one").do(sum).write_to("output") + + app = Pregel( + nodes={ + "one": one, + "two": two, + "chain_three": chain_three, + }, + channels={ + "inbox_one": Topic(int), + "outbox_one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + for _ in range(10): + assert app.invoke([2, 3]) == 27 + + with ThreadPoolExecutor() as executor: + assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10 + + # add checkpointer + app.checkpointer = sync_checkpointer + # subgraph is called twice in the same node, but that works + assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 + + # set inner graph checkpointer NeverCheckpoint + inner_app.checkpointer = False + # subgraph still called twice, but checkpointing for inner graph is disabled + assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 + + +def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = ( + NodeBuilder().subscribe_only("input").do(add_one).write_to("output", "between") + ) + two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + stream_channels=["output", "between"], + input_channels="input", + output_channels="output", + ) + + assert [c for c in app.stream(2, stream_mode="updates")] == [ + {"one": {"between": 3, "output": 3}}, + {"two": {"output": 4}}, + ] + assert [c for c in app.stream(2)] == [ + {"between": 3, "output": 3}, + {"between": 3, "output": 4}, + ] + + +def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between") + two = NodeBuilder().subscribe_only("between").do(add_one) + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + # It finishes executing (once no more messages being published) + # but returns nothing, as nothing was published to OUT topic + assert app.invoke(2) is None + + +def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("between").do(add_one) + + with pytest.raises(TypeError): + Pregel(nodes={"one": one, "two": two}) + + +def test_conditional_entrypoint_to_multiple_state_graph( + snapshot: SnapshotAssertion, +) -> None: + class OverallState(TypedDict): + locations: list[str] + results: Annotated[list[str], operator.add] + + def get_weather(state: OverallState) -> OverallState: + location = state["location"] + weather = "sunny" if len(location) > 2 else "cloudy" + return {"results": [f"It's {weather} in {location}"]} + + def continue_to_weather(state: OverallState) -> list[Send]: + return [ + Send("get_weather", {"location": location}) + for location in state["locations"] + ] + + workflow = StateGraph(OverallState) + + workflow.add_node("get_weather", get_weather) + workflow.add_edge("get_weather", END) + workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"]) + + app = workflow.compile() + + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"locations": ["sf", "nyc"]}, debug=True) == { + "locations": ["sf", "nyc"], + "results": ["It's cloudy in sf", "It's sunny in nyc"], + } + + assert [*app.stream({"locations": ["sf", "nyc"]}, stream_mode="values")][-1] == { + "locations": ["sf", "nyc"], + "results": ["It's cloudy in sf", "It's sunny in nyc"], + } + + +def test_conditional_state_graph_with_list_edge_inputs(snapshot: SnapshotAssertion): + class State(TypedDict): + foo: Annotated[list[str], operator.add] + + graph_builder = StateGraph(State) + graph_builder.add_node("A", lambda x: {"foo": ["A"]}) + graph_builder.add_node("B", lambda x: {"foo": ["B"]}) + graph_builder.add_edge(START, "A") + graph_builder.add_edge(START, "B") + graph_builder.add_edge(["A", "B"], END) + + app = graph_builder.compile() + assert app.invoke({"foo": []}) == {"foo": ["A", "B"]} + + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) -> None: + from langchain_core.language_models.fake import FakeStreamingListLLM + from langchain_core.prompts import PromptTemplate + from langchain_core.tools import tool + + class BaseState(TypedDict): + input: str + agent_outcome: AgentAction | AgentFinish | None + + class AgentState(BaseState, total=False): + intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] + + assert get_type_hints(AgentState).keys() == { + "input", + "agent_outcome", + "intermediate_steps", + } + + class Context(TypedDict, total=False): + tools: list[str] + + # Assemble the tools + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + # Construct the agent + prompt = PromptTemplate.from_template("Hello!") + + llm = FakeStreamingListLLM( + responses=[ + "tool:search_api:query", + "tool:search_api:another", + "finish:answer", + ] + ) + + def agent_parser(input: str) -> dict[str, AgentAction | AgentFinish]: + if input.startswith("finish"): + _, answer = input.split(":") + return { + "agent_outcome": AgentFinish( + return_values={"answer": answer}, log=input + ) + } + else: + _, tool_name, tool_input = input.split(":") + return { + "agent_outcome": AgentAction( + tool=tool_name, tool_input=tool_input, log=input + ) + } + + agent = prompt | llm | agent_parser + + # Define tool execution logic + def execute_tools(data: AgentState) -> dict: + agent_action: AgentAction = data.pop("agent_outcome") + observation = {t.name: t for t in tools}[agent_action.tool].invoke( + agent_action.tool_input + ) + return {"intermediate_steps": [(agent_action, observation)]} + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if isinstance(data["agent_outcome"], AgentFinish): + return "exit" + else: + return "continue" + + # Define a new graph + builder = StateGraph(AgentState, Context) + + builder.add_node("agent", agent) + builder.add_node("tools", execute_tools) + + builder.set_entry_point("agent") + + builder.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + + builder.add_edge("tools", "agent") + + app = builder.compile() + + assert json.dumps(app.get_context_jsonschema()) == snapshot + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + + assert builder.channels.keys() == {"input", "agent_outcome", "intermediate_steps"} + + assert app.invoke({"input": "what is weather in sf"}) == { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + "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", + ), + ], + } + + +def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None: + class AgentState(TypedDict, total=False): + input: str + output: str + steps: Annotated[list[str], operator.add] + + def left(data: AgentState) -> AgentState: + return {"output": data["input"] + "->left"} + + def right(data: AgentState) -> AgentState: + return {"output": data["input"] + "->right"} + + def should_start(data: AgentState) -> str: + assert data["steps"] == [], "Expected input to be read from the state" + # Logic to decide where to start + if len(data["input"]) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END, {END: END}) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert json.dumps(app.get_input_jsonschema()) == snapshot + assert json.dumps(app.get_output_jsonschema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + "steps": [], + } + + assert [*app.stream({"input": "what is weather in sf"})] == [ + {"right": {"output": "what is weather in sf->right"}}, + ] + + +def test_in_one_fan_out_state_graph_waiting_edge( + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + workflow = StateGraph(State) + + @workflow.add_node + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) # to ensure stream order + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow.add_node(analyzer_one) + workflow.add_node(retriever_one) + workflow.add_node(retriever_two) + workflow.add_node(qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["qa"], + ) + config = {"configurable": {"thread_id": "2"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state(config, {"docs": ["doc5"]}) + expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[ + -1 + ].config + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], + }, + tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),), + next=("qa",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + }, + parent_config=expected_parent_config, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, + ] + + +@pytest.mark.parametrize("use_waiting_edge", (True, False)) +def test_in_one_fan_out_state_graph_defer_node( + snapshot: SnapshotAssertion, + sync_checkpointer: BaseCheckpointSaver, + use_waiting_edge: bool, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + workflow = StateGraph(State) + + @workflow.add_node + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) # to ensure stream order + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow.add_node(analyzer_one) + workflow.add_node(retriever_one) + workflow.add_node(retriever_two) + workflow.add_node(qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("retriever_one", "analyzer_one") + workflow.add_edge("rewrite_query", "retriever_two") + if use_waiting_edge: + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + else: + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert [*app.stream({"query": "what is weather in sf"}, stream_mode="debug")] == [ + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "input": {"query": "what is weather in sf", "docs": []}, + "triggers": ("branch:to:rewrite_query",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "rewrite_query", + "error": None, + "result": { + "query": "query: what is weather in sf", + }, + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_one",), + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ("branch:to:retriever_two",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_one", + "error": None, + "result": { + "docs": ["doc1", "doc2"], + }, + "interrupts": [], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "error": None, + "result": { + "docs": ["doc3", "doc4"], + }, + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "input": { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:analyzer_one",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "analyzer_one", + "error": None, + "result": { + "query": "analyzed: query: what is weather in sf", + }, + "interrupts": [], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "input": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + "triggers": ("branch:to:qa", "join:retriever_one+retriever_two:qa") + if use_waiting_edge + else ("branch:to:qa",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 4, + "payload": { + "id": AnyStr(), + "name": "qa", + "error": None, + "result": { + "answer": "doc1,doc2,doc3,doc4", + }, + "interrupts": [], + }, + }, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["analyzer_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_before=["qa"], + ) + config = {"configurable": {"thread_id": "2"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"__interrupt__": ()}, + ] + + app_w_interrupt.update_state(config, {"docs": ["doc5"]}) + expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[ + -1 + ].config + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], + }, + tasks=(PregelTask(AnyStr(), "qa", (PULL, "qa")),), + next=("qa",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + }, + parent_config=expected_parent_config, + interrupts=(), + ) + + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, + ] + + +def test_in_one_fan_out_state_graph_waiting_edge_via_branch( + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + def rewrite_query_then(data: State) -> Literal["retriever_two"]: + return "retriever_two" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges("rewrite_query", rewrite_query_then) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + assert app.invoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( + snapshot: SnapshotAssertion, + sync_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class InnerObject(BaseModel): + yo: int + + class State(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + query: str + inner: Annotated[InnerObject, lambda x, y: y] + answer: str | None = None + docs: Annotated[list[str], sorted_add] + + class StateUpdate(BaseModel): + query: str | None = None + answer: str | None = None + docs: list[str] | None = None + + class UpdateDocs34(BaseModel): + docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"]) + + class Input(BaseModel): + query: str + inner: InnerObject + + class Output(BaseModel): + answer: str + docs: list[str] + + def rewrite_query(data: State) -> State: + assert isinstance(data.inner, InnerObject) + return {"query": f"query: {data.query}"} + + def analyzer_one(data: State) -> State: + assert isinstance(data.inner, InnerObject) + return StateUpdate(query=f"analyzed: {data.query}") + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return UpdateDocs34() + + def qa(data: State) -> State: + return {"answer": ",".join(data.docs)} + + def decider(data: State) -> str: + assert isinstance(data, State) + return "retriever_two" + + workflow = StateGraph(State, input_schema=Input, output_schema=Output) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", decider, {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if isinstance(sync_checkpointer, InMemorySaver): + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert app.get_input_jsonschema() == snapshot + assert app.get_output_jsonschema() == snapshot + + with pytest.raises(ValidationError): + app.invoke({"query": {}}) + + assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert app_w_interrupt.update_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } + } + + +def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class InnerObject(BaseModel): + yo: int + + class QueryModel(BaseModel): + query: str + + class State(QueryModel): + inner: InnerObject + answer: str | None = None + docs: Annotated[list[str], sorted_add] + + class StateUpdate(BaseModel): + query: str | None = None + answer: str | None = None + docs: list[str] | None = None + + class Input(QueryModel): + inner: InnerObject + + class Output(BaseModel): + answer: str + docs: list[str] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data.query}"} + + def analyzer_one(data: State) -> State: + return StateUpdate(query=f"analyzed: {data.query}") + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data.docs)} + + def decider(data: State) -> str: + assert isinstance(data, State) + return "retriever_two" + + workflow = StateGraph(State, input_schema=Input, output_schema=Output) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", decider, {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke( + Input(query="what is weather in sf", inner=InnerObject(yo=1)) + ) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [ + *app.stream(Input(query="what is weather in sf", inner=InnerObject(yo=1))) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + Input(query="what is weather in sf", inner=InnerObject(yo=1)), config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert app_w_interrupt.update_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } + } + + +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + time.sleep(0.1) + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.2) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + # silly edge, to make sure having been triggered before doesn't break + # semantics of named barrier (== waiting edges) + workflow.add_edge("rewrite_query", "qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] in ( + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ], + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ], + ) + + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] in ( + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +@pytest.mark.parametrize("with_cache", [True, False]) +def test_in_one_fan_out_state_graph_waiting_edge_multiple( + with_cache: bool, cache: BaseCache +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + rewrite_query_count = 0 + + def rewrite_query(data: State) -> State: + nonlocal rewrite_query_count + rewrite_query_count += 1 + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node( + "rewrite_query", + rewrite_query, + cache_policy=CachePolicy() if with_cache else None, + ) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile(cache=cache) + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + { + "rewrite_query": {"query": "query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"} + }, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + } + }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + ] + assert rewrite_query_count == 2 if with_cache else 4 + + # clear the cache + if with_cache: + app.clear_cache() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + assert rewrite_query_count == 4 + + +def test_callable_in_conditional_edges_with_no_path_map() -> None: + class State(TypedDict, total=False): + query: str + + def rewrite(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyze(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + class ChooseAnalyzer: + def __call__(self, data: State) -> str: + return "analyzer" + + workflow = StateGraph(State) + workflow.add_node("rewriter", rewrite) + workflow.add_node("analyzer", analyze) + workflow.add_conditional_edges("rewriter", ChooseAnalyzer()) + workflow.set_entry_point("rewriter") + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + } + + +def test_function_in_conditional_edges_with_no_path_map() -> None: + class State(TypedDict, total=False): + query: str + + def rewrite(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyze(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def choose_analyzer(data: State) -> str: + return "analyzer" + + workflow = StateGraph(State) + workflow.add_node("rewriter", rewrite) + workflow.add_node("analyzer", analyze) + workflow.add_conditional_edges("rewriter", choose_analyzer) + workflow.set_entry_point("rewriter") + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + } + + +def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def retriever_picker(data: State) -> list[str]: + return ["analyzer_one", "retriever_two"] + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_conditional_edges("rewrite_query", retriever_picker) + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + } + }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + ] + + +def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def up(state: State): + pass + + def side(state: State): + pass + + def other(state: State): + return {"my_key": "_more"} + + def down(state: State): + pass + + graph = StateGraph(State) + + graph.add_node("up", up) + graph.add_node("side", side) + graph.add_node("other", other) + graph.add_node("down", down) + + graph.set_entry_point("up") + graph.add_edge("up", "side") + graph.add_edge("up", "other") + graph.add_edge(["up", "side"], "down") + graph.set_finish_point("down") + + app = graph.compile() + + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value_more"} + assert [*app.stream({"my_key": "my_value"})] in ( + [ + {"up": None}, + {"side": None}, + {"other": {"my_key": "_more"}}, + {"down": None}, + ], + [ + {"up": None}, + {"other": {"my_key": "_more"}}, + {"side": None}, + {"down": None}, + ], + ) + + +def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def logic(state: State): + pass + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", logic) + tool_two_graph.add_node("tool_two_fast", logic) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + ["tool_two_slow", "tool_two_fast"], + ) + tool_two = tool_two_graph.compile() + + graph = StateGraph(State) + graph.add_node("tool_one", logic) + graph.add_node("tool_two", tool_two) + graph.add_node("tool_three", logic) + graph.set_conditional_entry_point( + lambda s: "tool_one", ["tool_one", "tool_two", "tool_three"] + ) + app = graph.compile() + + assert app.get_graph(xray=True).to_json() == snapshot + assert app.get_graph(xray=True).draw_mermaid() == snapshot + + +def test_nested_graph(snapshot: SnapshotAssertion) -> None: + def never_called_fn(state: Any): + assert 0, "This function should never be called" + + never_called = RunnableLambda(never_called_fn) + + class InnerState(TypedDict): + my_key: str + my_other_key: str + + def up(state: InnerState): + return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]} + + inner = StateGraph(InnerState) + inner.add_node("up", up) + inner.set_entry_point("up") + inner.set_finish_point("up") + + class State(TypedDict): + my_key: str + never_called: Any + + def side(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile()) + graph.add_node("side", side) + graph.set_entry_point("inner") + graph.add_edge("inner", "side") + graph.set_finish_point("side") + + app = graph.compile() + + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert app.get_graph(xray=True).draw_mermaid() == snapshot + assert app.invoke( + {"my_key": "my value", "never_called": never_called}, + print_mode=["values", "updates"], + ) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [*app.stream({"my_key": "my value", "never_called": never_called})] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + assert [ + *app.stream( + {"my_key": "my value", "never_called": never_called}, stream_mode="values" + ) + ] == [ + { + "my_key": "my value", + "never_called": never_called, + }, + { + "my_key": "my value there", + "never_called": never_called, + }, + { + "my_key": "my value there and back again", + "never_called": never_called, + }, + ] + + chain = app | RunnablePassthrough() + + assert chain.invoke({"my_key": "my value", "never_called": never_called}) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [*chain.stream({"my_key": "my value", "never_called": never_called})] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + + +def test_subgraph_checkpoint_true( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [ + c + for c in app.stream( + {"my_key": ""}, config, subgraphs=True, durability=durability + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + checkpoints = list(app.get_state_history(config)) + if durability != "exit": + assert len(checkpoints) == 4 + else: + assert len(checkpoints) == 1 + + +def test_subgraph_durability_inherited(durability: Durability) -> None: + sync_checkpointer = InMemorySaver() + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + inner_app = inner.compile(checkpointer=sync_checkpointer) + graph = StateGraph(State) + graph.add_node("inner", inner_app) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=sync_checkpointer) + thread_id = str(uuid.uuid4()) + config = {"configurable": {"thread_id": thread_id}} + app.invoke({"my_key": ""}, config, subgraphs=True, durability=durability) + if durability != "exit": + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 12 + else: + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 1 + + +def test_subgraph_checkpoint_true_interrupt( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + # Define subgraph + class SubgraphState(TypedDict): + # note that none of these keys are shared with the parent graph state + bar: str + baz: str + + def subgraph_node_1(state: SubgraphState): + baz_value = interrupt("Provide baz value") + return {"baz": baz_value} + + def subgraph_node_2(state: SubgraphState): + return {"bar": state["bar"] + state["baz"]} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + foo: str + + def node_1(state: ParentState): + return {"foo": "hi! " + state["foo"]} + + def node_2(state: ParentState): + response = subgraph.invoke({"bar": state["foo"]}) + return {"foo": response["bar"]} + + builder = StateGraph(ParentState) + builder.add_node("node_1", node_1) + builder.add_node("node_2", node_2) + builder.add_edge(START, "node_1") + builder.add_edge("node_1", "node_2") + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert graph.invoke({"foo": "foo"}, config, durability=durability) == { + "foo": "hi! foo", + "__interrupt__": [ + Interrupt( + value="Provide baz value", + id=AnyStr(), + ) + ], + } + assert graph.get_state(config, subgraphs=True).tasks[0].state.values == { + "bar": "hi! foo" + } + assert graph.invoke(Command(resume="baz"), config, durability=durability) == { + "foo": "hi! foobaz" + } + + +def test_stream_subgraphs_during_execution( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": "got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + time.sleep(0.5) + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def outer_1(state: State): + time.sleep(0.2) + return {"my_key": " and parallel"} + + def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile()) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.add_edge("outer_2", END) + + app = graph.compile(checkpointer=sync_checkpointer) + + start = time.perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + for c in app.stream({"my_key": ""}, config, subgraphs=True): + chunks.append((round(time.perf_counter() - start, 1), c)) + for idx in range(len(chunks)): + elapsed, c = chunks[idx] + chunks[idx] = (round(elapsed - chunks[0][0], 1), c) + + assert chunks == [ + # arrives before "inner" finishes + ( + FloatBetween(0.0, 0.1), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ), + (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})), + ( + FloatBetween(0.5, 0.8), + ( + (AnyStr("inner:"),), + {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, + ), + ), + (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), + (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), + ] + + +def test_stream_buffering_single_node(sync_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def node(state: State, writer: StreamWriter): + writer("Before sleep") + time.sleep(0.2) + writer("After sleep") + return {"my_key": "got here"} + + builder = StateGraph(State) + builder.add_node("node", node) + builder.add_edge(START, "node") + builder.add_edge("node", END) + graph = builder.compile(checkpointer=sync_checkpointer) + + start = time.perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + for c in graph.stream({"my_key": ""}, config, stream_mode="custom"): + chunks.append((round(time.perf_counter() - start, 1), c)) + + assert chunks == [ + (FloatBetween(0.0, 0.1), "Before sleep"), + (FloatBetween(0.2, 0.3), "After sleep"), + ] + + +def test_nested_graph_interrupts_parallel( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + time.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def outer_1(state: State): + return {"my_key": " and parallel"} + + def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=sync_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": ""}, config, durability=durability) == { + "my_key": " and parallel", + } + + assert app.invoke(None, config, durability=durability) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [ + *app.stream({"my_key": ""}, config, subgraphs=True, durability=durability) + ] == [ + # we got to parallel node first + ((), {"outer_1": {"my_key": " and parallel"}}), + ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), + ((), {"__interrupt__": ()}), + ] + assert [*app.stream(None, config, durability=durability)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=sync_checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [{"my_key": ""}] + # while we're waiting for the node w/ interrupt inside to finish + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=sync_checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + ] + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + +def test_doubly_nested_graph_interrupts( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=sync_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert app.invoke({"my_key": "my value"}, config, durability=durability) == { + "my_key": "hi my value", + } + + assert app.invoke(None, config, durability=durability) == { + "my_key": "hi my value here and there and back again", + } + + # test stream updates w/ nested interrupt + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } + assert [*app.stream({"my_key": "my value"}, config, durability=durability)] == [ + {"parent_1": {"my_key": "hi my value"}}, + {"__interrupt__": ()}, + ] + assert nodes == ["parent_1", "grandchild_1"] + assert [*app.stream(None, config, durability=durability)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] + assert nodes == [ + "parent_1", + "grandchild_1", + "grandchild_2", + "child_1", + "child", + "parent_2", + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + *app.stream( + {"my_key": "my value"}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": "my value"}, + {"my_key": "hi my value"}, + ] + assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [ + {"my_key": "hi my value"}, + {"my_key": "hi my value here and there"}, + {"my_key": "hi my value here and there and back again"}, + ] + + +def test_repeat_condition(snapshot: SnapshotAssertion) -> None: + class AgentState(TypedDict): + hello: str + + def router(state: AgentState) -> str: + return "hmm" + + workflow = StateGraph(AgentState) + workflow.add_node("Researcher", lambda x: x) + workflow.add_node("Chart Generator", lambda x: x) + workflow.add_node("Call Tool", lambda x: x) + workflow.add_conditional_edges( + "Researcher", + router, + { + "redo": "Researcher", + "continue": "Chart Generator", + "call_tool": "Call Tool", + "end": END, + }, + ) + workflow.add_conditional_edges( + "Chart Generator", + router, + {"continue": "Researcher", "call_tool": "Call Tool", "end": END}, + ) + workflow.add_conditional_edges( + "Call Tool", + # Each agent node updates the 'sender' field + # the tool calling node does not, meaning + # this edge will route back to the original agent + # who invoked the tool + lambda x: x["sender"], + { + "Researcher": "Researcher", + "Chart Generator": "Chart Generator", + }, + ) + workflow.set_entry_point("Researcher") + + app = workflow.compile() + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_checkpoint_metadata(sync_checkpointer: BaseCheckpointSaver) -> None: + """This test verifies that a run's configurable fields are merged with the + previous checkpoint config for each step in the run. + """ + # set up test + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, AnyMessage + from langchain_core.prompts import ChatPromptTemplate + from langchain_core.tools import tool + + # graph state + class BaseState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + # initialize graph nodes + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You are a nice assistant."), + ("placeholder", "{messages}"), + ] + ) + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + @traceable(run_type="llm") + def agent(state: BaseState) -> BaseState: + formatted = prompt.invoke(state) + response = model.invoke(formatted) + return {"messages": response, "usage_metadata": {"total_tokens": 123}} + + def should_continue(data: BaseState) -> str: + # Logic to decide whether to continue in the loop or exit + if not data["messages"][-1].tool_calls: + return "exit" + else: + return "continue" + + # define graphs w/ and w/o interrupt + workflow = StateGraph(BaseState) + workflow.add_node("agent", agent) + workflow.add_node("tools", ToolNode(tools)) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + workflow.add_edge("tools", "agent") + + # graph w/o interrupt + app = workflow.compile(checkpointer=sync_checkpointer) + + # graph w/ interrupt + app_w_interrupt = workflow.compile( + checkpointer=sync_checkpointer, interrupt_before=["tools"] + ) + + # assertions + + # invoke graph w/o interrupt + assert app.invoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "1", + "test_config_1": "foo", + "test_config_2": "bar", + }, + }, + ) == { + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="result for query", + name="search_api", + tool_call_id="tool_call123", + ), + _AnyIdAIMessage(content="answer"), + ] + } + + config = {"configurable": {"thread_id": "1"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_1 = sync_checkpointer.get_tuple(config).metadata + assert chkpnt_metadata_1["test_config_1"] == "foo" + assert chkpnt_metadata_1["test_config_2"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_1 = sync_checkpointer.list(config) + for chkpnt_tuple in chkpnt_tuples_1: + assert chkpnt_tuple.metadata["test_config_1"] == "foo" + assert chkpnt_tuple.metadata["test_config_2"] == "bar" + + # invoke graph, but interrupt before tool call + app_w_interrupt.invoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "2"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_2 = sync_checkpointer.get_tuple(config).metadata + assert chkpnt_metadata_2["test_config_3"] == "foo" + assert chkpnt_metadata_2["test_config_4"] == "bar" + + # resume graph execution + app_w_interrupt.invoke( + input=None, + config={ + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + } + }, + ) + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_3 = sync_checkpointer.get_tuple(config).metadata + assert chkpnt_metadata_3["test_config_3"] == "foo" + assert chkpnt_metadata_3["test_config_4"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_2 = sync_checkpointer.list(config) + for chkpnt_tuple in chkpnt_tuples_2: + assert chkpnt_tuple.metadata["test_config_3"] == "foo" + assert chkpnt_tuple.metadata["test_config_4"] == "bar" + + +def test_remove_message_via_state_update( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage + + workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] + workflow.add_node( + "chatbot", + lambda state: [ + AIMessage( + content="Hello! How can I help you", + ) + ], + ) + + workflow.set_entry_point("chatbot") + workflow.add_edge("chatbot", END) + + app = workflow.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + output = app.invoke([HumanMessage(content="Hi")], config=config) + app.update_state(config, values=[RemoveMessage(id=output[-1].id)]) + + updated_state = app.get_state(config) + + assert len(updated_state.values) == 1 + assert updated_state.values[-1].content == "Hi" + + app.checkpointer.delete_thread(config["configurable"]["thread_id"]) + + # Verify that the message was removed from the checkpointer + assert app.checkpointer.get_tuple(config) is None + assert [*app.get_state_history(config)] == [] + + +def test_remove_message_from_node(): + from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage + + workflow = StateGraph(state_schema=Annotated[list[AnyMessage], add_messages]) # type: ignore[arg-type] + workflow.add_node( + "chatbot", + lambda state: [ + AIMessage( + content="Hello!", + ), + AIMessage( + content="How can I help you?", + ), + ], + ) + workflow.add_node("delete_messages", lambda state: [RemoveMessage(id=state[-2].id)]) + workflow.set_entry_point("chatbot") + workflow.add_edge("chatbot", "delete_messages") + workflow.add_edge("delete_messages", END) + + app = workflow.compile() + output = app.invoke([HumanMessage(content="Hi")]) + assert len(output) == 2 + assert output[-1].content == "How can I help you?" + + +def test_xray_lance(snapshot: SnapshotAssertion): + from langchain_core.messages import AnyMessage, HumanMessage + + class Analyst(BaseModel): + affiliation: str = Field( + description="Primary affiliation of the investment analyst.", + ) + name: str = Field( + description="Name of the investment analyst.", + pattern=r"^[a-zA-Z0-9_-]{1,64}$", + ) + role: str = Field( + description="Role of the investment analyst in the context of the topic.", + ) + description: str = Field( + description="Description of the investment analyst focus, concerns, and motives.", + ) + + @property + def persona(self) -> str: + return f"Name: {self.name}\nRole: {self.role}\nAffiliation: {self.affiliation}\nDescription: {self.description}\n" + + class Perspectives(BaseModel): + analysts: list[Analyst] = Field( + description="Comprehensive list of investment analysts with their roles and affiliations.", + ) + + class Section(BaseModel): + section_title: str = Field(..., title="Title of the section") + context: str = Field( + ..., title="Provide a clear summary of the focus area that you researched." + ) + findings: str = Field( + ..., + title="Give a clear and detailed overview of your findings based upon the expert interview.", + ) + thesis: str = Field( + ..., + title="Give a clear and specific investment thesis based upon these findings.", + ) + + class InterviewState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + analyst: Analyst + section: Section + + class ResearchGraphState(TypedDict): + analysts: list[Analyst] + topic: str + max_analysts: int + sections: list[Section] + interviews: Annotated[list, operator.add] + + # Conditional edge + def route_messages(state): + return "ask_question" + + def generate_question(state): + return ... + + def generate_answer(state): + return ... + + # Add nodes and edges + interview_builder = StateGraph(InterviewState) + interview_builder.add_node("ask_question", generate_question) + interview_builder.add_node("answer_question", generate_answer) + + # Flow + interview_builder.add_edge(START, "ask_question") + interview_builder.add_edge("ask_question", "answer_question") + interview_builder.add_conditional_edges( + "answer_question", route_messages, ["ask_question", END] + ) + + # Interview + interview_graph = interview_builder.compile().with_config( + run_name="Conduct Interviews" + ) + + # View + assert interview_graph.get_graph().to_json() == snapshot + + def run_all_interviews(state: ResearchGraphState): + """Edge to run the interview sub-graph using Send""" + return [ + Send( + "conduct_interview", + { + "analyst": Analyst(), + "messages": [ + HumanMessage( + content="So you said you were writing an article on ...?" + ) + ], + }, + ) + for s in state["analysts"] + ] + + def generate_sections(state: ResearchGraphState): + return ... + + def generate_analysts(state: ResearchGraphState): + return ... + + builder = StateGraph(ResearchGraphState) + builder.add_node("generate_analysts", generate_analysts) + builder.add_node("conduct_interview", interview_builder.compile()) + builder.add_node("generate_sections", generate_sections) + + builder.add_edge(START, "generate_analysts") + builder.add_conditional_edges( + "generate_analysts", run_all_interviews, ["conduct_interview"] + ) + builder.add_edge("conduct_interview", "generate_sections") + builder.add_edge("generate_sections", END) + + graph = builder.compile() + + # View + assert graph.get_graph().to_json() == snapshot + assert graph.get_graph(xray=1).to_json() == snapshot + + +def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None: + config = {"configurable": {"thread_id": "1"}} + chain = NodeBuilder().subscribe_only("input").write_to("output") + app = Pregel( + nodes={ + "one": chain, + }, + channels={ + "ephemeral": EphemeralValue(Any), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels=["input", "ephemeral"], + output_channels="output", + checkpointer=sync_checkpointer, + ) + app.invoke({"input": 1, "ephemeral": "meow"}, config) + assert sync_checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1} + + +def test_xray_issue(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: Annotated[list, add_messages] + + def node(name): + def _node(state: State): + return {"messages": [("human", f"entered {name} node")]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + + child.add_edge("__start__", "c_one") + child.add_edge("c_two", "c_one") + + child.add_conditional_edges( + "c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"} + ) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + + parent.add_edge("__start__", "p_one") + parent.add_edge("p_two", "p_one") + + parent.add_conditional_edges( + "p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"} + ) + + app = parent.compile() + + assert app.get_graph(xray=True).draw_mermaid() == snapshot + + +def test_xray_bool(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: Annotated[list, add_messages] + + def node(name): + def _node(state: State): + return {"messages": [("human", f"entered {name} node")]} + + return _node + + grand_parent = StateGraph(State) + + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + + child.add_edge("__start__", "c_one") + child.add_edge("c_two", "c_one") + + child.add_conditional_edges( + "c_one", lambda x: str(randrange(0, 2)), {"0": "c_two", "1": "__end__"} + ) + + parent = StateGraph(State) + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge("__start__", "p_one") + parent.add_edge("p_two", "p_one") + parent.add_conditional_edges( + "p_one", lambda x: str(randrange(0, 2)), {"0": "p_two", "1": "__end__"} + ) + + grand_parent.add_node("gp_one", node("gp_one")) + grand_parent.add_node("gp_two", parent.compile()) + grand_parent.add_edge("__start__", "gp_one") + grand_parent.add_edge("gp_two", "gp_one") + grand_parent.add_conditional_edges( + "gp_one", lambda x: str(randrange(0, 2)), {"0": "gp_two", "1": "__end__"} + ) + + app = grand_parent.compile() + assert app.get_graph(xray=True).draw_mermaid() == snapshot + + +def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: Annotated[list, add_messages] + + subgraph_builder = StateGraph(State) + subgraph_builder.add_node("one", lambda x: x) + subgraph_builder.add_node("two", lambda x: x) + subgraph_builder.add_node("three", lambda x: x) + subgraph_builder.add_edge("__start__", "one") + subgraph_builder.add_conditional_edges("one", lambda x: "two", ["two", "three"]) + subgraph = subgraph_builder.compile() + + builder = StateGraph(State) + builder.add_node("uno", lambda x: x) + builder.add_node("dos", lambda x: x) + builder.add_node("subgraph", subgraph) + builder.add_edge("__start__", "uno") + builder.add_conditional_edges("uno", lambda x: "dos", ["dos", "subgraph"]) + + app = builder.compile() + assert app.get_graph(xray=True).draw_mermaid() == snapshot + + +def test_store_injected( + sync_checkpointer: BaseCheckpointSaver, sync_store: BaseStore +) -> None: + class State(TypedDict): + count: Annotated[int, operator.add] + + doc_id = str(uuid.uuid4()) + doc = {"some-key": "this-is-a-val"} + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") + thread_1 = str(uuid.uuid4()) + thread_2 = str(uuid.uuid4()) + + class Node: + def __init__(self, i: int | None = None): + self.i = i + + def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore): + assert isinstance(store, BaseStore) + store.put( + ( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar") + ), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} + + builder = StateGraph(State) + builder.add_node("node", Node()) + builder.add_edge("__start__", "node") + N = 50 + M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + + graph = builder.compile(store=sync_store, checkpointer=sync_checkpointer) + + results = graph.batch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], + ) + result = results[-1] + assert result == {"count": N + 1} + returned_doc = sync_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} + assert len(sync_store.search(namespace)) == 1 + # Check results after another turn of the same thread + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) + assert result == {"count": (N + 1) * 2} + returned_doc = sync_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len(sync_store.search(namespace)) == 1 + + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}}) + assert result == {"count": N + 1} + returned_doc = sync_store.get(namespace, doc_id).value + assert returned_doc == { + **doc, + "from_thread": thread_2, + "some_val": 0, + } # Overwrites the whole doc + assert len(sync_store.search(namespace)) == 1 # still overwriting the same one + + +def test_enum_node_names(): + class NodeName(str, enum.Enum): + BAZ = "baz" + + class State(TypedDict): + foo: str + bar: str + + def baz(state: State): + return {"bar": state["foo"] + "!"} + + graph = StateGraph(State) + graph.add_node(NodeName.BAZ, baz) + graph.add_edge(START, NodeName.BAZ) + graph.add_edge(NodeName.BAZ, END) + graph = graph.compile() + + assert graph.invoke({"foo": "hello"}) == {"foo": "hello", "bar": "hello!"} + + +def test_debug_retry(sync_checkpointer: BaseCheckpointSaver): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + builder = StateGraph(State) + builder.add_node("one", node("one")) + builder.add_node("two", node("two")) + builder.add_edge(START, "one") + builder.add_edge("one", "two") + builder.add_edge("two", END) + + graph = builder.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + graph.invoke({"messages": []}, config=config, durability="async") + + # re-run step: 1 + target_config = next( + c.parent_config + for c in sync_checkpointer.list(config) + if c.metadata["step"] == 1 + ) + update_config = graph.update_state(target_config, values=None) + + events = [ + *graph.stream( + None, config=update_config, stream_mode="debug", durability="async" + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + + checkpoint_history = { + c.config["configurable"]["checkpoint_id"]: c + for c in graph.get_state_history(config) + } + + def lax_normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + return config["configurable"] + + for stream in checkpoint_events: + stream_conf = lax_normalize_config(stream["config"]) + stream_parent_conf = lax_normalize_config(stream["parent_config"]) + assert stream_conf != stream_parent_conf + + # ensure the streamed checkpoint == checkpoint from checkpointer.list() + history = checkpoint_history[stream["config"]["configurable"]["checkpoint_id"]] + history_conf = lax_normalize_config(history.config) + assert stream_conf == history_conf + + history_parent_conf = lax_normalize_config(history.parent_config) + assert stream_parent_conf == history_parent_conf + + +def test_debug_subgraphs( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + graph = parent.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + events = [ + *graph.stream( + {"messages": []}, + config=config, + stream_mode="debug", + durability=durability, + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + if durability == "exit": + checkpoint_events = checkpoint_events[:1] + checkpoint_history = list(graph.get_state_history(config)) + + assert len(checkpoint_events) == len(checkpoint_history) + + def lax_normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + return config["configurable"] + + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert lax_normalize_config(stream["config"]) == lax_normalize_config( + history.config + ) + assert lax_normalize_config(stream["parent_config"]) == lax_normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +def test_debug_nested_subgraphs( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +): + from collections import defaultdict + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + grand_parent = StateGraph(State) + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + grand_parent.add_node("gp_one", node("gp_one")) + grand_parent.add_node("gp_two", parent.compile()) + grand_parent.add_edge(START, "gp_one") + grand_parent.add_edge("gp_one", "gp_two") + grand_parent.add_edge("gp_two", END) + + graph = grand_parent.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + events = [ + *graph.stream( + {"messages": []}, + config=config, + stream_mode="debug", + subgraphs=True, + durability=durability, + ) + ] + + stream_ns: dict[tuple, dict] = defaultdict(list) + for ns, e in events: + if e["type"] == "checkpoint": + stream_ns[ns].append(e["payload"]) + + assert list(stream_ns.keys()) == [ + (), + (AnyStr("gp_two:"),), + (AnyStr("gp_two:"), AnyStr("p_two:")), + ] + + history_ns = { + ns: list( + graph.get_state_history( + {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}} + ) + )[::-1] + for ns in stream_ns.keys() + } + + def normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + + clean_config = {} + clean_config["thread_id"] = config["configurable"]["thread_id"] + clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] + clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] + if "checkpoint_map" in config["configurable"]: + clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"] + + return clean_config + + for checkpoint_events, checkpoint_history, ns in zip( + stream_ns.values(), history_ns.values(), stream_ns.keys() + ): + if durability == "exit": + checkpoint_events = checkpoint_events[-1:] + if ns: # Save no checkpoints for subgraphs when durability="exit" + assert not checkpoint_history + continue + assert len(checkpoint_events) == len(checkpoint_history) + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config( + history.config + ) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +def test_add_sequence(): + class State(TypedDict): + foo: Annotated[list[str], operator.add] + bar: str + + def step1(state: State): + return {"foo": ["step1"], "bar": "baz"} + + def step2(state: State): + return {"foo": ["step2"]} + + # test raising if less than 1 steps + with pytest.raises(ValueError): + StateGraph(State).add_sequence([]) + + # test raising if duplicate step names + with pytest.raises(ValueError): + StateGraph(State).add_sequence([step1, step1]) + + with pytest.raises(ValueError): + StateGraph(State).add_sequence([("foo", step1), ("foo", step1)]) + + # test unnamed steps + builder = StateGraph(State) + builder.add_sequence([step1, step2]) + builder.add_edge(START, "step1") + graph = builder.compile() + result = graph.invoke({"foo": []}) + assert result == {"foo": ["step1", "step2"], "bar": "baz"} + stream_chunks = list(graph.stream({"foo": []})) + assert stream_chunks == [ + {"step1": {"foo": ["step1"], "bar": "baz"}}, + {"step2": {"foo": ["step2"]}}, + ] + + # test named steps + builder_named_steps = StateGraph(State) + builder_named_steps.add_sequence([("meow1", step1), ("meow2", step2)]) + builder_named_steps.add_edge(START, "meow1") + graph_named_steps = builder_named_steps.compile() + result = graph_named_steps.invoke({"foo": []}) + stream_chunks = list(graph_named_steps.stream({"foo": []})) + assert result == {"foo": ["step1", "step2"], "bar": "baz"} + assert stream_chunks == [ + {"meow1": {"foo": ["step1"], "bar": "baz"}}, + {"meow2": {"foo": ["step2"]}}, + ] + + builder_named_steps = StateGraph(State) + builder_named_steps.add_sequence( + [ + ("meow1", lambda state: {"foo": ["foo"]}), + ("meow2", lambda state: {"bar": state["foo"][0] + "bar"}), + ], + ) + builder_named_steps.add_edge(START, "meow1") + graph_named_steps = builder_named_steps.compile() + result = graph_named_steps.invoke({"foo": []}) + stream_chunks = list(graph_named_steps.stream({"foo": []})) + # filtered by output schema + assert result == {"bar": "foobar", "foo": ["foo"]} + assert stream_chunks == [ + {"meow1": {"foo": ["foo"]}}, + {"meow2": {"bar": "foobar"}}, + ] + + # test two sequences + + def a(state: State): + return {"foo": ["a"]} + + def b(state: State): + return {"foo": ["b"]} + + builder_two_sequences = StateGraph(State) + builder_two_sequences.add_sequence([a]) + builder_two_sequences.add_sequence([b]) + builder_two_sequences.add_edge(START, "a") + builder_two_sequences.add_edge("a", "b") + graph_two_sequences = builder_two_sequences.compile() + + result = graph_two_sequences.invoke({"foo": []}) + assert result == {"foo": ["a", "b"]} + + stream_chunks = list(graph_two_sequences.stream({"foo": []})) + assert stream_chunks == [ + {"a": {"foo": ["a"]}}, + {"b": {"foo": ["b"]}}, + ] + + # test mixed nodes and sequences + + def c(state: State): + return {"foo": ["c"]} + + def d(state: State): + return {"foo": ["d"]} + + def e(state: State): + return {"foo": ["e"]} + + def foo(state: State): + if state["foo"][0] == "a": + return "d" + else: + return "c" + + builder_complex = StateGraph(State) + builder_complex.add_sequence([a, b]) + builder_complex.add_conditional_edges("b", foo) + builder_complex.add_node(c) + builder_complex.add_sequence([d, e]) + builder_complex.add_edge(START, "a") + graph_complex = builder_complex.compile() + + result = graph_complex.invoke({"foo": []}) + assert result == {"foo": ["a", "b", "d", "e"]} + + result = graph_complex.invoke({"foo": ["start"]}) + assert result == {"foo": ["start", "a", "b", "c"]} + + stream_chunks = list(graph_complex.stream({"foo": []})) + assert stream_chunks == [ + {"a": {"foo": ["a"]}}, + {"b": {"foo": ["b"]}}, + {"d": {"foo": ["d"]}}, + {"e": {"foo": ["e"]}}, + ] + + +def test_runnable_passthrough_node_graph() -> None: + class State(TypedDict): + changeme: str + + async def dummy(state): + return state + + agent = dummy | RunnablePassthrough.assign(prediction=RunnableLambda(lambda x: x)) + + graph_builder = StateGraph(State) + + graph_builder.add_node("agent", agent) + graph_builder.add_edge(START, "agent") + + graph = graph_builder.compile() + + assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() + + +@pytest.mark.parametrize("subgraph_persist", [True, False]) +def test_parent_command( + sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: + from langchain_core.messages import BaseMessage + from langchain_core.tools import tool + + @tool(return_direct=True) + def get_user_name() -> Command: + """Retrieve user name""" + return Command(update={"user_name": "Meow"}, graph=Command.PARENT) + + subgraph_builder = StateGraph(MessagesState) + subgraph_builder.add_node("tool", get_user_name) + subgraph_builder.add_edge(START, "tool") + subgraph = subgraph_builder.compile(checkpointer=subgraph_persist) + + class CustomParentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + # this key is not available to the child graph + user_name: str + + builder = StateGraph(CustomParentState) + builder.add_node("alice", subgraph) + builder.add_edge(START, "alice") + graph = builder.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert graph.invoke( + {"messages": [("user", "get user name")]}, config, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage( + content="get user name", additional_kwargs={}, response_metadata={} + ), + ], + "user_name": "Meow", + } + assert graph.get_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage( + content="get user name", additional_kwargs={}, response_metadata={} + ), + ], + "user_name": "Meow", + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(), + interrupts=(), + ) + + +def test_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver): + class State(TypedDict): + baz: str + + def foo(state): + return {"baz": "foo"} + + def bar(state): + value = interrupt("Please provide baz value:") + return {"baz": value} + + child_builder = StateGraph(State) + child_builder.add_node(bar) + child_builder.add_edge(START, "bar") + + builder = StateGraph(State) + builder.add_node(foo) + builder.add_node("bar", child_builder.compile()) + builder.add_edge(START, "foo") + builder.add_edge("foo", "bar") + graph = builder.compile(checkpointer=sync_checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert graph.invoke({"baz": ""}, thread1) + # Resume with answer + assert graph.invoke(Command(resume="bar"), thread1) + + +@pytest.mark.parametrize("resume_style", ["null", "map"]) +def test_interrupt_multiple( + sync_checkpointer: BaseCheckpointSaver, resume_style: Literal["null", "map"] +): + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def node(s: State) -> State: + answer = interrupt({"value": 1}) + answer2 = interrupt({"value": 2}) + return {"my_key": answer + " " + answer2} + + builder = StateGraph(State) + builder.add_node("node", node) + builder.add_edge(START, "node") + + graph = builder.compile(checkpointer=sync_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + + result = [e for e in graph.stream({"my_key": "DE", "market": "DE"}, thread1)] + assert result == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 1}, + id=AnyStr(), + ), + ) + } + ] + + result = [ + event + for event in graph.stream( + Command( + resume="answer 1" + if resume_style == "null" + else {result[0]["__interrupt__"][0].id: "answer 1"}, + update={"my_key": " foofoo "}, + ), + thread1, + ) + ] + assert result == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 2}, + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + for event in graph.stream( + Command( + resume="answer 2" + if resume_style == "null" + else {result[0]["__interrupt__"][0].id: "answer 2"} + ), + thread1, + stream_mode="values", + ) + ] == [ + {"my_key": "DE foofoo "}, + {"my_key": "DE foofoo answer 1 answer 2"}, + ] + + +def test_interrupt_loop(sync_checkpointer: BaseCheckpointSaver): + class State(TypedDict): + age: int + other: str + + def ask_age(s: State): + """Ask an expert for help.""" + question = "How old are you?" + value = None + for _ in range(10): + value: str = interrupt(question) + if not value.isdigit() or int(value) < 18: + question = "invalid response" + value = None + else: + break + + return {"age": int(value)} + + builder = StateGraph(State) + builder.add_node("node", ask_age) + builder.add_edge(START, "node") + + graph = builder.compile(checkpointer=sync_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + + assert [e for e in graph.stream({"other": ""}, thread1)] == [ + { + "__interrupt__": ( + Interrupt( + value="How old are you?", + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + for event in graph.stream( + Command(resume="13"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + for event in graph.stream( + Command(resume="15"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + id=AnyStr(), + ), + ) + } + ] + + assert [event for event in graph.stream(Command(resume="19"), thread1)] == [ + {"node": {"age": 19}}, + ] + + +def test_interrupt_functional( + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion +) -> None: + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + @entrypoint(checkpointer=sync_checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**fut_foo.result(), "b": value} + fut_bar = bar(bar_input) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert graph.invoke({"a": ""}, config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + id=AnyStr(), + ) + ] + } + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +def test_interrupt_task_functional( + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion +) -> None: + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + @entrypoint(checkpointer=sync_checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + fut_bar = bar(fut_foo.result()) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert graph.invoke({"a": ""}, config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + id=AnyStr(), + ), + ] + } + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + # Test that we can interrupt the same task multiple times + config = {"configurable": {"thread_id": "2"}} + + @entrypoint(checkpointer=sync_checkpointer) + def graph(inputs: dict) -> dict: + foo_result = foo(inputs).result() + bar_result = bar(foo_result).result() + baz_result = bar(bar_result).result() + return baz_result + + # First run, interrupted at bar + assert graph.invoke({"a": ""}, config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + id=AnyStr(), + ), + ] + } + # Provide resumes + graph.invoke(Command(resume="bar"), config) + assert graph.invoke(Command(resume="baz"), config) == {"a": "foobarbaz"} + + +def test_root_mixed_return() -> None: + def my_node(state: list[str]): + return [Command(update=["a"]), ["b"]] + + graph = StateGraph(Annotated[list[str], operator.add]) + + graph.add_node(my_node) + graph.add_edge(START, "my_node") + graph = graph.compile() + + assert graph.invoke([]) == ["a", "b"] + + +def test_dict_mixed_return() -> None: + class State(TypedDict): + foo: Annotated[str, operator.add] + + def my_node(state: State): + return [Command(update={"foo": "a"}), {"foo": "b"}] + + graph = StateGraph(State) + graph.add_node(my_node) + graph.add_edge(START, "my_node") + graph = graph.compile() + + assert graph.invoke({"foo": ""}) == {"foo": "ab"} + + +def test_command_pydantic_dataclass() -> None: + class PydanticState(BaseModel): + foo: str + + @dataclass + class DataclassState: + foo: str + + for State in (PydanticState, DataclassState): + + def node_a(state) -> Command[Literal["node_b"]]: + return Command( + update=State(foo="foo"), + goto="node_b", + ) + + def node_b(state): + return {"foo": state.foo + "bar"} + + builder = StateGraph(State) + builder.add_edge(START, "node_a") + builder.add_node(node_a) + builder.add_node(node_b) + graph = builder.compile() + assert graph.invoke(State(foo="")) == {"foo": "foobar"} + + +def test_command_with_static_breakpoints( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that we can use Command to resume and update with static breakpoints.""" + + class State(TypedDict): + """The graph state.""" + + foo: str + + def node1(state: State): + return { + "foo": state["foo"] + "|node-1", + } + + def node2(state: State): + return { + "foo": state["foo"] + "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"]) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + graph.invoke({"foo": "abc"}, config) + result = graph.invoke(Command(resume="node1"), config) + assert result == {"foo": "abc|node-1|node-2"} + + +def test_multistep_plan(sync_checkpointer: BaseCheckpointSaver): + from langchain_core.messages import AnyMessage + + class State(TypedDict, total=False): + plan: list[str | list[str]] + messages: Annotated[list[AnyMessage], add_messages] + + def planner(state: State): + if state.get("plan") is None: + # create plan somehow + plan = ["step1", ["step2", "step3"], "step4"] + # pick the first step to execute next + first_step, *plan = plan + # put the rest of plan in state + return Command(goto=first_step, update={"plan": plan}) + elif state["plan"]: + # go to the next step of the plan + next_step, *next_plan = state["plan"] + return Command(goto=next_step, update={"plan": next_plan}) + else: + # the end of the plan + pass + + def step1(state: State): + return Command(goto="planner", update={"messages": [("human", "step1")]}) + + def step2(state: State): + return Command(goto="planner", update={"messages": [("human", "step2")]}) + + def step3(state: State): + return Command(goto="planner", update={"messages": [("human", "step3")]}) + + def step4(state: State): + return Command(goto="planner", update={"messages": [("human", "step4")]}) + + builder = StateGraph(State) + builder.add_node(planner) + builder.add_node(step1) + builder.add_node(step2) + builder.add_node(step3) + builder.add_node(step4) + builder.add_edge(START, "planner") + graph = builder.compile(checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert graph.invoke({"messages": [("human", "start")]}, config) == { + "messages": [ + _AnyIdHumanMessage(content="start"), + _AnyIdHumanMessage(content="step1"), + _AnyIdHumanMessage(content="step2"), + _AnyIdHumanMessage(content="step3"), + _AnyIdHumanMessage(content="step4"), + ], + "plan": [], + } + + +def test_command_goto_with_static_breakpoints( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Use Command goto with static breakpoints.""" + + class State(TypedDict): + """The graph state.""" + + foo: Annotated[str, operator.add] + + def node1(state: State): + return { + "foo": "|node-1", + } + + def node2(state: State): + return { + "foo": "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"]) + + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + graph.invoke({"foo": "abc"}, config) + result = graph.invoke(Command(goto=["node2"]), config) + assert result == {"foo": "abc|node-1|node-2|node-2"} + + +def test_parallel_node_execution(): + """Test that parallel nodes execute concurrently.""" + + class State(TypedDict): + results: Annotated[list[str], operator.add] + + def slow_node(state: State): + time.sleep(1) + return {"results": ["slow"]} + + def fast_node(state: State): + time.sleep(2) + return {"results": ["fast"]} + + builder = StateGraph(State) + builder.add_node("slow", slow_node) + builder.add_node("fast", fast_node) + builder.add_edge(START, "slow") + builder.add_edge(START, "fast") + + graph = builder.compile() + + start = time.perf_counter() + result = graph.invoke({"results": []}) + duration = time.perf_counter() - start + + # Fast node result should be available first + assert "fast" in result["results"][0] + + # Total duration should be less than sum of both nodes + assert duration < 3.0 + + +def test_multiple_interrupt_state_persistence( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that state is preserved correctly across multiple interrupts.""" + + class State(TypedDict): + steps: Annotated[list[str], operator.add] + + def interruptible_node(state: State): + first = interrupt("First interrupt") + second = interrupt("Second interrupt") + return {"steps": [first, second]} + + builder = StateGraph(State) + builder.add_node("node", interruptible_node) + builder.add_edge(START, "node") + + app = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First execution - should hit first interrupt + app.invoke({"steps": []}, config) + + # State should still be empty since node hasn't returned + state = app.get_state(config) + assert state.values == {"steps": []} + + # Resume after first interrupt - should hit second interrupt + app.invoke(Command(resume="step1"), config) + + # State should still be empty since node hasn't returned + state = app.get_state(config) + assert state.values == {"steps": []} + + # Resume after second interrupt - node should complete + result = app.invoke(Command(resume="step2"), config) + + # Now state should contain both steps since node returned + assert result["steps"] == ["step1", "step2"] + state = app.get_state(config) + assert state.values["steps"] == ["step1", "step2"] + + +def test_concurrent_execution_thread_safety(): + """Test thread safety during concurrent execution.""" + + class State(TypedDict): + counter: Annotated[int, operator.add] + + results = deque() # thread-safe queue + threads: list[threading.Thread] = [] + + def slow_node(state: State): + time.sleep(0.1) + return {"counter": 1} + + builder = StateGraph(State) + builder.add_node("node", slow_node) + builder.add_edge(START, "node") + graph = builder.compile() + + def run_graph(): + result = graph.invoke({"counter": 0}) + results.append(result) + + # Start multiple threads + for _ in range(10): + thread = threading.Thread(target=run_graph) + thread.start() + threads.append(thread) + + # Wait for all threads + for thread in threads: + thread.join() + + # Verify results are independent + assert len(results) == 10 + for result in results: + assert result["counter"] == 1 + + +def test_checkpoint_recovery( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +): + """Test recovery from checkpoints after failures.""" + + class State(TypedDict): + steps: Annotated[list[str], operator.add] + attempt: int # Track number of attempts + + def failing_node(state: State): + # Fail on first attempt, succeed on retry + if state["attempt"] == 1: + raise RuntimeError("Simulated failure") + return {"steps": ["node1"]} + + def second_node(state: State): + return {"steps": ["node2"]} + + builder = StateGraph(State) + builder.add_node("node1", failing_node) + builder.add_node("node2", second_node) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First attempt should fail + with pytest.raises(RuntimeError): + graph.invoke( + {"steps": ["start"], "attempt": 1}, + config, + durability=durability, + ) + + # Verify checkpoint state + state = graph.get_state(config) + assert state is not None + assert state.values == {"steps": ["start"], "attempt": 1} # input state saved + assert state.next == ("node1",) # Should retry failed node + assert "RuntimeError('Simulated failure')" in state.tasks[0].error + + # Retry with updated attempt count + result = graph.invoke({"steps": [], "attempt": 2}, config, durability=durability) + assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} + + # Verify checkpoint history shows both attempts + history = list(graph.get_state_history(config)) + if durability != "exit": + assert len(history) == 6 # Initial + failed attempt + successful attempt + else: + assert len(history) == 2 # error + success + + # Verify the error was recorded in checkpoint + failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) + assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error + + # Verify delete leaves it empty + graph.checkpointer.delete_thread(config["configurable"]["thread_id"]) + assert graph.checkpointer.get_tuple(config) is None + assert [*graph.get_state_history(config)] == [] + + +def test_multiple_updates_root() -> None: + def node_a(state): + return [Command(update="a1"), Command(update="a2")] + + def node_b(state): + return "b" + + graph = ( + StateGraph(Annotated[str, operator.add]) + .add_sequence([node_a, node_b]) + .add_edge(START, "node_a") + .compile() + ) + + assert graph.invoke("") == "a1a2b" + + # only streams the last update from node_a + assert [c for c in graph.stream("", stream_mode="updates")] == [ + {"node_a": ["a1", "a2"]}, + {"node_b": "b"}, + ] + + +def test_multiple_updates() -> None: + class State(TypedDict): + foo: Annotated[str, operator.add] + + def node_a(state): + return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})] + + def node_b(state): + return {"foo": "b"} + + graph = ( + StateGraph(State) + .add_sequence([node_a, node_b]) + .add_edge(START, "node_a") + .compile() + ) + + assert graph.invoke({"foo": ""}) == { + "foo": "a1a2b", + } + + # only streams the last update from node_a + assert [c for c in graph.stream({"foo": ""}, stream_mode="updates")] == [ + {"node_a": [{"foo": "a1"}, {"foo": "a2"}]}, + {"node_b": {"foo": "b"}}, + ] + + +def test_falsy_return_from_task(sync_checkpointer: BaseCheckpointSaver): + """Test with a falsy return from a task.""" + + @task + def falsy_task() -> bool: + return False + + @entrypoint(checkpointer=sync_checkpointer) + def graph(state: dict) -> dict: + """React tool.""" + falsy_task().result() + interrupt("test") + + configurable = {"configurable": {"thread_id": uuid.uuid4()}} + assert [ + chunk + for chunk in graph.stream( + {"a": 5}, configurable, stream_mode="debug", durability="exit" + ) + ] == [ + { + "payload": { + "config": { + "configurable": { + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + "thread_id": AnyStr(), + }, + }, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + }, + "next": [ + "graph", + ], + "parent_config": None, + "tasks": [ + { + "id": AnyStr(), + "interrupts": (), + "name": "graph", + "state": None, + }, + ], + "values": None, + }, + "step": -1, + "timestamp": AnyStr(), + "type": "checkpoint", + }, + { + "payload": { + "id": AnyStr(), + "input": { + "a": 5, + }, + "name": "graph", + "triggers": ("__start__",), + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task", + }, + { + "payload": { + "id": AnyStr(), + "input": ( + (), + {}, + ), + "name": "falsy_task", + "triggers": ("__pregel_push",), + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task", + }, + { + "payload": { + "error": None, + "id": AnyStr(), + "interrupts": [], + "name": "falsy_task", + "result": { + "__return__": False, + }, + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task_result", + }, + { + "payload": { + "error": None, + "id": AnyStr(), + "interrupts": [ + { + "id": AnyStr(), + "value": "test", + }, + ], + "name": "graph", + "result": {}, + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task_result", + }, + ] + assert [ + c + for c in graph.stream( + Command(resume="123"), + configurable, + stream_mode="debug", + durability="exit", + ) + ] == [ + { + "payload": { + "config": { + "configurable": { + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + "thread_id": AnyStr(), + }, + }, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + }, + "next": [ + "graph", + ], + "parent_config": None, + "tasks": [ + { + "id": AnyStr(), + "interrupts": ( + { + "id": AnyStr(), + "value": "test", + }, + ), + "name": "graph", + "state": None, + }, + ], + "values": None, + }, + "step": -1, + "timestamp": AnyStr(), + "type": "checkpoint", + }, + { + "payload": { + "id": AnyStr(), + "input": { + "a": 5, + }, + "name": "graph", + "triggers": ("__start__",), + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task", + }, + { + "payload": { + "id": AnyStr(), + "input": ( + (), + {}, + ), + "name": "falsy_task", + "triggers": ("__pregel_push",), + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task", + }, + { + "payload": { + "error": None, + "id": AnyStr(), + "interrupts": [], + "name": "graph", + "result": { + "__end__": None, + }, + }, + "step": 0, + "timestamp": AnyStr(), + "type": "task_result", + }, + { + "payload": { + "config": { + "configurable": { + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + "thread_id": AnyStr(), + }, + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + }, + "next": [], + "parent_config": None, + "tasks": [], + "values": None, + }, + "step": 0, + "timestamp": AnyStr(), + "type": "checkpoint", + }, + ] + + +def test_multiple_interrupts_functional(sync_checkpointer: BaseCheckpointSaver): + """Test multiple interrupts with functional API.""" + + counter = 0 + + @task + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=sync_checkpointer) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 2, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + result = graph.invoke(Command(resume="c"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 4, "b", 6, "c"], + } + assert counter == 3 + + +def test_multiple_interrupts_functional_cache( + sync_checkpointer: BaseCheckpointSaver, cache: BaseCache +): + """Test multiple interrupts with functional API.""" + + counter = 0 + + @task(cache_policy=CachePolicy()) + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=sync_checkpointer, cache=cache) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 1, 2, 2, 3, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + # should all be cached now + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + # clear cache + double.clear_cache(cache) + + # should recompute now + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + graph.invoke(Command(resume="c"), configurable) + graph.invoke(Command(resume="d"), configurable) + graph.invoke(Command(resume="e"), configurable) + result = graph.invoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 6 + + +def test_double_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None: + class AgentState(TypedDict): + input: str + + def node_1(state: AgentState): + result = interrupt("interrupt node 1") + return {"input": result} + + def node_2(state: AgentState): + result = interrupt("interrupt node 2") + return {"input": result} + + subgraph_builder = ( + StateGraph(AgentState) + .add_node("node_1", node_1) + .add_node("node_2", node_2) + .add_edge(START, "node_1") + .add_edge("node_1", "node_2") + .add_edge("node_2", END) + ) + + # invoke the sub graph + subgraph = subgraph_builder.compile(checkpointer=sync_checkpointer) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + assert [c for c in subgraph.stream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + id=AnyStr(), + ), + ) + }, + ] + # resume from the first interrupt + assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [ + { + "node_1": {"input": "123"}, + }, + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + id=AnyStr(), + ), + ) + }, + ] + # resume from the second interrupt + assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [ + { + "node_2": {"input": "123"}, + }, + ] + + subgraph = subgraph_builder.compile() + + def invoke_sub_agent(state: AgentState): + return subgraph.invoke(state) + + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + parent_agent = ( + StateGraph(AgentState) + .add_node("invoke_sub_agent", invoke_sub_agent) + .add_edge(START, "invoke_sub_agent") + .add_edge("invoke_sub_agent", END) + .compile(checkpointer=sync_checkpointer) + ) + + assert [c for c in parent_agent.stream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + id=AnyStr(), + ), + ) + }, + ] + + # resume from the first interrupt + assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + id=AnyStr(), + ), + ) + } + ] + + # resume from 2nd interrupt + assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [ + { + "invoke_sub_agent": {"input": True}, + }, + ] + + +def test_multi_resume(sync_checkpointer: BaseCheckpointSaver) -> None: + class ChildState(TypedDict): + prompt: str + human_input: str + human_inputs: list[str] + + def get_human_input(state: ChildState): + human_input = interrupt(state["prompt"]) + + return { + "human_input": human_input, + "human_inputs": [human_input], + } + + child_graph = ( + StateGraph(ChildState) + .add_node("get_human_input", get_human_input) + .add_edge(START, "get_human_input") + .add_edge("get_human_input", END) + .compile(checkpointer=sync_checkpointer) + ) + + class ParentState(TypedDict): + prompts: list[str] + human_inputs: Annotated[list[str], operator.add] + + def assign_workers(state: ParentState) -> list[Send]: + return [ + Send( + "child_graph", + {"prompt": prompt}, + ) + for prompt in state["prompts"] + ] + + def cleanup(state: ParentState): + assert len(state["human_inputs"]) == len(state["prompts"]) + + parent_graph = ( + StateGraph(ParentState) + .add_node("child_graph", child_graph) + .add_node("cleanup", cleanup) + .add_conditional_edges(START, assign_workers, ["child_graph"]) + .add_edge("child_graph", "cleanup") + .add_edge("cleanup", END) + .compile(checkpointer=sync_checkpointer) + ) + + thread_config: RunnableConfig = { + "configurable": { + "thread_id": uuid.uuid4(), + }, + } + + prompts = ["a", "b", "c", "d", "e"] + + events = parent_graph.invoke( + {"prompts": prompts}, thread_config, stream_mode="values" + ) + + assert len(events["__interrupt__"]) == len(prompts) + interrupt_values = {i.value for i in events["__interrupt__"]} + assert interrupt_values == set(prompts) + + resume_map: dict[str, str] = { + i.id: f"human input for prompt {i.value}" + for i in parent_graph.get_state(thread_config).interrupts + } + + result = parent_graph.invoke(Command(resume=resume_map), thread_config) + assert result == { + "prompts": prompts, + "human_inputs": [f"human input for prompt {prompt}" for prompt in prompts], + } + + +def test_sync_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.05 + + @task() + def slow() -> dict: + time.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": time.time()} + + @entrypoint() + def graph(inputs: dict) -> list: + first = slow().result() + second = slow().result() + return [first, second] + + arrival_times = [] + + for chunk in graph.stream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(time.time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay + + +def test_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert foo.invoke({"a": "1"}, config) == {"a": "1"} + + @entrypoint() + def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + + +def test_entrypoint_stateful(sync_checkpointer: BaseCheckpointSaver) -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=sync_checkpointer) + def foo(inputs, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + } + assert foo.invoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + }, + } + assert states == [ + None, + {"current": {"a": "1"}, "previous": None}, + {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}}, + ] + + # Test stream + @entrypoint(checkpointer=sync_checkpointer) + def foo(inputs, *, previous: Any) -> Any: + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "2"}} + items = [item for item in foo.stream({"a": "1"}, config)] + assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}] + + +def test_entrypoint_stateful_update_state( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=sync_checkpointer) + def foo(inputs, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + # assert print(foo.input_channels) + foo.update_state(config, {"a": "-1"}) + assert foo.invoke({"a": "1"}, config) == { + "current": {"a": "1"}, + "previous": {"a": "-1"}, + } + assert foo.invoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + } + assert foo.invoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + }, + } + + # update state + foo.update_state(config, {"a": "3"}) + + # Test stream + assert [item for item in foo.stream({"a": "1"}, config)] == [ + {"foo": {"current": {"a": "1"}, "previous": {"a": "3"}}} + ] + assert states == [ + {"a": "-1"}, + {"current": {"a": "1"}, "previous": {"a": "-1"}}, + { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + }, + {"a": "3"}, + ] + + +def test_entrypoint_from_sync_generator() -> None: + """@entrypoint does not support sync generators.""" + previous_return_values = [] + + with pytest.raises(NotImplementedError): + + @entrypoint() + def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + +def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + def call_same_subgraph(state): + result = add_subgraph.invoke(state) + another_result = add_subgraph.invoke({"a": result["result"], "b": 10}) + return another_result + + parent_call_same_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=sync_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + def call_multiple_subgraphs(state): + add_result = add_subgraph.invoke(state) + multiply_result = multiply_subgraph.invoke(state) + return { + "add_result": add_result["result"], + "multiply_result": multiply_result["result"], + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output_schema=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=sync_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } + + +def test_multiple_subgraphs_functional(sync_checkpointer: BaseCheckpointSaver) -> None: + # Define addition subgraph + @entrypoint() + def add(inputs: tuple[int, int]): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + def multiply_task(a, b): + return a * b + + @entrypoint() + def multiply(inputs: tuple[int, int]): + return multiply_task(*inputs).result() + + # Test calling the same subgraph multiple times + @task + def call_same_subgraph(a, b): + result = add.invoke([a, b]) + another_result = add.invoke([result, 10]) + return another_result + + @entrypoint(checkpointer=sync_checkpointer) + def parent_call_same_subgraph(inputs): + return call_same_subgraph(*inputs).result() + + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + def call_multiple_subgraphs(a, b): + add_result = add.invoke([a, b]) + multiply_result = multiply.invoke([a, b]) + return [add_result, multiply_result] + + @entrypoint(checkpointer=sync_checkpointer) + def parent_call_multiple_subgraphs(inputs): + return call_multiple_subgraphs(*inputs).result() + + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] + + +def test_multiple_subgraphs_mixed_entrypoint( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test calling multiple StateGraph subgraphs from an entrypoint.""" + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + @task + def call_same_subgraph(a, b): + result = add_subgraph.invoke({"a": a, "b": b})["result"] + another_result = add_subgraph.invoke({"a": result, "b": 10})["result"] + return another_result + + @entrypoint(checkpointer=sync_checkpointer) + def parent_call_same_subgraph(inputs): + return call_same_subgraph(*inputs).result() + + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + def call_multiple_subgraphs(a, b): + add_result = add_subgraph.invoke({"a": a, "b": b})["result"] + multiply_result = multiply_subgraph.invoke({"a": a, "b": b})["result"] + return [add_result, multiply_result] + + @entrypoint(checkpointer=sync_checkpointer) + def parent_call_multiple_subgraphs(inputs): + return call_multiple_subgraphs(*inputs).result() + + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] + + +def test_multiple_subgraphs_mixed_state_graph( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define addition subgraph + @entrypoint() + def add(inputs: tuple[int, int]): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + def multiply_task(a, b): + return a * b + + @entrypoint() + def multiply(inputs: tuple[int, int]): + return multiply_task(*inputs).result() + + # Test calling the same subgraph multiple times + def call_same_subgraph(state): + result = add.invoke([state["a"], state["b"]]) + another_result = add.invoke([result, 10]) + return {"result": another_result} + + parent_call_same_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=sync_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + def call_multiple_subgraphs(state): + add_result = add.invoke([state["a"], state["b"]]) + multiply_result = multiply.invoke([state["a"], state["b"]]) + return { + "add_result": add_result, + "multiply_result": multiply_result, + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output_schema=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=sync_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } + + +def test_multiple_subgraphs_checkpointer( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + sub_counter: Annotated[int, operator.add] + + def subgraph_node(state): + return {"sub_counter": 2} + + sub_graph_1 = ( + StateGraph(SubgraphState) + .add_node(subgraph_node) + .add_edge(START, "subgraph_node") + .compile(checkpointer=True) + ) + + class OtherSubgraphState(TypedDict): + other_sub_counter: Annotated[int, operator.add] + + def other_subgraph_node(state): + return {"other_sub_counter": 3} + + sub_graph_2 = ( + StateGraph(OtherSubgraphState) + .add_node(other_subgraph_node) + .add_edge(START, "other_subgraph_node") + .compile() + ) + + class ParentState(TypedDict): + parent_counter: int + + def parent_node(state): + result = sub_graph_1.invoke({"sub_counter": state["parent_counter"]}) + other_result = sub_graph_2.invoke({"other_sub_counter": result["sub_counter"]}) + return {"parent_counter": other_result["other_sub_counter"]} + + parent_graph = ( + StateGraph(ParentState) + .add_node(parent_node) + .add_edge(START, "parent_node") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 5} + assert parent_graph.invoke({"parent_counter": 0}, config) == {"parent_counter": 7} + config = {"configurable": {"thread_id": "2"}} + assert [ + c + for c in parent_graph.stream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 5}}), + ] + assert [ + c + for c in parent_graph.stream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 7}}), + ] + + +def test_entrypoint_output_schema_with_return_and_save() -> None: + """Test output schema inference with entrypoint.final.""" + + # Un-parameterized entrypoint.final is interpreted as entrypoint.final[Any, Any] + @entrypoint() + def foo2(inputs, *, previous: Any) -> entrypoint.final: + return entrypoint.final(value="foo", save=1) + + assert foo2.get_output_jsonschema() == { + "title": "LangGraphOutput", + } + + @entrypoint() + def foo(inputs, *, previous: Any) -> entrypoint.final[str, int]: + return entrypoint.final(value="foo", save=1) + + assert foo.get_output_jsonschema() == { + "title": "LangGraphOutput", + "type": "string", + } + + with pytest.raises(TypeError): + # Raise an exception on an improperly parameterized entrypoint.final + # User is attempting to parameterize in this case, so we'll offer + # a bit of help if it's not done correctly. + @entrypoint() + def foo(inputs, *, previous: Any) -> entrypoint.final[int]: + return entrypoint.final(value=1, save=1) # type: ignore + + +def test_entrypoint_with_return_and_save( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test entrypoint with return and save.""" + previous_ = None + + @entrypoint(checkpointer=sync_checkpointer) + def foo(msg: str, *, previous: Any) -> entrypoint.final[int, list[str]]: + nonlocal previous_ + previous_ = previous + previous = previous or [] + return entrypoint.final(value=len(previous), save=previous + [msg]) + + assert foo.get_output_jsonschema() == { + "title": "LangGraphOutput", + "type": "integer", + } + + config = {"configurable": {"thread_id": "1"}} + assert foo.invoke("hello", config) == 0 + assert previous_ is None + assert foo.invoke("goodbye", config) == 1 + assert previous_ == ["hello"] + assert foo.invoke("definitely", config) == 2 + assert previous_ == ["hello", "goodbye"] + + +def test_overriding_injectable_args_with_tasks(sync_store: BaseStore) -> None: + """Test overriding injectable args in tasks.""" + + @task + def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: + assert store is value + assert writer is value + + @entrypoint(store=sync_store) + def main(inputs, store: BaseStore) -> str: + assert store is not None + foo(store=None, writer=None, value=None).result() + foo(store="hello", writer="hello", value="hello").result() + return "OK" + + assert main.invoke({}) == "OK" + + +def test_named_tasks_functional() -> None: + class Foo: + def foo(self, value: str) -> dict: + return value + "foo" + + f = Foo() + + # class method task + foo = task(f.foo, name="custom_foo") + other_foo = task(f.foo, name="other_foo") + + # regular function task + @task(name="custom_bar") + def bar(value: str) -> dict: + return value + "|bar" + + def baz(update: str, value: str) -> dict: + return value + f"|{update}" + + # partial function task (unnamed) + baz_task = task(functools.partial(baz, "baz")) + # partial function task (named_) + custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz") + + class Qux: + def __call__(self, value: str) -> dict: + return value + "|qux" + + qux_task = task(Qux(), name="qux") + + @entrypoint() + def workflow(inputs: dict) -> dict: + foo_result = foo(inputs).result() + other_foo(inputs).result() + fut_bar = bar(foo_result) + fut_baz = baz_task(fut_bar.result()) + fut_custom_baz = custom_baz_task(fut_baz.result()) + fut_qux = qux_task(fut_custom_baz.result()) + return fut_qux.result() + + assert list(workflow.stream("", stream_mode="updates")) == [ + {"custom_foo": "foo"}, + {"other_foo": "foo"}, + {"custom_bar": "foo|bar"}, + {"baz": "foo|bar|baz"}, + {"custom_baz": "foo|bar|baz|custom_baz"}, + {"qux": "foo|bar|baz|custom_baz|qux"}, + {"workflow": "foo|bar|baz|custom_baz|qux"}, + ] + + +def test_tags_stream_mode_messages() -> None: + model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"]) + graph = ( + StateGraph(MessagesState) + .add_node( + "call_model", lambda state: {"messages": model.invoke(state["messages"])} + ) + .add_edge(START, "call_model") + .compile() + ) + assert list( + graph.stream( + { + "messages": "hi", + }, + stream_mode="messages", + ) + ) == [ + ( + _AnyIdAIMessageChunk(content="foo", chunk_position="last"), + { + "langgraph_step": 1, + "langgraph_node": "call_model", + "langgraph_triggers": ("branch:to:call_model",), + "langgraph_path": ("__pregel_pull", "call_model"), + "langgraph_checkpoint_ns": AnyStr("call_model:"), + "checkpoint_ns": AnyStr("call_model:"), + "ls_provider": "genericfakechatmodel", + "ls_model_type": "chat", + "tags": ["meow"], + }, + ) + ] + + +def test_stream_mode_messages_command() -> None: + from langchain_core.messages import HumanMessage + + def my_node(state): + return {"messages": HumanMessage(content="foo")} + + def my_other_node(state): + return Command(update={"messages": HumanMessage(content="bar")}) + + def my_last_node(state): + return [Command(update={"messages": HumanMessage(content="baz")})] + + graph = ( + StateGraph(MessagesState) + .add_sequence([my_node, my_other_node, my_last_node]) + .add_edge(START, "my_node") + .compile() + ) + assert list( + graph.stream( + { + "messages": [], + }, + stream_mode="messages", + ) + ) == [ + ( + _AnyIdHumanMessage(content="foo"), + { + "langgraph_step": 1, + "langgraph_node": "my_node", + "langgraph_triggers": ("branch:to:my_node",), + "langgraph_path": ("__pregel_pull", "my_node"), + "langgraph_checkpoint_ns": AnyStr("my_node:"), + }, + ), + ( + _AnyIdHumanMessage(content="bar"), + { + "langgraph_step": 2, + "langgraph_node": "my_other_node", + "langgraph_triggers": ("branch:to:my_other_node",), + "langgraph_path": ("__pregel_pull", "my_other_node"), + "langgraph_checkpoint_ns": AnyStr("my_other_node:"), + }, + ), + ( + _AnyIdHumanMessage(content="baz"), + { + "langgraph_step": 3, + "langgraph_node": "my_last_node", + "langgraph_triggers": ("branch:to:my_last_node",), + "langgraph_path": ("__pregel_pull", "my_last_node"), + "langgraph_checkpoint_ns": AnyStr("my_last_node:"), + }, + ), + ] + + +def test_node_destinations() -> None: + class State(TypedDict): + foo: Annotated[str, operator.add] + + def node_a(state: State): + value = state["foo"] + if value == "a": + goto = "node_b" + else: + goto = "node_c" + + return Command( + update={"foo": value}, + goto=goto, + graph=Command.PARENT, + ) + + subgraph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() + + # test calling subgraph inside a node function + def call_subgraph(state: State): + return subgraph.invoke(state) + + def node_b(state: State): + return {"foo": "b"} + + def node_c(state: State): + return {"foo": "c"} + + for subgraph_node in (subgraph, call_subgraph): + # destinations w/ tuples + builder = StateGraph(State) + builder.add_edge(START, "child") + builder.add_node("child", subgraph_node, destinations=("node_b", "node_c")) + builder.add_node(node_b) + builder.add_node(node_c) + compiled_graph = builder.compile() + assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} + + graph = compiled_graph.get_graph() + assert [ + Edge(source="__start__", target="child", data=None, conditional=False), + Edge(source="child", target="node_b", data=None, conditional=True), + Edge(source="child", target="node_c", data=None, conditional=True), + Edge(source="node_b", target="__end__", data=None, conditional=False), + Edge(source="node_c", target="__end__", data=None, conditional=False), + ] == graph.edges + + # destinations w/ dicts + builder = StateGraph(State) + builder.add_edge(START, "child") + builder.add_node( + "child", subgraph_node, destinations={"node_b": "foo", "node_c": "bar"} + ) + builder.add_node(node_b) + builder.add_node(node_c) + compiled_graph = builder.compile() + assert compiled_graph.invoke({"foo": ""}) == {"foo": "c"} + + graph = compiled_graph.get_graph() + assert [ + Edge(source="__start__", target="child", data=None, conditional=False), + Edge(source="child", target="node_b", data="foo", conditional=True), + Edge(source="child", target="node_c", data="bar", conditional=True), + Edge(source="node_b", target="__end__", data=None, conditional=False), + Edge(source="node_c", target="__end__", data=None, conditional=False), + ] == graph.edges + + +def test_pydantic_none_state_update() -> None: + class State(BaseModel): + foo: str | None + + def node_a(state: State) -> State: + return State(foo=None) + + graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() + assert graph.invoke({"foo": ""}) == {"foo": None} + + +def test_pydantic_state_update_command() -> None: + class State(BaseModel): + foo: str | None + + def node_a(state: State) -> State: + return Command(update=State(foo=None)) + + graph = StateGraph(State).add_node(node_a).add_edge(START, "node_a").compile() + assert graph.invoke({"foo": ""}) == {"foo": None} + + class State(BaseModel): + foo: str | None = None + bar: str | None = None + + def node_a(state: State): + return State(foo="foo") + + def node_b(state: State): + return Command(update=State(bar="bar")) + + builder = StateGraph(State) + builder.add_node(node_a) + builder.add_node(node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_b", END) + graph = builder.compile() + + assert graph.invoke(State()) == {"foo": "foo", "bar": "bar"} + + +def test_pydantic_state_mutation() -> None: + class Inner(BaseModel): + a: int = 0 + + class State(BaseModel): + inner: Inner = Inner() + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return state + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + # test w/ default_factory + class State(BaseModel): + inner: Inner = Field(default_factory=Inner) + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return state + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + +def test_pydantic_state_mutation_command() -> None: + class Inner(BaseModel): + a: int = 0 + + class State(BaseModel): + inner: Inner = Inner() + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return Command(update=state) + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + # test w/ default_factory + class State(BaseModel): + inner: Inner = Field(default_factory=Inner) + outer: int = 0 + + def my_node(state: State) -> State: + state.inner.a = 5 + state.outer = 10 + return Command(update=state) + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + + assert graph.invoke({"outer": 1}) == {"outer": 10, "inner": Inner(a=5)} + + +def test_get_stream_writer() -> None: + class State(TypedDict): + foo: str + + def my_node(state): + writer = get_stream_writer() + writer("custom!") + return state + + graph = StateGraph(State).add_node(my_node).add_edge(START, "my_node").compile() + assert list(graph.stream({"foo": "bar"}, stream_mode="custom")) == ["custom!"] + assert list(graph.stream({"foo": "bar"}, stream_mode="values")) == [ + {"foo": "bar"}, + {"foo": "bar"}, + ] + assert list(graph.stream({"foo": "bar"}, stream_mode=["custom", "updates"])) == [ + ( + "custom", + "custom!", + ), + ( + "updates", + { + "my_node": { + "foo": "bar", + }, + }, + ), + ] + + +def test_stream_messages_dedupe_inputs() -> None: + from langchain_core.messages import AIMessage + + def call_model(state): + return {"messages": AIMessage("hi", id="1")} + + def route(state): + return Command(goto="node_2", graph=Command.PARENT) + + subgraph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_node(route) + .add_edge(START, "call_model") + .add_edge("call_model", "route") + .compile() + ) + + graph = ( + StateGraph(MessagesState) + .add_node("node_1", subgraph) + .add_node("node_2", lambda state: state) + .add_edge(START, "node_1") + .compile() + ) + + chunks = [ + chunk + for ns, chunk in graph.stream( + {"messages": "hi"}, stream_mode="messages", subgraphs=True + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("hi", id="1") + assert chunks[0][1]["langgraph_node"] == "call_model" + + +def test_stream_messages_dedupe_state(sync_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.messages import AIMessage + + to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] + + def call_model(state): + return {"messages": to_emit.pop(0)} + + def route(state): + return Command(goto="node_2", graph=Command.PARENT) + + subgraph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_node(route) + .add_edge(START, "call_model") + .add_edge("call_model", "route") + .compile() + ) + + graph = ( + StateGraph(MessagesState) + .add_node("node_1", subgraph) + .add_node("node_2", lambda state: state) + .add_edge(START, "node_1") + .compile(checkpointer=sync_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + + chunks = [ + chunk + for ns, chunk in graph.stream( + {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye", id="1") + assert chunks[0][1]["langgraph_node"] == "call_model" + + chunks = [ + chunk + for ns, chunk in graph.stream( + {"messages": "hi again"}, + thread1, + stream_mode="messages", + subgraphs=True, + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye again", id="2") + assert chunks[0][1]["langgraph_node"] == "call_model" + + +def test_stream_messages_dedupe_pydantic_subgraph_interrupt( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Pydantic BaseModel state should not cause duplicate messages when + streaming from subgraphs that use interrupts. Regression test for a bug + where ``on_chain_start`` only populated the ``seen`` set for dict inputs, + skipping Pydantic model inputs entirely.""" + + class PydanticState(BaseModel): + messages: Annotated[list[AnyMessage], add_messages] = Field( + default_factory=list + ) + + def subgraph_proposal(state) -> Command[Literal["subgraph_approval"]]: + return Command( + goto="subgraph_approval", + update={"messages": [AIMessage(content="Proposal", id="proposal_msg")]}, + ) + + def subgraph_approval(state) -> Command[Literal["__end__"]]: + resume_value = interrupt({"message": "Waiting for approval"}) + user_msg = resume_value.get("user_message", "") + msgs = [HumanMessage(content=user_msg)] if user_msg else [] + return Command(goto="__end__", update={"messages": msgs}) + + subgraph = ( + StateGraph(PydanticState) + .add_node("proposal", subgraph_proposal) + .add_node("subgraph_approval", subgraph_approval) + .add_edge(START, "proposal") + .compile(checkpointer=sync_checkpointer) + ) + + def finalize(state) -> Command[Literal["__end__"]]: + return Command( + goto="__end__", + update={"messages": [AIMessage(content="Finalized", id="finalize_msg")]}, + ) + + graph = ( + StateGraph(PydanticState) + .add_node("subgraph", subgraph) + .add_node("finalize", finalize) + .add_edge(START, "subgraph") + .add_edge("subgraph", "finalize") + .compile(checkpointer=sync_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + + # First stream: should hit interrupt after proposal + chunks_req0 = [ + (ns, chunk) + for ns, chunk in graph.stream( + {"messages": [HumanMessage(content="Create a proposal")]}, + thread1, + stream_mode="messages", + subgraphs=True, + ) + ] + + msg_ids_req0 = {chunk[0].id for _, chunk in chunks_req0} + assert "proposal_msg" in msg_ids_req0 + + # Verify interrupted + state = graph.get_state(thread1) + assert state.next + + # Second stream: resume — should NOT duplicate messages from first stream + chunks_req1 = [ + (ns, chunk) + for ns, chunk in graph.stream( + Command(resume={"user_message": "Yes"}), + thread1, + stream_mode="messages", + subgraphs=True, + ) + ] + + msg_ids_req1 = {chunk[0].id for _, chunk in chunks_req1} + assert "finalize_msg" in msg_ids_req1 + + # The key assertion: no message IDs from request 0 should appear in request 1 + duplicates = msg_ids_req0 & msg_ids_req1 + assert not duplicates, f"Duplicate message IDs across requests: {duplicates}" + + +def test_interrupt_subgraph_reenter_checkpointer_true( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + foo: str + bar: str + + class ParentState(TypedDict): + foo: str + counter: int + + called = [] + bar_values = [] + + def subnode_1(state: SubgraphState): + called.append("subnode_1") + bar_values.append(state.get("bar")) + return {"foo": "subgraph_1"} + + def subnode_2(state: SubgraphState): + called.append("subnode_2") + value = interrupt("Provide value") + value += "baz" + return {"foo": "subgraph_2", "bar": value} + + subgraph = ( + StateGraph(SubgraphState) + .add_node(subnode_1) + .add_node(subnode_2) + .add_edge(START, "subnode_1") + .add_edge("subnode_1", "subnode_2") + .compile(checkpointer=True) + ) + + def call_subgraph(state: ParentState): + called.append("call_subgraph") + return subgraph.invoke(state) + + def node(state: ParentState): + called.append("parent") + if state["counter"] < 1: + return Command( + goto="call_subgraph", update={"counter": state["counter"] + 1} + ) + + return {"foo": state["foo"] + "|" + "parent"} + + parent = ( + StateGraph(ParentState) + .add_node(call_subgraph) + .add_node(node) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "node") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert parent.invoke({"foo": "", "counter": 0}, config) == { + "foo": "", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + assert parent.invoke(Command(resume="bar"), config) == { + "foo": "subgraph_2", + "counter": 1, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + assert parent.invoke(Command(resume="qux"), config) == { + "foo": "subgraph_2|parent", + "counter": 1, + } + assert called == [ + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + ] + + # invoke parent again (new turn) + assert parent.invoke({"foo": "meow", "counter": 0}, config) == { + "foo": "meow", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + # confirm that we preserve the state values from the previous invocation + assert bar_values == [None, "barbaz", "quxbaz"] + + +def test_empty_invoke() -> None: + def reducer_merge_dicts( + dict1: dict[Any, Any], dict2: dict[Any, Any] + ) -> dict[Any, Any]: + merged = {**dict1, **dict2} + return merged + + class SimpleGraphState(BaseModel): + x1: Annotated[list[str], operator.add] = [] + x2: Annotated[dict[str, Any], reducer_merge_dicts] = {} + + def update_x1_1(state: SimpleGraphState): + print(state) + return {"x1": ["111"]} + + def update_x1_2(state: SimpleGraphState): + print(state) + state.x1.append("222") + return {"x1": ["222"]} + + def update_x2_1(state: SimpleGraphState): + print(state) + return {"x2": {"111": 111}} + + def update_x2_2(state: SimpleGraphState): + print(state) + return {"x2": {"222": 222}} + + graph = StateGraph(SimpleGraphState) + graph.add_node("x1_1_node", update_x1_1) + graph.add_node("x1_2_node", update_x1_2) + graph.add_node("x2_1_node", update_x2_1) + graph.add_node("x2_2_node", update_x2_2) + graph.add_edge("x1_1_node", "x1_2_node") + graph.add_edge("x1_2_node", "x2_1_node") + graph.add_edge("x2_1_node", "x2_2_node") + + graph.add_edge(START, "x1_1_node") + graph.add_edge("x2_2_node", END) + + compiled = graph.compile() + + assert compiled.invoke(SimpleGraphState()).get("x2") == { + "111": 111, + "222": 222, + } + + +def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None: + # --- CHILD GRAPH --- + + class ChildState(BaseModel): + prompt: str = Field(..., description="What is going to be asked to the user?") + human_input: str | None = Field(None, description="What the human said") + human_inputs: Annotated[list[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def get_human_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_input=human_input, # update child state + human_inputs=[human_input], # update parent state + ) + + child_graph_builder = StateGraph(ChildState) + child_graph_builder.add_node("get_human_input", get_human_input) + child_graph_builder.add_edge(START, "get_human_input") + child_graph_builder.add_edge("get_human_input", END) + child_graph = child_graph_builder.compile() + + # --- PARENT GRAPH --- + + class ParentState(BaseModel): + prompts: list[str] = Field( + ..., description="What is going to be asked to the user?" + ) + human_inputs: Annotated[list[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def assign_workers(state: ParentState): + return [ + Send( + "child_graph", + dict( + prompt=prompt, + ), + ) + for prompt in state.prompts + ] + + def cleanup(state: ParentState): + assert len(state.human_inputs) == len(state.prompts) + + parent_graph_builder = StateGraph(ParentState) + parent_graph_builder.add_node("child_graph", child_graph) + parent_graph_builder.add_node("cleanup", cleanup) + + parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"]) + parent_graph_builder.add_edge("child_graph", "cleanup") + parent_graph_builder.add_edge("cleanup", END) + + parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer) + + # --- CLIENT INVOCATION --- + + thread_config = dict( + configurable=dict( + thread_id=str(uuid.uuid4()), + ) + ) + current_input = dict( + prompts=["a", "b"], + ) + + invokes = 0 + events: dict[int, list[dict]] = {} + while invokes < 10: + # reset interrupt + invokes += 1 + events[invokes] = [] + current_interrupts: list[Interrupt] = [] + + # start / resume the graph + for event in parent_graph.stream( + input=current_input, + config=thread_config, + stream_mode="updates", + ): + events[invokes].append(event) + # handle the interrupt + if "__interrupt__" in event: + current_interrupts.extend(event["__interrupt__"]) + # assume that it breaks here, because it is an interrupt + + # get human input and resume + if len(current_interrupts) > 0: + # we resume one at a time to preserve original test behavior, + # but we could also resume all at once if we wanted + # with a single dict mapping of interrupt ids to resume values + resume = {current_interrupts[0].id: f"Resume #{invokes}"} + current_input = Command(resume=resume) + + # not more human input required, must be completed + else: + break + else: + assert False, "Detected infinite loop" + + assert invokes == 3 + assert len(events) == 3 + + assert events[1] == UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="a", + id=AnyStr(), + ), + ) + }, + { + "__interrupt__": ( + Interrupt( + value="b", + id=AnyStr(), + ), + ) + }, + ) + assert events[2] in ( + UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="a", + id=AnyStr(), + ), + ) + }, + {"child_graph": {"human_inputs": ["Resume #1"]}}, + ), + UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="b", + id=AnyStr(), + ), + ) + }, + {"child_graph": {"human_inputs": ["Resume #1"]}}, + ), + ) + assert events[3] == UnsortedSequence( + { + "child_graph": {"human_inputs": ["Resume #1"]}, + "__metadata__": {"cached": True}, + }, + {"child_graph": {"human_inputs": ["Resume #2"]}}, + {"cleanup": None}, + ) + + +def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None: + # --- CHILD GRAPH --- + + class ChildState(BaseModel): + prompt: str = Field(..., description="What is going to be asked to the user?") + human_input: str | None = Field(None, description="What the human said") + human_inputs: Annotated[list[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def get_human_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_inputs=[human_input], # update parent state + ) + + def get_dolphin_input(state: ChildState): + human_input = interrupt(state.prompt) + + return dict( + human_inputs=[human_input], # update parent state + ) + + child_graph_builder = StateGraph(ChildState) + child_graph_builder.add_node("get_human_input", get_human_input) + child_graph_builder.add_node("get_dolphin_input", get_dolphin_input) + child_graph_builder.add_edge(START, "get_human_input") + child_graph_builder.add_edge(START, "get_dolphin_input") + child_graph = child_graph_builder.compile() + + # --- PARENT GRAPH --- + + class ParentState(BaseModel): + prompts: list[str] = Field( + ..., description="What is going to be asked to the user?" + ) + human_inputs: Annotated[list[str], operator.add] = Field( + default_factory=list, description="All of my messages" + ) + + def assign_workers(state: ParentState): + return [ + Send( + "child_graph", + dict( + prompt=prompt, + ), + ) + for prompt in state.prompts + ] + + def cleanup(state: ParentState): + assert len(state.human_inputs) == len(state.prompts) * 2 + + parent_graph_builder = StateGraph(ParentState) + parent_graph_builder.add_node("child_graph", child_graph) + parent_graph_builder.add_node("cleanup", cleanup) + + parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"]) + parent_graph_builder.add_edge("child_graph", "cleanup") + parent_graph_builder.add_edge("cleanup", END) + + parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer) + + # --- CLIENT INVOCATION --- + + thread_config = dict( + configurable=dict( + thread_id=str(uuid.uuid4()), + ) + ) + current_input = dict( + prompts=["a", "b"], + ) + + invokes = 0 + events: dict[int, list[dict]] = {} + while invokes < 10: + # reset interrupt + invokes += 1 + events[invokes] = [] + current_interrupts: list[Interrupt] = [] + + # start / resume the graph + for event in parent_graph.stream( + input=current_input, + config=thread_config, + stream_mode="updates", + ): + events[invokes].append(event) + # handle the interrupt + if "__interrupt__" in event: + current_interrupts.extend(event["__interrupt__"]) + # assume that it breaks here, because it is an interrupt + + # get human input and resume + if len(current_interrupts) > 0: + # we resume one at a time to preserve original test behavior, + # but we could also resume all at once if we wanted + # with a single dict mapping of interrupt ids to resume values + resume = {current_interrupts[0].id: f"Resume #{invokes}"} + current_input = Command(resume=resume) + + # not more human input required, must be completed + else: + break + else: + assert False, "Detected infinite loop" + + assert invokes == 5 + assert len(events) == 5 + + +def test_pregel_loop_refcount(): + gc.collect() + try: + gc.disable() + + class State(TypedDict): + messages: Annotated[list, add_messages] + + graph_builder = StateGraph(State) + + def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} + + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + graph.invoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable() + + +def test_bulk_state_updates( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "bar"}, as_node="node_a"), + ] + ], + ) + + # Then bulk update with both nodes + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "updated"}, as_node="node_a"), + StateUpdate(values={"baz": "new"}, as_node="node_b"), + ] + ], + ) + + state = graph.get_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 2 + + # perform multiple steps at the same time + config = {"configurable": {"thread_id": "2"}} + + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "bar"}, as_node="node_a"), + ], + [ + StateUpdate(values={"foo": "updated"}, as_node="node_a"), + StateUpdate(values={"baz": "new"}, as_node="node_b"), + ], + ], + ) + + state = graph.get_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + checkpoints = list(sync_checkpointer.list(config)) + assert len(checkpoints) == 2 + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "error"}, as_node=None), + StateUpdate(values={"bar": "error"}, as_node=None), + ] + ], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No supersteps provided"): + graph.bulk_update_state(config, []) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + graph.bulk_update_state(config, [[], []]) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + graph.bulk_update_state( + config, + [ + [ + StateUpdate(values=None, as_node="__end__"), + StateUpdate(values=None, as_node="__copy__"), + ], + ], + ) + + +def test_pregel_node_copy() -> None: + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile() + ) + + graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) + graph.copy() + graph.nodes["agent"].copy({}) + + +def test_update_as_input( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=sync_checkpointer) + ) + + assert graph.invoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == {"foo": "tool"} + + assert graph.invoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == {"foo": "tool"} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + if durability != "exit": + assert new_history == history + else: + assert [new_history[0], new_history[4]] == history + + +def test_batch_update_as_input( + sync_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=sync_checkpointer) + ) + + assert graph.invoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == { + "foo": "map", + "tasks": [0, 1, 2], + } + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "1"}}) + ] + + graph.bulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = graph.get_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + for s in graph.get_state_history({"configurable": {"thread_id": "2"}}) + ] + + if durability != "exit": + assert new_history == history + else: + assert new_history[:1] == history + + +def test_migration_graph(snapshot: SnapshotAssertion) -> None: + class DummyState(BaseModel): + pass_count: int = 0 + + def increment_pass_count(state: DummyState): + state.pass_count += 1 + return state + + def route_b(state: DummyState): + if state.pass_count == 0: + return "X" + else: + return "Y" + + migration_graph = StateGraph(DummyState) + + migration_graph.add_node("B", increment_pass_count) + migration_graph.add_node("C", increment_pass_count) + migration_graph.add_node("D", increment_pass_count) + + migration_graph.add_edge(START, "B") + + migration_graph.add_conditional_edges( + "B", + route_b, + { + "X": "C", + "Y": "D", + }, + ) + + migration_graph.add_edge("D", "B") + migration_graph.add_edge("C", END) + + app = migration_graph.compile() + + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_get_graph_loop(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + foo: str + + def human_node(state: State) -> State: + value = interrupt() + return {"foo": value} + + def agent_node(state: State) -> State: + return {"foo": "Hi " + state["foo"]} + + workflow = StateGraph(State) + workflow.add_node("human", human_node) + workflow.add_node("agent", agent_node) + workflow.add_edge(START, "human") + workflow.add_edge("human", "agent") + workflow.add_edge("agent", "human") + + app = workflow.compile() + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_get_graph_self_loop(snapshot: SnapshotAssertion) -> None: + import random + + subgraph_builder = StateGraph(MessagesState) + subgraph_builder.add_node("agent", lambda x: x) + subgraph_builder.add_edge(START, "agent") + subgraph = subgraph_builder.compile() + + def worker_node(state: MessagesState) -> Command[Literal["worker_node", "__end__"]]: + subgraph_result = subgraph.invoke(state) + + if random.choice([True, False]): + next_node_name = "worker_node" + else: + next_node_name = END + + return Command(update=subgraph_result, goto=next_node_name) + + self_loop_builder = StateGraph(MessagesState) + self_loop_builder.add_node("worker_node", worker_node) + self_loop_builder.add_edge(START, "worker_node") + self_loop_graph = self_loop_builder.compile() + + assert json.dumps(self_loop_graph.get_graph().to_json(), indent=2) == snapshot + assert self_loop_graph.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None: + child_builder = StateGraph(list) + child_builder.add_node("child_node", lambda x: x) + child_builder.add_edge(START, "child_node") + child_graph = child_builder.compile() + + graph_builder = StateGraph(list) + graph_builder.add_node("child", child_graph) + graph_builder.add_edge(START, "child") + graph = graph_builder.compile() + + assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot + assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_imp_exception( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + @task() + def my_task(number: int): + time.sleep(0.1) + return number * 2 + + @task() + def task_with_exception(number: int): + time.sleep(0.1) + raise Exception("This is a test exception") + + @entrypoint(checkpointer=sync_checkpointer) + def my_workflow(number: int): + my_task(number).result() + try: + task_with_exception(number).result() + except Exception as e: + print(f"Exception caught: {e}") + my_task(number).result() + return "done" + + thread1 = {"configurable": {"thread_id": "1"}} + assert my_workflow.invoke(1, thread1) == "done" + + assert [c for c in my_workflow.stream(1, thread1)] == [ + {"my_task": 2}, + {"my_task": 2}, + {"my_workflow": "done"}, + ] + + +@pytest.mark.parametrize("with_timeout", [False, "inner", "outer", "both"]) +@pytest.mark.parametrize("subgraph_persist", [True, False]) +def test_parent_command_goto( + sync_checkpointer: BaseCheckpointSaver, subgraph_persist: bool, with_timeout: bool +) -> None: + class State(TypedDict): + dialog_state: Annotated[list[str], operator.add] + + def node_a_child(state): + return {"dialog_state": ["a_child_state"]} + + def node_b_child(state): + return Command( + graph=Command.PARENT, + goto="node_b_parent", + update={"dialog_state": ["b_child_state"]}, + ) + + sub_builder = StateGraph(State) + sub_builder.add_node(node_a_child) + sub_builder.add_node(node_b_child) + sub_builder.add_edge(START, "node_a_child") + sub_builder.add_edge("node_a_child", "node_b_child") + sub_graph = sub_builder.compile(checkpointer=subgraph_persist) + if with_timeout in ("inner", "both"): + sub_graph.step_timeout = 1 + + def node_b_parent(state): + return {"dialog_state": ["node_b_parent"]} + + main_builder = StateGraph(State) + main_builder.add_node(node_b_parent) + main_builder.add_edge(START, "subgraph_node") + main_builder.add_node("subgraph_node", sub_graph, destinations=("node_b_parent",)) + main_graph = main_builder.compile(sync_checkpointer, name="parent") + if with_timeout in ("outer", "both"): + main_graph.step_timeout = 1 + + config = {"configurable": {"thread_id": 1}} + + assert main_graph.invoke(input={"dialog_state": ["init_state"]}, config=config) == { + "dialog_state": ["init_state", "b_child_state", "node_b_parent"] + } + + +@pytest.mark.parametrize("subgraph_persist", [True, False]) +def test_parent_command_goto_deeply_nested( + sync_checkpointer: BaseCheckpointSaver, + subgraph_persist: bool, +) -> None: + """Test Command.PARENT in a 3-level nested subgraph. + + Command.PARENT should jump to sub_child_3 in the immediate parent (sub_graph). + + Note: With operator.add, subgraph state (including its input) is merged with + parent state, causing the input to appear multiple times. This is expected. + """ + + class State(TypedDict): + dialog_state: Annotated[list[str], operator.add] + + # Level 3: Deepest subgraph that issues Command.PARENT + def sub_sub_child_node(state): + # Jump to immediate parent (sub_graph) + return Command( + graph=Command.PARENT, + goto="sub_child_3", + update={"dialog_state": ["sub_sub_child"]}, + ) + + sub_sub_builder = StateGraph(State) + sub_sub_builder.add_node("sub_sub_child", sub_sub_child_node) + sub_sub_builder.add_edge(START, "sub_sub_child") + sub_sub_graph = sub_sub_builder.compile( + name="sub_sub_graph", checkpointer=subgraph_persist + ) + + # Level 2: Middle subgraph containing Level 3 + def sub_child_1(state): + return {"dialog_state": ["sub_child_1"]} + + def sub_child_3(state): + return {"dialog_state": ["sub_child_3"]} + + sub_builder = StateGraph(State) + sub_builder.add_node("sub_child_1", sub_child_1) + sub_builder.add_node("sub_child_2", sub_sub_graph, destinations=("sub_child_3",)) + sub_builder.add_node("sub_child_3", sub_child_3) + sub_builder.add_edge(START, "sub_child_1") + sub_builder.add_edge("sub_child_1", "sub_child_2") + sub_graph = sub_builder.compile(name="sub_graph", checkpointer=subgraph_persist) + + # Level 1: Main graph containing Level 2 + def child_1(state): + return {"dialog_state": ["child_1"]} + + builder = StateGraph(State) + builder.add_node("child_1", child_1) + builder.add_node("child_2", sub_graph) + builder.add_edge(START, "child_1") + builder.add_edge("child_1", "child_2") + graph = builder.compile(name="main_graph", checkpointer=sync_checkpointer) + + config = {"configurable": {"thread_id": 1}} + + result = graph.invoke(input={"dialog_state": ["init"]}, config=config) + + # Command.PARENT from sub_sub_child jumps to sub_child_3 in immediate parent + # State duplication occurs due to operator.add merging behavior + assert result == { + "dialog_state": [ + "init", + "child_1", + "init", + "child_1", + "sub_child_1", + "sub_sub_child", + "sub_child_3", + ] + } + + +@pytest.mark.parametrize("with_timeout", [True, False]) +def test_timeout_with_parent_command( + sync_checkpointer: BaseCheckpointSaver, with_timeout: bool +) -> None: + """Test that parent commands are properly propagated during timeouts.""" + + class State(TypedDict): + value: str + + def parent_command_node(state: State) -> State: + time.sleep(0.1) # Add some delay before raising + return Command(graph=Command.PARENT, goto="test_cmd", update={"key": "value"}) + + builder = StateGraph(State) + builder.add_node("parent_cmd", parent_command_node) + builder.set_entry_point("parent_cmd") + graph = builder.compile(checkpointer=sync_checkpointer) + if with_timeout: + graph.step_timeout = 1 + + # Should propagate parent command, not timeout + thread1 = {"configurable": {"thread_id": "1"}} + with pytest.raises(ParentCommand) as exc_info: + graph.invoke({"value": "start"}, thread1) + assert exc_info.value.args[0].goto == "test_cmd" + assert exc_info.value.args[0].update == {"key": "value"} + + +def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None: + """Test forking and updating task results with state history.""" + + def checkpoint(values: dict[str, Any]): + return ("checkpoint", {"values": values}) + + def task(name: str, result: Any): + return ("task", {"name": name, "result": result}) + + def get_tree(history: list[StateSnapshot]) -> list: + """Build a tree structure from state history for comparison.""" + if not history: + return [] + + # Build a tree structure similar to renderForks + node_map: dict[str, dict] = {} + root_nodes: list[dict] = [] + + # Second pass: establish parent-child relationships + for item in reversed(history): + checkpoint_id = item.config["configurable"]["checkpoint_id"] + parent_checkpoint_id = ( + item.parent_config["configurable"]["checkpoint_id"] + if item.parent_config + else None + ) + node_map[checkpoint_id] = {"item": item, "children": []} + + parent = node_map.get(parent_checkpoint_id) + (parent["children"] if parent else root_nodes).append( + node_map[checkpoint_id] + ) + + def node_to_tree(node: dict) -> list: + """Convert a node to tree structure.""" + result = [ + checkpoint(node["item"].values), + ] + [ + task(task_info.name, task_info.result) + for task_info in node["item"].tasks + ] + + if len(node["children"]) > 1: + branches = [node_to_tree(child) for child in node["children"]] + return result + [branches] + elif len(node["children"]) == 1: + return result + node_to_tree(node["children"][0]) + else: + return result + + if len(root_nodes) == 1: + # Process all root nodes + return node_to_tree(root_nodes[0]) + + elif len(root_nodes) > 1: + # Multiple root nodes - treat as branches + branches = [node_to_tree(node) for node in root_nodes] + return branches + else: + return [] + + class State(TypedDict): + name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b] + + # Define the graph with a sequence of nodes + def one(state: State) -> Command: + return Command(goto=[Send("two", {})], update={"name": "one"}) + + def two(state: State) -> State: + return {"name": "two"} + + def three(state: State) -> State: + return {"name": "three"} + + graph = ( + StateGraph(State) + .add_node("one", one) + .add_node("two", two) + .add_node("three", three) + .add_edge(START, "one") + .add_edge("one", "two") + .add_edge("two", "three") + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + history: list[StateSnapshot] = [] + + # Initial run + graph.invoke({"name": "start"}, config) + history = list(graph.get_state_history(config)) + + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ] + + # Update the start state + graph.invoke( + None, + graph.update_state( + history[4].config, + values=[StateUpdate(values={"name": "start*"}, as_node="__start__")], + as_node="__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + ] + + # Fork from task "one" + # Start from the checkpoint that has the task "one" + assert history[3].values == {"name": "start*"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "one" + + graph.invoke( + None, + graph.update_state( + history[3].config, + [StateUpdate(values={"name": "one*"}, as_node="one")], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + [ + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one*"}), + checkpoint({"name": "start* > one*"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one* > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one* > two > three"}), + ], + ], + ], + ] + + config = {"configurable": {"thread_id": "2"}} + + # Initialize the thread once again + graph.invoke({"name": "start"}, config) + history = list(graph.get_state_history(config)) + + # Fork from task "two" + # Start from the checkpoint that has the task "two" + assert history[2].values == {"name": "start > one"} + + graph.invoke( + None, + graph.update_state( + history[2].config, + [ + StateUpdate(values={"name": "two"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), + ], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ] + + # Fork task three + assert history[1].values == {"name": "start > one > two > two"} + assert len(history[1].tasks) == 1 + assert history[1].tasks[0].name == "three" + + graph.invoke( + None, + graph.update_state( + history[1].config, + [StateUpdate(values={"name": "three*"}, as_node="three")], + "__copy__", + ), + ) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + ], + ], + ], + ] + + # Regenerate task three + assert history[3].values == {"name": "start > one > two > two"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "three" + + graph.invoke(None, graph.update_state(history[3].config, None, "__copy__")) + + history = list(graph.get_state_history(config)) + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ], + ], + ] + + +def test_subgraph_streaming_sync() -> None: + """Test subgraph streaming when used as a node in sync version""" + + # Create a fake chat model that returns a simple response + model = GenericFakeChatModel(messages=iter(["The weather is sunny today."])) + + # Create a subgraph that uses the fake chat model + def call_model_node(state: MessagesState, config: RunnableConfig) -> MessagesState: + """Node that calls the model with the last message.""" + messages = state["messages"] + last_message = messages[-1].content if messages else "" + response = model.invoke([("user", last_message)], config) + return {"messages": [response]} + + # Build the subgraph + subgraph = StateGraph(MessagesState) + subgraph.add_node("call_model", call_model_node) + subgraph.add_edge(START, "call_model") + compiled_subgraph = subgraph.compile() + + class SomeCustomState(TypedDict): + last_chunk: NotRequired[str] + num_chunks: NotRequired[int] + + # Will invoke a subgraph as a function + def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict: + """Node that runs the subgraph.""" + msgs = {"messages": [("user", "What is the weather in Tokyo?")]} + events = [] + for event in compiled_subgraph.stream(msgs, config, stream_mode="messages"): + events.append(event) + ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events] + return { + "last_chunk": ai_msg_chunks[-1], + "num_chunks": len(ai_msg_chunks), + } + + # Build the main workflow + workflow = StateGraph(SomeCustomState) + workflow.add_node("subgraph", parent_node) + workflow.add_edge(START, "subgraph") + compiled_workflow = workflow.compile() + + # Test the basic functionality + result = compiled_workflow.invoke({}) + + assert result["last_chunk"].content == "today." + assert result["num_chunks"] == 9 + + +def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None: + class State(TypedDict): + messages: list[str] + + def chatbot_node(state: State) -> State: + return {"messages": state["messages"] + ["chatbot"]} + + def tools_node(state: State) -> State: + return {"messages": state["messages"] + ["tools"]} + + def human_node(state: State) -> State: + return {"messages": state["messages"] + ["human"]} + + def tools_condition(_: State) -> str: + return "tools" + + def end_condition(_: State) -> str: + return "chatbot" + + workflow = StateGraph(State) + workflow.add_node("chatbot", chatbot_node) + workflow.add_node("tools", tools_node) + workflow.add_node("human", human_node) + + workflow.add_edge(START, "human") + workflow.add_edge("tools", "chatbot") + + workflow.add_conditional_edges( + "chatbot", tools_condition, {"tools": "tools", "human": "human"} + ) + workflow.add_conditional_edges( + "human", end_condition, {"chatbot": "chatbot", END: END} + ) + + app = workflow.compile() + graph = app.get_graph() + graph_json = graph.to_json() + + assert json.dumps(graph_json, indent=2, sort_keys=True) == snapshot + + +def test_null_resume_disallowed_with_multiple_interrupts( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + text_1: str + text_2: str + + def human_node_1(state: State): + value = interrupt({"text_to_revise": state["text_1"]}) + return {"text_1": value} + + def human_node_2(state: State): + value = interrupt({"text_to_revise": state["text_2"]}) + return {"text_2": value} + + graph_builder = StateGraph(State) + graph_builder.add_node("human_node_1", human_node_1) + graph_builder.add_node("human_node_2", human_node_2) + + # Add both nodes in parallel from START + graph_builder.add_edge(START, "human_node_1") + graph_builder.add_edge(START, "human_node_2") + + checkpointer = InMemorySaver() + graph = graph_builder.compile(checkpointer=checkpointer) + + thread_id = str(uuid.uuid4()) + config: RunnableConfig = {"configurable": {"thread_id": thread_id}} + graph.invoke( + {"text_1": "original text 1", "text_2": "original text 2"}, config=config + ) + + resume_map = { + i.id: f"resume for prompt: {i.value['text_to_revise']}" + for i in graph.get_state(config).interrupts + } + with pytest.raises( + RuntimeError, + match="When there are multiple pending interrupts, you must specify the interrupt id when resuming.", + ): + graph.invoke(Command(resume="singular resume"), config=config) + + assert graph.invoke(Command(resume=resume_map), config=config) == { + "text_1": "resume for prompt: original text 1", + "text_2": "resume for prompt: original text 2", + } + + +def test_interrupt_stream_mode_values(sync_checkpointer: BaseCheckpointSaver): + """Test that interrupts are surfaced on 'values' stream mode""" + + class State(TypedDict): + robot_input: str + human_input: str + + def robot_input_node(state: State) -> State: + return {"robot_input": "beep boop i am a robot"} + + def human_input_node(state: State) -> Command: + human_input = interrupt("interrupt") + return Command(update={"human_input": human_input}) + + builder = StateGraph(State) + builder.add_node(robot_input_node) + builder.add_node(human_input_node) + builder.add_edge(START, "robot_input_node") + builder.add_edge("robot_input_node", "human_input_node") + app = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + result = [*app.stream(State(), config, stream_mode=["updates", "values"])] + assert len(result) == 4 + assert result == [ + ("updates", {"robot_input_node": {"robot_input": "beep boop i am a robot"}}), + ("values", {"robot_input": "beep boop i am a robot"}), + ("updates", {"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}), + ( + "values", + { + "robot_input": "beep boop i am a robot", + "__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),), + }, + ), + ] + resume_result = [ + *app.stream( + Command(resume="i am a human"), config, stream_mode=["updates", "values"] + ) + ] + assert resume_result == [ + ("values", {"robot_input": "beep boop i am a robot"}), + ("updates", {"human_input_node": {"human_input": "i am a human"}}), + ( + "values", + {"robot_input": "beep boop i am a robot", "human_input": "i am a human"}, + ), + ] + + +def test_supersteps_populate_task_results( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + num: int + text: str + + def double(state: State) -> State: + return {"num": state["num"] * 2, "text": state["text"] * 2} + + graph = ( + StateGraph(State) + .add_node("double", double) + .add_edge(START, "double") + .add_edge("double", END) + .compile(checkpointer=sync_checkpointer) + ) + + def first_task_result(history: list[StateSnapshot], node: str) -> Any: + for s in history: + for t in s.tasks: + if t.name == node: + return t.result + return None + + # reference run with invoke + ref_cfg = {"configurable": {"thread_id": "ref"}} + graph.invoke({"num": 1, "text": "one"}, ref_cfg) + ref_history = list(graph.get_state_history(ref_cfg)) + + ref_start_result = first_task_result(ref_history, "__start__") + ref_double_result = first_task_result(ref_history, "double") + assert ref_start_result == {"num": 1, "text": "one"} + assert ref_double_result == {"num": 2, "text": "oneone"} + + # using supersteps + bulk_cfg = {"configurable": {"thread_id": "bulk"}} + graph.bulk_update_state( + bulk_cfg, + [ + [StateUpdate(values={}, as_node="__input__")], + [StateUpdate(values={"num": 1, "text": "one"}, as_node="__start__")], + [StateUpdate(values={"num": 2, "text": "oneone"}, as_node="double")], + ], + ) + bulk_history = list(graph.get_state_history(bulk_cfg)) + + bulk_start_result = first_task_result(bulk_history, "__start__") + bulk_double_result = first_task_result(bulk_history, "double") + + assert bulk_start_result == ref_start_result == {"num": 1, "text": "one"} + assert bulk_double_result == ref_double_result == {"num": 2, "text": "oneone"} + + +def test_multiple_writes_same_channel_from_same_node( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that a node can write multiple times to the same channel and that writes are ordered, reduced, and reflected in streamed events and state history.""" + + class State(TypedDict): + foo: Annotated[str, lambda a, b: ", ".join([x for x in [a, b] if x])] + + def one(_: State) -> Command: + return Command(update=[("foo", "one.0"), ("foo", "one.1")]) + + def two(_: State) -> State: + return {"foo": "two"} + + graph = ( + StateGraph(State) + .add_node("one", one) + .add_node("two", two) + .add_edge(START, "one") + .add_edge("one", "two") + .add_edge("two", END) + .compile(checkpointer=sync_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + events = [ + (ns, ev) + for ns, ev in graph.stream( + {"foo": "input"}, config, stream_mode=["updates", "tasks"] + ) + ] + + assert events == [ + ( + "tasks", + { + "id": AnyStr(), + "name": "one", + "input": {"foo": "input"}, + "triggers": ("branch:to:one",), + }, + ), + ("updates", {"one": [{"foo": "one.0"}, {"foo": "one.1"}]}), + ( + "tasks", + { + "id": AnyStr(), + "name": "one", + "error": None, + "result": {"foo": {"$writes": ["one.0", "one.1"]}}, + "interrupts": [], + }, + ), + ( + "tasks", + { + "id": AnyStr(), + "name": "two", + "input": {"foo": "input, one.0, one.1"}, + "triggers": ("branch:to:two",), + }, + ), + ("updates", {"two": {"foo": "two"}}), + ( + "tasks", + { + "id": AnyStr(), + "name": "two", + "error": None, + "result": {"foo": "two"}, + "interrupts": [], + }, + ), + ] + + def map_snapshot(s: StateSnapshot) -> dict: + return { + "tasks": [{"name": t.name, "result": t.result} for t in s.tasks], + "values": s.values, + } + + history = [map_snapshot(s) for s in graph.get_state_history(config)] + + assert history == [ + { + "tasks": [], + "values": {"foo": "input, one.0, one.1, two"}, + }, + { + "tasks": [{"name": "two", "result": {"foo": "two"}}], + "values": {"foo": "input, one.0, one.1"}, + }, + { + "tasks": [ + {"name": "one", "result": {"foo": {"$writes": ["one.0", "one.1"]}}} + ], + "values": {"foo": "input"}, + }, + { + "tasks": [{"name": "__start__", "result": {"foo": "input"}}], + "values": {"foo": ""}, + }, + ] + + +def test_send_with_untracked_value(sync_checkpointer: BaseCheckpointSaver): + """Test that Send objects work correctly with untracked values in state.""" + + class UnserializableResource: + def __init__(self, name: str): + self.name = name + self.lock = threading.Lock() + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + session_resource: Annotated[UnserializableResource, UntrackedValue] + + def setup_node(state: State) -> State: + resource = UnserializableResource("test_session") + return {"messages": ["setup complete"], "session_resource": resource} + + def send_to_tool(state: State): + return [Send("tool_node", state)] + + def tool_node(state: State) -> State: + resource = state["session_resource"] + assert isinstance(resource, UnserializableResource) + assert resource.name == "test_session" + + new_resource = UnserializableResource("new_session") + + return { + "messages": [f"tool used resource: {resource.name}"], + "session_resource": new_resource, + } + + graph = StateGraph(State) + graph.add_node("setup", setup_node) + graph.add_node("tool_node", tool_node) + graph.add_edge(START, "setup") + graph.add_conditional_edges("setup", send_to_tool) + + app = graph.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + result = app.invoke({}, config) + + assert len(result["messages"]) == 2 + assert result["messages"][0] == "setup complete" + assert result["messages"][1] == "tool used resource: test_session" + assert result["session_resource"].name == "new_session" + + state = app.get_state(config) + assert "session_resource" not in state.values + + +def test_send_with_untracked_value_overlapping_keys( + sync_checkpointer: BaseCheckpointSaver, +): + """Test that Send objects work correctly with untracked values in state.""" + + class State(TypedDict): + dictionary: dict + session_resource: Annotated[str, UntrackedValue] + + def setup_node(state: State) -> State: + return {} + + def send_to_tool(state: State): + return [ + Send( + "tool_node", + { + "dictionary": {"session_resource": "legal_value"}, + "session_resource": "illegal_value", + }, + ) + ] + + def tool_node(state: State) -> State: + print(f"STATE: {state}") + assert state["dictionary"] == {"session_resource": "legal_value"} + assert state["session_resource"] == "illegal_value" + + return { + "dictionary": state["dictionary"], + "session_resource": "new_illegal_value", + } + + graph = StateGraph(State) + graph.add_node("setup", setup_node) + graph.add_node("tool_node", tool_node) + graph.add_edge(START, "setup") + graph.add_conditional_edges("setup", send_to_tool) + + app = graph.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + result = app.invoke({}, config) + + assert result["session_resource"] == "new_illegal_value" + state = app.get_state(config) + assert "session_resource" not in state.values + assert state.values.get("dictionary") == {"session_resource": "legal_value"} + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_overwrite_sequential( + sync_checkpointer: BaseCheckpointSaver, as_json: bool +) -> None: + """Test a sequential chain of nodes where the last node uses Overwrite to bypass a reducer and write a value directly to the channel.""" + + class State(TypedDict): + messages: Annotated[list, operator.add] + + def node_a(state: State): + return {"messages": ["a"]} + + def node_b(state: State): + overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"]) + return {"messages": overwrite} + + builder = StateGraph(State) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + result = graph.invoke({"messages": ["START"]}, config) + # a is overwritten by b + assert result == {"messages": ["b"]} + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_overwrite_parallel( + sync_checkpointer: BaseCheckpointSaver, as_json: bool +) -> None: + """Test parallel nodes where max one node uses Overwrite to bypass a reducer and write a value directly to the channel.""" + + class State(TypedDict): + messages: Annotated[list, operator.add] + + def node_a(state: State): + return {"messages": ["a"]} + + def node_b(state: State): + overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"]) + return {"messages": overwrite} + + def node_c(state: State): + return {"messages": ["c"]} + + def node_d(state: State): + return {"messages": ["d"]} + + builder = StateGraph(State) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_node("node_c", node_c) + builder.add_node("node_d", node_d) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_a", "node_c") + builder.add_edge("node_b", "node_d") + builder.add_edge("node_c", "node_d") + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + result = graph.invoke({"messages": ["START"]}, config) + # a, c are overwritten by b, then d is written + assert result == {"messages": ["b", "d"]} + + +@pytest.mark.parametrize("as_json", [False, True]) +def test_overwrite_parallel_error( + sync_checkpointer: BaseCheckpointSaver, as_json: bool +) -> None: + """Test parallel nodes where more than one node uses Overwrite to bypass a reducer and write a value directly to the channel. In this case, InvalidUpdateError should be raised.""" + + class State(TypedDict): + messages: Annotated[list, operator.add] + + def node_a(state: State): + return {"messages": ["a"]} + + def node_b(state: State): + overwrite = {"__overwrite__": ["b"]} if as_json else Overwrite(["b"]) + return {"messages": overwrite} + + def node_c(state: State): + overwrite = {"__overwrite__": ["c"]} if as_json else Overwrite(["c"]) + return {"messages": overwrite} + + builder = StateGraph(State) + builder.add_node("node_a", node_a) + builder.add_node("node_b", node_b) + builder.add_node("node_c", node_c) + builder.add_edge(START, "node_a") + builder.add_edge("node_a", "node_b") + builder.add_edge("node_a", "node_c") + builder.add_edge("node_b", END) + builder.add_edge("node_c", END) + + graph = builder.compile(checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "1"}} + with pytest.raises( + InvalidUpdateError, match="Can receive only one Overwrite value per super-step." + ): + graph.invoke({"messages": ["START"]}, config) + + +def test_fork_does_not_apply_pending_writes( + sync_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that forking with update_state does not apply pending writes from original execution.""" + + class State(TypedDict): + value: Annotated[int, operator.add] + + def node_a(state: State) -> State: + return {"value": 10} + + def node_b(state: State) -> State: + return {"value": 100} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=sync_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + graph.invoke({"value": 1}, thread1) + + history = list(graph.get_state_history(thread1)) + checkpoint_before_a = next(s for s in history if s.next == ("node_a",)) + + fork_config = graph.update_state( + checkpoint_before_a.config, {"value": 20}, as_node="node_a" + ) + + # Continue from fork (should run node_b) + result = graph.invoke(None, fork_config) + + # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 + assert result == {"value": 121} diff --git a/langgraph/source/libs/langgraph/tests/test_pregel_async.py b/langgraph/source/libs/langgraph/tests/test_pregel_async.py new file mode 100644 index 0000000000000000000000000000000000000000..d3a655bcf886adf6b7b0a4b8caafcda198abf2b9 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_pregel_async.py @@ -0,0 +1,9349 @@ +import asyncio +import enum +import functools +import gc +import logging +import operator +import random +import sys +import uuid +from collections import Counter, deque +from dataclasses import replace +from time import perf_counter +from typing import ( + Annotated, + Any, + Literal, + Optional, +) +from uuid import UUID + +import pytest +from langchain_core.language_models import GenericFakeChatModel +from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough +from langchain_core.utils.aiter import aclosing +from langgraph.cache.base import BaseCache +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + ChannelVersions, + Checkpoint, + CheckpointMetadata, + CheckpointTuple, +) +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer +from langgraph.prebuilt.tool_node import ToolNode +from langgraph.store.base import BaseStore +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion +from typing_extensions import NotRequired, TypedDict + +from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL +from langgraph._internal._queue import AsyncQueue +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.last_value import LastValue +from langgraph.channels.topic import Topic +from langgraph.errors import ( + GraphRecursionError, + InvalidUpdateError, + ParentCommand, +) +from langgraph.func import entrypoint, task +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import MessagesState, add_messages +from langgraph.pregel import NodeBuilder, Pregel +from langgraph.pregel._loop import AsyncPregelLoop +from langgraph.pregel._runner import PregelRunner +from langgraph.types import ( + CachePolicy, + Command, + Durability, + Interrupt, + PregelTask, + RetryPolicy, + Send, + StateSnapshot, + StateUpdate, + StreamWriter, + interrupt, +) +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.fake_tracer import FakeTracer +from tests.memory_assert import MemorySaverNoPending +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + +logger = logging.getLogger(__name__) + +pytestmark = pytest.mark.anyio + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + + +async def test_checkpoint_errors() -> None: + class FaultyGetCheckpointer(InMemorySaver): + async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + raise ValueError("Faulty get_tuple") + + class FaultyPutCheckpointer(InMemorySaver): + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + raise ValueError("Faulty put") + + class FaultyPutWritesCheckpointer(InMemorySaver): + async def aput_writes( + self, config: RunnableConfig, writes: list[tuple[str, Any]], task_id: str + ) -> RunnableConfig: + raise ValueError("Faulty put_writes") + + class FaultyVersionCheckpointer(InMemorySaver): + def get_next_version(self, current: int | None, channel: None) -> int: + raise ValueError("Faulty get_next_version") + + class FaultySerializer(JsonPlusSerializer): + def dumps_typed(self, obj: Any) -> tuple[str, bytes]: + raise ValueError("Faulty serializer") + + def logic(inp: str) -> str: + return "" + + builder = StateGraph(Annotated[str, operator.add]) + builder.add_node("agent", logic) + builder.add_edge(START, "agent") + + graph = builder.compile(checkpointer=InMemorySaver(serde=FaultySerializer())) + with pytest.raises(ValueError, match="Faulty serializer"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty serializer"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty serializer"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + graph = builder.compile(checkpointer=FaultyGetCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_tuple"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty get_tuple"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty get_tuple"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + graph = builder.compile(checkpointer=FaultyPutCheckpointer()) + with pytest.raises(ValueError, match="Faulty put"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty put"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty put"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_next_version"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty get_next_version"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty get_next_version"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + # add a parallel node + builder.add_node("parallel", logic) + builder.add_edge(START, "parallel") + graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer()) + with pytest.raises(ValueError, match="Faulty put_writes"): + await graph.ainvoke( + "", {"configurable": {"thread_id": "thread-1"}}, durability="async" + ) + with pytest.raises(ValueError, match="Faulty put_writes"): + async for _ in graph.astream( + "", {"configurable": {"thread_id": "thread-2"}}, durability="async" + ): + pass + with pytest.raises(ValueError, match="Faulty put_writes"): + async for _ in graph.astream_events( + "", + {"configurable": {"thread_id": "thread-3"}}, + version="v2", + durability="async", + ): + pass + + def faulty_reducer(a: Any, b: Any) -> Any: + raise ValueError("Faulty reducer") + + builder = StateGraph(Annotated[str, faulty_reducer]) + builder.add_node("agent", logic) + builder.add_edge(START, "agent") + graph = builder.compile(checkpointer=InMemorySaver()) + + with pytest.raises(ValueError, match="Faulty reducer"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty reducer"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty reducer"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + +async def test_py_async_with_cancel_behavior() -> None: + """This test confirms that in all versions of Python we support, __aexit__ + is not cancelled when the coroutine containing the async with block is cancelled.""" + + logs: list[str] = [] + + class MyContextManager: + async def __aenter__(self): + logs.append("Entering") + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + logs.append("Starting exit") + try: + # Simulate some cleanup work + await asyncio.sleep(2) + logs.append("Cleanup completed") + except asyncio.CancelledError: + logs.append("Cleanup was cancelled!") + raise + logs.append("Exit finished") + + async def main(): + try: + async with MyContextManager(): + logs.append("In context") + await asyncio.sleep(1) + logs.append("This won't print if cancelled") + except asyncio.CancelledError: + logs.append("Context was cancelled") + raise + + # create task + t = asyncio.create_task(main()) + # cancel after 0.2 seconds + await asyncio.sleep(0.2) + t.cancel() + # check logs before cancellation is handled + assert logs == [ + "Entering", + "In context", + ], "Cancelled before cleanup started" + # wait for task to finish + try: + await t + except asyncio.CancelledError: + # check logs after cancellation is handled + assert logs == [ + "Entering", + "In context", + "Starting exit", + "Cleanup completed", + "Exit finished", + "Context was cancelled", + ], "Cleanup started and finished after cancellation" + else: + assert False, "Task should be cancelled" + + +async def test_checkpoint_put_after_cancellation() -> None: + logs: list[str] = [] + + class LongPutCheckpointer(InMemorySaver): + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + logs.append("checkpoint.aput.start") + try: + await asyncio.sleep(1) + return await super().aput(config, checkpoint, metadata, new_versions) + finally: + logs.append("checkpoint.aput.end") + + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + logs.append("awhile.start") + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + finally: + logs.append("awhile.end") + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile(checkpointer=LongPutCheckpointer()) + thread1 = {"configurable": {"thread_id": "1"}} + + # start the task + t = asyncio.create_task( + graph.ainvoke({"hello": "world"}, thread1, durability="exit") + ) + # cancel after 0.2 seconds + await asyncio.sleep(0.2) + t.cancel() + # check logs before cancellation is handled + assert sorted(logs) == [ + "awhile.start", + ], "Cancelled before checkpoint put started" + # wait for task to finish + try: + await t + except asyncio.CancelledError: + # check logs after cancellation is handled + assert sorted(logs) == [ + "awhile.end", + "awhile.start", + "checkpoint.aput.end", + "checkpoint.aput.start", + ], "Checkpoint put is not cancelled" + else: + assert False, "Task should be cancelled" + + +async def test_checkpoint_put_after_cancellation_stream_anext() -> None: + logs: list[str] = [] + + class LongPutCheckpointer(InMemorySaver): + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + logs.append("checkpoint.aput.start") + try: + await asyncio.sleep(1) + return await super().aput(config, checkpoint, metadata, new_versions) + finally: + logs.append("checkpoint.aput.end") + + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + logs.append("awhile.start") + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + finally: + logs.append("awhile.end") + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile(checkpointer=LongPutCheckpointer()) + thread1 = {"configurable": {"thread_id": "1"}} + + # start the task + s = graph.astream({"hello": "world"}, thread1, durability="exit") + t = asyncio.create_task(s.__anext__()) + # cancel after 0.2 seconds + await asyncio.sleep(0.2) + t.cancel() + # check logs before cancellation is handled + assert sorted(logs) == [ + "awhile.start", + ], "Cancelled before checkpoint put started" + # wait for task to finish + try: + await t + except asyncio.CancelledError: + # check logs after cancellation is handled + assert sorted(logs) == [ + "awhile.end", + "awhile.start", + "checkpoint.aput.end", + "checkpoint.aput.start", + ], "Checkpoint put is not cancelled" + else: + assert False, "Task should be cancelled" + + +async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None: + logs: list[str] = [] + + class LongPutCheckpointer(InMemorySaver): + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + logs.append("checkpoint.aput.start") + try: + await asyncio.sleep(1) + return await super().aput(config, checkpoint, metadata, new_versions) + finally: + logs.append("checkpoint.aput.end") + + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + logs.append("awhile.start") + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + finally: + logs.append("awhile.end") + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile(checkpointer=LongPutCheckpointer()) + thread1 = {"configurable": {"thread_id": "1"}} + + # start the task + s = graph.astream_events( + {"hello": "world"}, + thread1, + version="v2", + include_names=["LangGraph"], + durability="exit", + ) + # skip first event (happens right away) + await s.__anext__() + # start the task for 2nd event + t = asyncio.create_task(s.__anext__()) + # cancel after 0.2 seconds + await asyncio.sleep(0.2) + t.cancel() + # check logs before cancellation is handled + assert logs == [ + "awhile.start", + ], "Cancelled before checkpoint put started" + # wait for task to finish + try: + await t + except asyncio.CancelledError: + # check logs after cancellation is handled + assert logs == [ + "awhile.start", + "awhile.end", + "checkpoint.aput.start", + "checkpoint.aput.end", + ], "Checkpoint put is not cancelled" + else: + assert False, "Task should be cancelled" + + +async def test_node_cancellation_on_external_cancel() -> None: + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile() + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(graph.ainvoke({"hello": "world"}), 0.5) + + assert inner_task_cancelled + + +async def test_node_cancellation_on_other_node_exception() -> None: + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def iambad(input: Any) -> None: + raise ValueError("I am bad") + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.add_node("bad", iambad) + builder.set_conditional_entry_point(lambda _: ["agent", "bad"]) + + graph = builder.compile() + + with pytest.raises(ValueError, match="I am bad"): + # This will raise ValueError, not TimeoutError + await asyncio.wait_for(graph.ainvoke({"hello": "world"}), 0.5) + + assert inner_task_cancelled + + +async def test_node_cancellation_on_other_node_exception_two() -> None: + async def awhile(input: Any) -> None: + await asyncio.sleep(1) + + async def iambad(input: Any) -> None: + raise ValueError("I am bad") + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.add_node("bad", iambad) + builder.set_conditional_entry_point(lambda _: ["agent", "bad"]) + + graph = builder.compile() + + with pytest.raises(ValueError, match="I am bad"): + # This will raise ValueError, not CancelledError + await graph.ainvoke({"hello": "world"}) + + +@NEEDS_CONTEXTVARS +async def test_dynamic_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_node_count = 0 + + async def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt is cleared, as well as the next tasks + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) + + +@NEEDS_CONTEXTVARS +async def test_dynamic_interrupt_subgraph( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + my_key: str + market: str + + tool_two_node_count = 0 + + def tool_two_node(s: SubgraphState) -> SubgraphState: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + subgraph = StateGraph(SubgraphState) + subgraph.add_node("do", tool_two_node, retry_policy=RetryPolicy()) + subgraph.add_edge(START, "do") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", subgraph.compile()) + tool_two_graph.add_edge(START, "tool_two") + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value", + "market": "DE", + "__interrupt__": [ + Interrupt( + value="Just because...", + id=AnyStr(), + ) + ], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer", "market": "DE"}}, + ] + + # flow: interrupt -> clear + thread1 = {"configurable": {"thread_id": "1"}} + thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt is cleared, as well as the next tasks + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + +@NEEDS_CONTEXTVARS +async def test_partial_pending_checkpoint( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def tool_one(s: State) -> State: + return {"my_key": " one"} + + tool_two_node_count = 0 + + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + def start(state: State) -> list[Send | str]: + return ["tool_two", Send("tool_one", state)] + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry_policy=RetryPolicy()) + tool_two_graph.add_node("tool_one", tool_one) + tool_two_graph.set_conditional_entry_point(start) + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True + ) == { + "my_key": "value one", + "market": "DE", + "__interrupt__": [Interrupt(value="Just because...", id=AnyStr())], + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value one"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value all good one", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ) + }, + { + "tool_one": {"my_key": " one"}, + }, + ) + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + { + "__metadata__": {"cached": True}, + "tool_one": {"my_key": " one"}, + }, + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke( + {"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit" + ) == { + "my_key": "value ⛰️ one", + "market": "DE", + "__interrupt__": [ + Interrupt( + value="Just because...", + id=AnyStr(), + ) + ], + } + + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + }, + ] + + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + name="tool_one", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result={"my_key": " one"}, + ), + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + }, + parent_config=None, + interrupts=( + Interrupt( + value="Just because...", + id=AnyStr(), + ), + ), + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared, finished tasks are kept + tup_upd = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=(), + tasks=(), + config=tup_upd.config, + created_at=tup_upd.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + }, + parent_config=tup.config, + interrupts=(), + ) + + +@NEEDS_CONTEXTVARS +async def test_node_not_cancelled_on_other_node_interrupted( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + hello: Annotated[str, operator.add] + + awhiles = 0 + inner_task_cancelled = False + + async def awhile(input: State) -> None: + nonlocal awhiles + + awhiles += 1 + try: + await asyncio.sleep(1) + return {"hello": " again"} + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def iambad(input: State) -> None: + return {"hello": interrupt("I am bad")} + + builder = StateGraph(State) + builder.add_node("agent", awhile) + builder.add_node("bad", iambad) + builder.set_conditional_entry_point(lambda _: ["agent", "bad"]) + + graph = builder.compile(checkpointer=async_checkpointer) + thread = {"configurable": {"thread_id": "1"}} + + # writes from "awhile" are applied to last chunk + assert await graph.ainvoke({"hello": "world"}, thread) == { + "hello": "world again", + "__interrupt__": [ + Interrupt( + value="I am bad", + id=AnyStr(), + ) + ], + } + + assert not inner_task_cancelled + assert awhiles == 1 + + assert await graph.ainvoke(None, thread, debug=True) == { + "hello": "world again", + "__interrupt__": [ + Interrupt( + value="I am bad", + id=AnyStr(), + ) + ], + } + + assert not inner_task_cancelled + assert awhiles == 1 + + # resume with answer + assert await graph.ainvoke(Command(resume=" okay"), thread) == { + "hello": "world again okay" + } + + assert not inner_task_cancelled + assert awhiles == 1 + + +@pytest.mark.parametrize("stream_hang_s", [0.3, 0.6]) +async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None: + inner_task_cancelled = False + + async def awhile(input: Any) -> None: + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + nonlocal inner_task_cancelled + inner_task_cancelled = True + raise + + async def alittlewhile(input: Any) -> None: + await asyncio.sleep(0.6) + return {"hello": "1"} + + class State(TypedDict): + hello: str + + builder = StateGraph(State) + builder.add_node(awhile) + builder.add_node(alittlewhile) + builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"]) + graph = builder.compile() + graph.step_timeout = 1 + + with pytest.raises(asyncio.TimeoutError): + async for chunk in graph.astream({"hello": "world"}, stream_mode="updates"): + assert chunk == {"alittlewhile": {"hello": "1"}} + await asyncio.sleep(stream_hang_s) + + assert inner_task_cancelled + + +async def test_cancel_graph_astream(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + value: Annotated[int, operator.add] + + class AwhileMaker: + def __init__(self) -> None: + self.reset() + + async def __call__(self, input: State) -> Any: + self.started = True + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + self.cancelled = True + raise + + def reset(self): + self.started = False + self.cancelled = False + + async def alittlewhile(input: State) -> None: + await asyncio.sleep(0.6) + return {"value": 2} + + awhile = AwhileMaker() + aparallelwhile = AwhileMaker() + builder = StateGraph(State) + builder.add_node("awhile", awhile) + builder.add_node("aparallelwhile", aparallelwhile) + builder.add_node(alittlewhile) + builder.add_edge(START, "alittlewhile") + builder.add_edge(START, "aparallelwhile") + builder.add_edge("alittlewhile", "awhile") + + graph = builder.compile(checkpointer=async_checkpointer) + + # test interrupting astream + got_event = False + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + async with aclosing(graph.astream({"value": 1}, thread1)) as stream: + async for chunk in stream: + assert chunk == {"alittlewhile": {"value": 2}} + got_event = True + break + + assert got_event + + # node aparallelwhile should start, but be cancelled + assert aparallelwhile.started is True + assert aparallelwhile.cancelled is True + + # node "awhile" should never start + assert awhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + # but we should have applied pending writes + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 3} # 1 + 2 + assert state.next == ("aparallelwhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + } + + +async def test_cancel_graph_astream_events_v2( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + value: int + + class AwhileMaker: + def __init__(self) -> None: + self.reset() + + async def __call__(self, input: State) -> Any: + self.started = True + try: + await asyncio.sleep(1.5) + except asyncio.CancelledError: + self.cancelled = True + raise + + def reset(self): + self.started = False + self.cancelled = False + + async def alittlewhile(input: State) -> None: + await asyncio.sleep(0.6) + return {"value": 2} + + awhile = AwhileMaker() + anotherwhile = AwhileMaker() + builder = StateGraph(State) + builder.add_node(alittlewhile) + builder.add_node("awhile", awhile) + builder.add_node("anotherwhile", anotherwhile) + builder.add_edge(START, "alittlewhile") + builder.add_edge("alittlewhile", "awhile") + builder.add_edge("awhile", "anotherwhile") + + graph = builder.compile(checkpointer=async_checkpointer) + + # test interrupting astream_events v2 + got_event = False + thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} + async with aclosing( + graph.astream_events({"value": 1}, thread2, version="v2") + ) as stream: + async for chunk in stream: + if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: + got_event = True + assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} + await asyncio.sleep(0.1) + break + + # did break + assert got_event + + # node "awhile" maybe starts (impl detail of astream_events) + # if it does start, it must be cancelled + if awhile.started: + assert awhile.cancelled is True + + # node "anotherwhile" should never start + assert anotherwhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + state = await graph.aget_state(thread2) + assert state is not None + assert state.values == {"value": 2} + assert state.next == ("awhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 1, + } + + +async def test_node_schemas_custom_output() -> None: + class State(TypedDict): + hello: str + bye: str + messages: Annotated[list[str], add_messages] + + class Output(TypedDict): + messages: list[str] + + class StateForA(TypedDict): + hello: str + messages: Annotated[list[str], add_messages] + + async def node_a(state: StateForA): + assert state == { + "hello": "there", + "messages": [_AnyIdHumanMessage(content="hello")], + } + + class StateForB(TypedDict): + bye: str + now: int + + async def node_b(state: StateForB): + assert state == { + "bye": "world", + } + return { + "now": 123, + "hello": "again", + } + + class StateForC(TypedDict): + hello: str + now: int + + async def node_c(state: StateForC): + assert state == { + "hello": "again", + "now": 123, + } + + builder = StateGraph(State, output_schema=Output) + builder.add_node("a", node_a) + builder.add_node("b", node_b) + builder.add_node("c", node_c) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", "c") + graph = builder.compile() + + assert await graph.ainvoke( + {"hello": "there", "bye": "world", "messages": "hello"} + ) == { + "messages": [_AnyIdHumanMessage(content="hello")], + } + + builder = StateGraph(State, output_schema=Output) + builder.add_node("a", node_a) + builder.add_node("b", node_b) + builder.add_node("c", node_c) + builder.add_edge(START, "a") + builder.add_edge("a", "b") + builder.add_edge("b", "c") + graph = builder.compile() + + assert await graph.ainvoke( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) == { + "messages": [_AnyIdHumanMessage(content="hello")], + } + + assert [ + c + async for c in graph.astream( + { + "hello": "there", + "bye": "world", + "messages": "hello", + "now": 345, # ignored because not in input schema + } + ) + ] == [ + {"a": None}, + {"b": {"hello": "again", "now": 123}}, + {"c": None}, + ] + + +async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={ + "one": chain, + }, + channels={ + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "integer", + } + assert await app.ainvoke(2) == 3 + assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3} + + +async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = ( + NodeBuilder() + .subscribe_only("input") + .do(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) + + app = Pregel( + nodes={"one": chain}, + channels={ + "input": LastValue(int), + "output": LastValue(int), + "fixed": LastValue(int), + "output_plus_one": LastValue(int), + }, + output_channels=["output", "fixed", "output_plus_one"], + input_channels="input", + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None}, + "fixed": {"title": "Fixed", "type": "integer", "default": None}, + "output_plus_one": { + "title": "Output Plus One", + "type": "integer", + "default": None, + }, + }, + } + assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} + + +async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels=["output"], + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } + assert await app.ainvoke(2) == {"output": 3} + + +async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels=["input"], + output_channels=["output"], + ) + + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer", "default": None}}, + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } + assert await app.ainvoke({"input": 2}) == {"output": 3} + + +async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + stream_channels=["inbox", "output"], + ) + + assert await app.ainvoke(2) == 4 + + with pytest.raises(GraphRecursionError): + await app.ainvoke(2, {"recursion_limit": 1}) + + step = 0 + async for values in app.astream(2): + step += 1 + if step == 1: + assert values == { + "inbox": 3, + } + elif step == 2: + assert values == { + "inbox": 3, + "output": 4, + } + assert step == 2 + + +async def test_batch_two_processes_in_out() -> None: + async def add_one_with_delay(inp: int) -> int: + await asyncio.sleep(inp / 10) + return inp + 1 + + one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one") + two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] + assert await app.abatch([3, 2, 1, 3, 5], output_keys=["output"]) == [ + {"output": 5}, + {"output": 4}, + {"output": 3}, + {"output": 5}, + {"output": 7}, + ] + + +async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: + test_size = 100 + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} + for i in range(test_size - 2): + nodes[str(i)] = ( + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) + ) + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") + + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + # No state is left over from previous invocations + for _ in range(10): + assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size + + # Concurrent invocations do not interfere with each other + assert await asyncio.gather( + *(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10)) + ) == [2 + test_size for _ in range(10)] + + +async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: + test_size = 100 + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} + for i in range(test_size - 2): + nodes[str(i)] = ( + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) + ) + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") + + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + # No state is left over from previous invocations + for _ in range(3): + # Then invoke pubsub + assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ + 2 + test_size, + 1 + test_size, + 3 + test_size, + 4 + test_size, + 5 + test_size, + ] + + # Concurrent invocations do not interfere with each other + assert await asyncio.gather( + *(app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) for _ in range(3)) + ) == [ + [2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size] + for _ in range(3) + ] + + +async def test_invoke_two_processes_two_in_two_out_invalid( + mocker: MockerFixture, +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={"output": LastValue(int), "input": LastValue(int)}, + input_channels="input", + output_channels="output", + ) + + with pytest.raises(InvalidUpdateError): + # LastValue channels can only be updated once per iteration + await app.ainvoke(2) + + +async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "output": Topic(int), + }, + input_channels="input", + output_channels="output", + ) + + # An Topic channel accumulates updates into a sequence + assert await app.ainvoke(2) == [3, 3] + + +async def test_invoke_checkpoint( + mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + errored_once = False + + def raise_if_above_10(input: int) -> int: + nonlocal errored_once + if input > 4: + if errored_once: + pass + else: + errored_once = True + raise ConnectionError("I will be retried") + if input > 10: + raise ValueError("Input is too large") + return input + + one = ( + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) + ) + + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + retry_policy=RetryPolicy(), + ) + + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 2 + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "2"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 5 + + +async def test_pending_writes_resume( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + value: Annotated[int, operator.add] + + class AwhileMaker: + def __init__(self, sleep: float, rtn: dict | Exception) -> None: + self.sleep = sleep + self.rtn = rtn + self.reset() + + async def __call__(self, input: State) -> Any: + self.calls += 1 + await asyncio.sleep(self.sleep) + if isinstance(self.rtn, Exception): + raise self.rtn + else: + return self.rtn + + def reset(self): + self.calls = 0 + + one = AwhileMaker(0.1, {"value": 2}) + two = AwhileMaker(0.2, ConnectionError("I'm not good")) + builder = StateGraph(State) + builder.add_node("one", one) + builder.add_node( + "two", + two, + retry_policy=RetryPolicy(max_attempts=2, initial_interval=0, jitter=False), + ) + builder.add_edge(START, "one") + builder.add_edge(START, "two") + graph = builder.compile(checkpointer=async_checkpointer) + + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ConnectionError, match="I'm not good"): + await graph.ainvoke({"value": 1}, thread1, durability=durability) + + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 2 + + # latest checkpoint should be before nodes "one", "two" + # but we should have applied pending writes from "one" + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 3} + assert state.next == ("two",) + assert state.tasks == ( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), + PregelTask( + AnyStr(), + "two", + (PULL, "two"), + 'ConnectionError("I\'m not good")', + ), + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + } + # get_state with checkpoint_id should not apply any pending writes + state = await graph.aget_state(state.config) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + # should contain pending write of "one" + checkpoint = await async_checkpointer.aget_tuple(thread1) + assert checkpoint is not None + # should contain error from "two" + expected_writes = [ + (AnyStr(), "value", 2), + (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), + ] + assert len(checkpoint.pending_writes) == 2 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] + + # resume execution + with pytest.raises(ConnectionError, match="I'm not good"): + await graph.ainvoke(None, thread1, durability=durability) + + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 4 + + # confirm no new checkpoints saved + state_two = await graph.aget_state(thread1) + assert state_two.metadata == state.metadata + + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert await graph.ainvoke(None, thread1, durability=durability) == {"value": 6} + + # check all final checkpoints + checkpoints = [c async for c in async_checkpointer.alist(thread1)] + # we should have 3 + assert len(checkpoints) == (3 if durability != "exit" else 2) + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": { + "one": { + "branch:to:one": AnyVersion(), + }, + "two": { + "branch:to:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": {"value": 6}, + "updated_channels": ["value"], + }, + metadata={ + "parents": {}, + "step": 1, + "source": "loop", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "branch:to:one": None, + "branch:to:two": None, + }, + "updated_channels": ["branch:to:one", "branch:to:two", "value"], + }, + metadata={ + "parents": {}, + "step": 0, + "source": "loop", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], + } + } + if durability != "exit" + else None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + (AnyStr(), "value", 3), + ) + if durability != "exit" + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved + ), + ) + if durability == "exit": + return + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 4, + "id": AnyStr(), + "ts": AnyStr(), + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, + "updated_channels": ["__start__"], + }, + metadata={ + "parents": {}, + "step": -1, + "source": "input", + }, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "branch:to:one", None), + (AnyStr(), "branch:to:two", None), + ), + ) + + +async def test_run_from_checkpoint_id_retains_previous_writes( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + async def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + graph = builder.compile(checkpointer=async_checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = await graph.ainvoke({"myval": 1}, thread1, durability="async") + assert result["myval"] == 4 + history = [c async for c in graph.aget_state_history(thread1)] + + assert len(history) == 4 + assert history[0].values == {"myval": 4, "otherval": False} + assert history[-1].values == {"myval": 0} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = await graph.ainvoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + +async def test_cond_edge_after_send() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + async def __call__(self, state): + return [self.name] + + async def send_for_fun(state): + return [Send("2", state), Send("2", state)] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + + assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] + + +async def test_concurrent_emit_sends() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + async def __call__(self, state): + return ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + + async def send_for_fun(state): + return [Send("2", 1), Send("2", 2), "3.1"] + + async def send_for_profit(state): + return [Send("2", 3), Send("2", 4)] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("1.1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_edge(START, "1.1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("1.1", send_for_profit) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert await graph.ainvoke(["0"]) == ( + [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + ) + + +async def test_send_sequences(async_checkpointer: BaseCheckpointSaver) -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + async def __call__(self, state): + update = ( + [self.name] + if isinstance(state, list) # or isinstance(state, Control) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Command): + return replace(state, update=update) + else: + return update + + async def send_for_fun(state): + return [ + Send("2", Command(goto=Send("2", 3))), + Send("2", Command(goto=Send("2", 4))), + "3.1", + ] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert await graph.ainvoke(["0"]) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["3.1"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1) == [ + "0", + "1", + ] + assert await graph.ainvoke(None, thread1) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + + +@NEEDS_CONTEXTVARS +async def test_imp_task( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + mapper_calls = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls + mapper_calls += 1 + await asyncio.sleep(0.1 * input) + return str(input) * 2 + + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] + + tracer = FakeTracer() + thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} + result = [c async for c in graph.astream([0, 1], thread1, durability=durability)] + # mapper tasks run concurrently so output order is non-deterministic + assert sorted(result[:-1], key=lambda d: str(d)) == [ + {"mapper": "00"}, + {"mapper": "11"}, + ] + assert result[-1] == { + "__interrupt__": ( + Interrupt( + value="question", + id=AnyStr(), + ), + ) + } + assert mapper_calls == 2 + assert len(tracer.runs) == 1 + assert len(tracer.runs[0].child_runs) == 1 + entrypoint_run = tracer.runs[0].child_runs[0] + assert entrypoint_run.name == "graph" + mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"] + assert len(mapper_runs) == 2 + assert any(r.inputs == {"input": 0} for r in mapper_runs) + assert any(r.inputs == {"input": 1} for r in mapper_runs) + + assert await graph.ainvoke( + Command(resume="answer"), thread1, durability=durability + ) == [ + "00answer", + "11answer", + ] + assert mapper_calls == 2 + + +@NEEDS_CONTEXTVARS +async def test_imp_nested( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + async def mynode(input: list[str]) -> list[str]: + return [it + "a" for it in input] + + builder = StateGraph(list[str]) + builder.add_node(mynode) + builder.add_edge(START, "mynode") + add_a = builder.compile() + + @task + def submapper(input: int) -> str: + return str(input) + + @task + async def mapper(input: int) -> str: + await asyncio.sleep(input / 100) + return await submapper(input) * 2 + + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + final = [m + answer for m in mapped] + return await add_a.ainvoke(final) + + assert graph.get_input_jsonschema() == { + "type": "array", + "items": {"type": "integer"}, + "title": "LangGraphInput", + } + assert graph.get_output_jsonschema() == { + "type": "array", + "items": {"type": "string"}, + "title": "LangGraphOutput", + } + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream([0, 1], thread1, durability=durability)] == [ + {"submapper": "0"}, + {"mapper": "00"}, + {"submapper": "1"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + id=AnyStr(), + ), + ) + }, + ] + + assert await graph.ainvoke( + Command(resume="answer"), thread1, durability=durability + ) == [ + "00answera", + "11answera", + ] + + +@NEEDS_CONTEXTVARS +async def test_imp_task_cancel( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + mapper_calls = 0 + mapper_cancels = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls, mapper_cancels + mapper_calls += 1 + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + mapper_cancels += 1 + raise + return str(input) * 2 + + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + await asyncio.sleep(0.1) + futures.pop().cancel() # cancel one + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream([0, 1], thread1, durability=durability)] == [ + {"mapper": "00"}, + { + "__interrupt__": ( + Interrupt( + value="question", + id=AnyStr(), + ), + ) + }, + ] + assert mapper_calls == 2 + assert mapper_cancels == 1 + + assert await graph.ainvoke( + Command(resume="answer"), thread1, durability=durability + ) == [ + "00answer", + ] + assert mapper_calls == 3 + assert mapper_cancels == 2 + + +@NEEDS_CONTEXTVARS +async def test_imp_sync_from_async( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task + def bar(a: str, b: str, c: str | None = None) -> dict: + return {"a": a + b, "c": (c or "") + "bark"} + + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @entrypoint(checkpointer=async_checkpointer) + def graph(state: dict) -> dict: + foo_result = foo(state).result() + fut_bar = bar(foo_result["a"], foo_result["b"]) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c async for c in graph.astream({"a": "0"}, thread1, durability=durability) + ] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + +@NEEDS_CONTEXTVARS +async def test_imp_stream_order( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + @task() + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task + async def bar(a: str, b: str, c: str | None = None) -> dict: + return {"a": a + b, "c": (c or "") + "bark"} + + @task() + async def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + foo_res = await foo(state) + + fut_bar = bar(foo_res["a"], foo_res["b"]) + fut_baz = baz(await fut_bar) + return await fut_baz + + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c async for c in graph.astream({"a": "0"}, thread1, durability=durability) + ] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Requires Python 3.11 or higher for context management", +) +async def test_send_dedupe_on_resume( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + interrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Command): + return replace(state, update=update) + else: + return update + + def send_for_fun(state): + return [ + Send("2", Command(goto=Send("2", 3))), + Send("2", Command(goto=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1, durability=durability) == { + "__interrupt__": [ + Interrupt( + value="Bahh", + id=AnyStr(), + ), + ], + } + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # resume execution + assert await graph.ainvoke(None, thread1, durability=durability) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check history + history = [c async for c in graph.aget_state_history(thread1)] + assert len(history) == (6 if durability != "exit" else 2) + expected_history = [ + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + "3", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 4, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "3", + "2|3", + "flaky|4", + ], + next=("3",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 3, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + ], + next=("2", "flaky", "3"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=("__pregel_push", 1, False), + error=None, + interrupts=(Interrupt(value="Bahh", id=AnyStr()),), + state=None, + result=["flaky|4"] if durability != "exit" else None, + ), + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + ), + interrupts=(Interrupt(value="Bahh", id=AnyStr()),), + ), + StateSnapshot( + values=["0", "1"], + next=("2", "2", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 1, False), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=["0"], + next=("1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + ), + interrupts=(), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + interrupts=(), + ), + ] + if durability != "exit": + assert history == expected_history + else: + assert history[0] == expected_history[0]._replace( + parent_config=history[1].config + ) + assert history[1] == expected_history[2]._replace(parent_config=None) + + +async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + async def agent(state): + return {"messages": ai_message} + + def route(state): + if isinstance(state["messages"][-1], AIMessage): + return [ + Send(call["name"], call) for call in state["messages"][-1].tool_calls + ] + + foo_called = 0 + + async def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + builder.add_conditional_edges("agent", route) + graph = builder.compile() + + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert await graph.ainvoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + interrupts=(), + ) + + # remove the tool call, clearing the pending task + await graph.aupdate_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert await graph.aget_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ), + ] + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + tasks=(), + interrupts=(), + ) + + # tool call not executed + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + # interrupt-update-resume flow, creating new Send in update call + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "3"}} + assert await graph.ainvoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + interrupts=(), + ) + + # replace the tool call, should clear previous send, create new one + await graph.aupdate_state( + thread1, + { + "messages": AIMessage( + "", + id=ai_message.id, + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ) + }, + ) + + # prev tool call no longer in pending tasks, new tool call is + assert await graph.aget_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + interrupts=(), + ) + + # prev tool call not executed, new tool call is + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"), + ] + } + assert foo_called == 1 + + +async def test_send_react_interrupt_control( + async_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + async def agent(state) -> Command[Literal["foo"]]: + return Command( + update={"messages": ai_message}, + goto=[Send(call["name"], call) for call in ai_message.tool_calls], + ) + + foo_called = 0 + + async def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + graph = builder.compile() + if isinstance(async_checkpointer, InMemorySaver): + assert graph.get_graph().draw_mermaid() == snapshot + + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert await graph.ainvoke( + {"messages": [HumanMessage("hello")]}, thread1, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + interrupts=(), + ) + + # remove the tool call, clearing the pending task + await graph.aupdate_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert await graph.aget_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ), + ] + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + tasks=(), + interrupts=(), + ) + + # tool call not executed + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + +async def test_max_concurrency(async_checkpointer: BaseCheckpointSaver) -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + self.currently = 0 + self.max_currently = 0 + + async def __call__(self, state): + self.currently += 1 + if self.currently > self.max_currently: + self.max_currently = self.currently + await asyncio.sleep(random.random() / 10) + self.currently -= 1 + return [state] + + def one(state): + return ["1"] + + def three(state): + return ["3"] + + async def send_to_many(state): + return [Send("2", idx) for idx in range(100)] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + node2 = Node("2") + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node("1", one) + builder.add_node(node2) + builder.add_node("3", three) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_to_many) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + + assert await graph.ainvoke(["0"]) == ["0", "1", *range(100), "3"] + assert node2.max_currently == 100 + assert node2.currently == 0 + node2.max_currently = 0 + + assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ + "0", + "1", + *range(100), + "3", + ] + assert node2.max_currently == 10 + assert node2.currently == 0 + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["2"]) + thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} + + assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"] + state = await graph.aget_state(thread1) + assert state.values == ["0", "1"] + assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] + + +async def test_max_concurrency_control(async_checkpointer: BaseCheckpointSaver) -> None: + async def node1(state) -> Command[Literal["2"]]: + return Command(update=["1"], goto=[Send("2", idx) for idx in range(100)]) + + node2_currently = 0 + node2_max_currently = 0 + + async def node2(state) -> Command[Literal["3"]]: + nonlocal node2_currently, node2_max_currently + node2_currently += 1 + if node2_currently > node2_max_currently: + node2_max_currently = node2_currently + await asyncio.sleep(0.1) + node2_currently -= 1 + + return Command(update=[state], goto="3") + + async def node3(state) -> Literal["3"]: + return ["3"] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node("1", node1) + builder.add_node("2", node2) + builder.add_node("3", node3) + builder.add_edge(START, "1") + graph = builder.compile() + + if isinstance(async_checkpointer, InMemorySaver): + assert ( + graph.get_graph().draw_mermaid() + == """--- +config: + flowchart: + curve: linear +--- +graph TD; + __start__([__start__
]):::first + 1(1) + 2(2) + 3(3) + __end__([__end__
]):::last + 1 -.-> 2; + 2 -.-> 3; + __start__ --> 1; + 3 --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc +""" + ) + + assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"] + assert node2_max_currently == 100 + assert node2_currently == 0 + node2_max_currently = 0 + + assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ + "0", + "1", + *range(100), + "3", + ] + assert node2_max_currently == 10 + assert node2_currently == 0 + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["2"]) + thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} + + assert await graph.ainvoke(["0"], thread1) == ["0", "1"] + assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] + + +async def test_invoke_checkpoint_three( + mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + + def raise_if_above_10(input: int) -> int: + if input > 10: + raise ValueError("Input is too large") + return input + + one = ( + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) + ) + + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + debug=True, + ) + + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1, durability="async") == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await async_checkpointer.aget(thread_1))["id"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1, durability="async") == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await async_checkpointer.aget(thread_1))["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1, durability="async") + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1, durability="async") == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () + + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == () + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () + + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config + ) + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await async_checkpointer.aget(thread_1_history[0].config))[ + "id" + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] + assert (await async_checkpointer.aget(thread_1_history[1].config))[ + "id" + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] + + thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # 1 more checkpoint in history + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state(thread_1_next_config) + + +async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) + + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_three = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_four = ( + NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output") + ) + + app = Pregel( + nodes={ + "one": one, + "chain_three": chain_three, + "chain_four": chain_four, + }, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + # Then invoke app + # We get a single array result as chain_four waits for all publishers to finish + # before operating on all elements published to topic_two as an array + for _ in range(100): + assert await app.ainvoke(2) == [13, 13] + + assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [ + [13, 13] for _ in range(100) + ] + + +async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + one = ( + NodeBuilder().subscribe_only("input").do(add_one).write_to("output", "between") + ) + two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + stream_channels=["output", "between"], + input_channels="input", + output_channels="output", + ) + + # Then invoke pubsub + assert [c async for c in app.astream(2)] == [ + {"between": 3, "output": 3}, + {"between": 3, "output": 4}, + ] + + +async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between") + two = NodeBuilder().subscribe_only("between").do(add_one) + + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) + + # It finishes executing (once no more messages being published) + # but returns nothing, as nothing was published to "output" topic + assert await app.ainvoke(2) is None + + +async def test_conditional_entrypoint_graph_state() -> None: + class AgentState(TypedDict, total=False): + input: str + output: str + steps: Annotated[list[str], operator.add] + + async def left(data: AgentState) -> AgentState: + return {"output": data["input"] + "->left"} + + async def right(data: AgentState) -> AgentState: + return {"output": data["input"] + "->right"} + + def should_start(data: AgentState) -> str: + assert data["steps"] == [], "Expected input to be read from the state" + # Logic to decide where to start + if len(data["input"]) > 10: + return "go-right" + else: + return "go-left" + + # Define a new graph + workflow = StateGraph(AgentState) + + workflow.add_node("left", left) + workflow.add_node("right", right) + + workflow.set_conditional_entry_point( + should_start, {"go-left": "left", "go-right": "right"} + ) + + workflow.add_conditional_edges("left", lambda data: END) + workflow.add_edge("right", END) + + app = workflow.compile() + + assert await app.ainvoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "output": "what is weather in sf->right", + "steps": [], + } + + assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + {"right": {"output": "what is weather in sf->right"}}, + ] + + +async def test_in_one_fan_out_state_graph_waiting_edge( + async_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +@pytest.mark.parametrize("use_waiting_edge", (True, False)) +async def test_in_one_fan_out_state_graph_defer_node( + async_checkpointer: BaseCheckpointSaver, use_waiting_edge: bool +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa, defer=True) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + if use_waiting_edge: + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + else: + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( + async_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", lambda _: "retriever_two", {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}, debug=True) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +async def test_nested_pydantic_models() -> None: + """Test that nested Pydantic models are properly constructed from leaf nodes up.""" + + class NestedModel(BaseModel): + value: int + name: str + something: str | None = None + + # Forward reference model + class RecursiveModel(BaseModel): + value: str + child: Optional["RecursiveModel"] = None + + # Discriminated union models + class Cat(BaseModel): + pet_type: Literal["cat"] + meow: str + + class Dog(BaseModel): + pet_type: Literal["dog"] + bark: str + + # Cyclic reference model + class Person(BaseModel): + id: str + name: str + friends: list[str] = Field(default_factory=list) # IDs of friends + + class MyEnum(enum.Enum): + A = 1 + B = 2 + + class MyTypedDict(TypedDict): + x: int + my_enum: MyEnum + + class State(BaseModel): + # Basic nested model tests + top_level: str + nested: NestedModel + optional_nested: NestedModel | None = None + dict_nested: dict[str, NestedModel] + my_set: set[int] + another_set: set + my_enum: MyEnum + list_nested: Annotated[ + dict | list[dict[str, NestedModel]], lambda x, y: (x or []) + [y] + ] + list_nested_reversed: Annotated[ + list[dict[str, NestedModel]] | NestedModel | dict | list, + lambda x, y: (x or []) + [y], + ] + tuple_nested: tuple[str, NestedModel] + tuple_list_nested: list[tuple[int, NestedModel]] + complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]] + my_typed_dict: MyTypedDict + + # Forward reference test + recursive: RecursiveModel + + # Discriminated union test + pet: Cat | Dog + + # Cyclic reference test + people: dict[str, Person] # Map of ID -> Person + + inputs = { + # Basic nested models + "top_level": "initial", + "nested": {"value": 42, "name": "test"}, + "optional_nested": {"value": 10, "name": "optional"}, + "my_set": [1, 2, 7], + "another_set": ["foo", 3], + "my_enum": MyEnum.B, + "my_typed_dict": {"x": 1, "my_enum": MyEnum.A}, + "dict_nested": {"a": {"value": 5, "name": "a"}}, + "list_nested": [{"a": {"value": 6, "name": "b"}}], + "list_nested_reversed": ["foo", "bar"], + "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}], + "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]], + "complex_tuple": [ + "complex", + {"nested": [9, {"value": 10, "name": "deep"}]}, + ], + # Forward reference + "recursive": {"value": "parent", "child": {"value": "child", "child": None}}, + # Discriminated union (using a cat in this case) + "pet": {"pet_type": "cat", "meow": "meow!"}, + # Cyclic references + "people": { + "1": { + "id": "1", + "name": "Alice", + "friends": ["2", "3"], # Alice is friends with Bob and Charlie + }, + "2": { + "id": "2", + "name": "Bob", + "friends": ["1"], # Bob is friends with Alice + }, + "3": { + "id": "3", + "name": "Charlie", + "friends": ["1", "2"], # Charlie is friends with Alice and Bob + }, + }, + } + + update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}} + + async def node_fn(state: State) -> dict: + assert state == State(**inputs) + return update + + builder = StateGraph(State) + builder.add_node("process", node_fn) + builder.set_entry_point("process") + builder.set_finish_point("process") + graph = builder.compile() + + result = await graph.ainvoke(inputs.copy()) + + assert result == {**inputs, **update} + + +async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( + async_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + query: str + answer: str | None = None + docs: Annotated[list[str], sorted_add] + + class Input(BaseModel): + query: str + + class Output(BaseModel): + answer: str + docs: list[str] + + class StateUpdate(BaseModel): + query: str | None = None + answer: str | None = None + docs: list[str] | None = None + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data.query}"} + + async def analyzer_one(data: State) -> State: + return StateUpdate(query=f"analyzed: {data.query}") + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data.docs)} + + async def decider(data: State) -> str: + assert isinstance(data, State) + return "retriever_two" + + workflow = StateGraph(State, input_schema=Input, output_schema=Output) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", decider, {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + with pytest.raises(ValidationError): + await app.ainvoke({"query": {}}) + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } + } + + +async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( + snapshot: SnapshotAssertion, async_checkpointer: BaseCheckpointSaver +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class InnerObject(BaseModel): + yo: int + + class State(BaseModel): + query: str + inner: InnerObject + answer: str | None = None + docs: Annotated[list[str], sorted_add] + + class StateUpdate(BaseModel): + query: str | None = None + answer: str | None = None + docs: list[str] | None = None + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data.query}"} + + async def analyzer_one(data: State) -> State: + return StateUpdate(query=f"analyzed: {data.query}") + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data.docs)} + + async def decider(data: State) -> str: + assert isinstance(data, State) + return "retriever_two" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", decider, {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + if isinstance(async_checkpointer, InMemorySaver): + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert app.get_input_jsonschema() == snapshot + assert app.get_output_jsonschema() == snapshot + + with pytest.raises(ValidationError): + await app.ainvoke({"query": {}}) + + assert await app.ainvoke( + {"query": "what is weather in sf", "inner": {"yo": 1}} + ) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + "inner": {"yo": 1}, + } + + assert [ + c + async for c in app.astream( + {"query": "what is weather in sf", "inner": {"yo": 1}} + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } + } + + +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + async_checkpointer: BaseCheckpointSaver, +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + await asyncio.sleep(0.1) + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.2) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + # silly edge, to make sure having been triggered before doesn't break + # semantics of named barrier (== waiting edges) + workflow.add_edge("rewrite_query", "qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + +@pytest.mark.parametrize("with_cache", [True, False]) +async def test_in_one_fan_out_state_graph_waiting_edge_multiple( + with_cache: bool, cache: BaseCache +) -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + rewrite_query_count = 0 + + async def rewrite_query(data: State) -> State: + nonlocal rewrite_query_count + rewrite_query_count += 1 + return {"query": f"query: {data['query']}"} + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + async def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node( + "rewrite_query", + rewrite_query, + cache_policy=CachePolicy() if with_cache else None, + ) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile(cache=cache) + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + assert rewrite_query_count == 2 + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + { + "rewrite_query": {"query": "query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}, + "__metadata__": {"cached": True}, + } + if with_cache + else { + "rewrite_query": {"query": "query: analyzed: query: what is weather in sf"} + }, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + } + }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + ] + assert rewrite_query_count == 2 if with_cache else 4 + + # clear the cache + if with_cache: + await app.aclear_cache() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + assert rewrite_query_count == 4 + + +async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: + def sorted_add(x: list[str], y: list[str] | list[tuple[str, str]]) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + async def retriever_picker(data: State) -> list[str]: + return ["analyzer_one", "retriever_two"] + + async def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + await asyncio.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + async def decider(data: State) -> None: + return None + + def decider_cond(data: State) -> str: + if data["query"].count("analyzed") > 1: + return "qa" + else: + return "rewrite_query" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("decider", decider) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_conditional_edges("rewrite_query", retriever_picker) + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_edge(["retriever_one", "retriever_two"], "decider") + workflow.add_conditional_edges("decider", decider_cond) + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: analyzed: query: what is weather in sf", + "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", + "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, + { + "analyzer_one": { + "query": "analyzed: query: analyzed: query: what is weather in sf" + } + }, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"decider": None}, + {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, + ] + + +async def test_nested_graph(snapshot: SnapshotAssertion) -> None: + def never_called_fn(state: Any): + assert 0, "This function should never be called" + + never_called = RunnableLambda(never_called_fn) + + class InnerState(TypedDict): + my_key: str + my_other_key: str + + def up(state: InnerState): + return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]} + + inner = StateGraph(InnerState) + inner.add_node("up", up) + inner.set_entry_point("up") + inner.set_finish_point("up") + + class State(TypedDict): + my_key: str + never_called: Any + + async def side(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile()) + graph.add_node("side", side) + graph.set_entry_point("inner") + graph.add_edge("inner", "side") + graph.set_finish_point("side") + + app = graph.compile() + + assert await app.ainvoke({"my_key": "my value", "never_called": never_called}) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [ + chunk + async for chunk in app.astream( + {"my_key": "my value", "never_called": never_called} + ) + ] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + assert [ + chunk + async for chunk in app.astream( + {"my_key": "my value", "never_called": never_called}, stream_mode="values" + ) + ] == [ + {"my_key": "my value", "never_called": never_called}, + {"my_key": "my value there", "never_called": never_called}, + {"my_key": "my value there and back again", "never_called": never_called}, + ] + times_called = 0 + async for event in app.astream_events( + {"my_key": "my value", "never_called": never_called}, + version="v2", + config={"run_id": UUID(int=0)}, + stream_mode="values", + ): + if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): + times_called += 1 + assert event["data"] == { + "output": { + "my_key": "my value there and back again", + "never_called": never_called, + } + } + assert times_called == 1 + times_called = 0 + async for event in app.astream_events( + {"my_key": "my value", "never_called": never_called}, + version="v2", + config={"run_id": UUID(int=0)}, + ): + if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): + times_called += 1 + assert event["data"] == { + "output": { + "my_key": "my value there and back again", + "never_called": never_called, + } + } + assert times_called == 1 + + chain = app | RunnablePassthrough() + + assert await chain.ainvoke( + {"my_key": "my value", "never_called": never_called} + ) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [ + chunk + async for chunk in chain.astream( + {"my_key": "my value", "never_called": never_called} + ) + ] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + times_called = 0 + async for event in chain.astream_events( + {"my_key": "my value", "never_called": never_called}, + version="v2", + config={"run_id": UUID(int=0)}, + ): + if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): + times_called += 1 + assert event["data"] == { + "output": {"side": {"my_key": "my value there and back again"}} + } + assert times_called == 1 + + +async def test_subgraph_checkpoint_true( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + + app = graph.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + durability=durability, + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + +async def test_subgraph_durability_inherited( + durability: Durability, +) -> None: + async_checkpointer = InMemorySaver() + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + inner_app = inner.compile(checkpointer=async_checkpointer) + graph = StateGraph(State) + graph.add_node("inner", inner_app) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=async_checkpointer) + thread_id = str(uuid.uuid4()) + config = {"configurable": {"thread_id": thread_id}} + await app.ainvoke({"my_key": ""}, config, subgraphs=True, durability=durability) + if durability != "exit": + checkpoints = list(async_checkpointer.list(config)) + assert len(checkpoints) == 12 + else: + checkpoints = list(async_checkpointer.list(config)) + assert len(checkpoints) == 1 + + +@NEEDS_CONTEXTVARS +async def test_subgraph_checkpoint_true_interrupt( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + # Define subgraph + class SubgraphState(TypedDict): + # note that none of these keys are shared with the parent graph state + bar: str + baz: str + + def subgraph_node_1(state: SubgraphState): + baz_value = interrupt("Provide baz value") + return {"baz": baz_value} + + def subgraph_node_2(state: SubgraphState): + return {"bar": state["bar"] + state["baz"]} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + foo: str + + def node_1(state: ParentState): + return {"foo": "hi! " + state["foo"]} + + async def node_2(state: ParentState, config: RunnableConfig): + response = await subgraph.ainvoke({"bar": state["foo"]}) + return {"foo": response["bar"]} + + builder = StateGraph(ParentState) + builder.add_node("node_1", node_1) + builder.add_node("node_2", node_2) + builder.add_edge(START, "node_1") + builder.add_edge("node_1", "node_2") + + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke({"foo": "foo"}, config, durability=durability) == { + "foo": "hi! foo", + "__interrupt__": [ + Interrupt( + value="Provide baz value", + id=AnyStr(), + ) + ], + } + assert (await graph.aget_state(config, subgraphs=True)).tasks[0].state.values == { + "bar": "hi! foo" + } + assert await graph.ainvoke( + Command(resume="baz"), config, durability=durability + ) == {"foo": "hi! foobaz"} + + +async def test_stream_subgraphs_during_execution( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + async def inner_1(state: InnerState): + return {"my_key": "got here", "my_other_key": state["my_key"]} + + async def inner_2(state: InnerState): + await asyncio.sleep(0.5) + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + async def outer_1(state: State): + await asyncio.sleep(0.2) + return {"my_key": " and parallel"} + + async def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile()) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.add_edge("outer_2", END) + + app = graph.compile(checkpointer=async_checkpointer) + + start = perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + async for c in app.astream({"my_key": ""}, config, subgraphs=True): + chunks.append((round(perf_counter() - start, 1), c)) + for idx in range(len(chunks)): + elapsed, c = chunks[idx] + chunks[idx] = (round(elapsed - chunks[0][0], 1), c) + + assert chunks == [ + # arrives before "inner" finishes + ( + FloatBetween(0.0, 0.1), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ), + (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), + ( + FloatBetween(0.5, 0.8), + ( + (AnyStr("inner:"),), + {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, + ), + ), + (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), + (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), + ] + + +@NEEDS_CONTEXTVARS +async def test_stream_buffering_single_node( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + + async def node(state: State, writer: StreamWriter): + writer("Before sleep") + await asyncio.sleep(0.2) + writer("After sleep") + return {"my_key": "got here"} + + builder = StateGraph(State) + builder.add_node("node", node) + builder.add_edge(START, "node") + builder.add_edge("node", END) + + graph = builder.compile(checkpointer=async_checkpointer) + + start = perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + async for c in graph.astream({"my_key": ""}, config, stream_mode="custom"): + chunks.append((round(perf_counter() - start, 1), c)) + + assert chunks == [ + (FloatBetween(0.0, 0.1), "Before sleep"), + (FloatBetween(0.2, 0.3), "After sleep"), + ] + + +async def test_nested_graph_interrupts_parallel( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + async def inner_1(state: InnerState): + await asyncio.sleep(0.1) + return {"my_key": "got here", "my_other_key": state["my_key"]} + + async def inner_2(state: InnerState): + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + async def outer_1(state: State): + return {"my_key": " and parallel"} + + async def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node( + "inner", + inner.compile(interrupt_before=["inner_2"]), + ) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=async_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": ""}, config, durability=durability) == { + "my_key": " and parallel", + } + + assert await app.ainvoke(None, config, durability=durability) == { + "my_key": "got here and there and parallel and back again", + } + + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + durability=durability, + ) + ] == [ + # we got to parallel node first + ((), {"outer_1": {"my_key": " and parallel"}}), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ((), {"__interrupt__": ()}), + ] + assert [c async for c in app.astream(None, config, durability=durability)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + # # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=async_checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": ""}, + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=async_checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + +async def test_doubly_nested_graph_interrupts( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + async def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + async def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + async def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + async def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=async_checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, durability=durability) == { + "my_key": "hi my value", + } + + assert await app.ainvoke(None, config, durability=durability) == { + "my_key": "hi my value here and there and back again", + } + + # test stream updates w/ nested interrupt + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, durability=durability + ) + ] == [ + {"parent_1": {"my_key": "hi my value"}}, + {"__interrupt__": ()}, + ] + assert nodes == ["parent_1", "grandchild_1"] + assert [c async for c in app.astream(None, config, durability=durability)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] + assert nodes == [ + "parent_1", + "grandchild_1", + "grandchild_2", + "child_1", + "child", + "parent_2", + ] + + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, + config, + stream_mode="values", + durability=durability, + ) + ] == [ + {"my_key": "my value"}, + {"my_key": "hi my value"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", durability=durability + ) + ] == [ + {"my_key": "hi my value"}, + {"my_key": "hi my value here and there"}, + {"my_key": "hi my value here and there and back again"}, + ] + + +async def test_checkpoint_metadata(async_checkpointer: BaseCheckpointSaver) -> None: + """This test verifies that a run's configurable fields are merged with the + previous checkpoint config for each step in the run. + """ + # set up test + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, AnyMessage + from langchain_core.prompts import ChatPromptTemplate + from langchain_core.tools import tool + + # graph state + class BaseState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + # initialize graph nodes + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You are a nice assistant."), + ("placeholder", "{messages}"), + ] + ) + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage(content="answer"), + ] + ) + + def agent(state: BaseState, config: RunnableConfig) -> BaseState: + formatted = prompt.invoke(state) + response = model.invoke(formatted) + return {"messages": response} + + def should_continue(data: BaseState) -> str: + # Logic to decide whether to continue in the loop or exit + if not data["messages"][-1].tool_calls: + return "exit" + else: + return "continue" + + # define graphs w/ and w/o interrupt + workflow = StateGraph(BaseState) + workflow.add_node("agent", agent) + workflow.add_node("tools", ToolNode(tools)) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", should_continue, {"continue": "tools", "exit": END} + ) + workflow.add_edge("tools", "agent") + + # graph w/o interrupt + app = workflow.compile(checkpointer=async_checkpointer) + + # graph w/ interrupt + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, interrupt_before=["tools"] + ) + + # assertions + + # invoke graph w/o interrupt + await app.ainvoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "1", + "test_config_1": "foo", + "test_config_2": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "1"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_1 = (await async_checkpointer.aget_tuple(config)).metadata + assert chkpnt_metadata_1["test_config_1"] == "foo" + assert chkpnt_metadata_1["test_config_2"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_1 = async_checkpointer.alist(config) + async for chkpnt_tuple in chkpnt_tuples_1: + assert chkpnt_tuple.metadata["test_config_1"] == "foo" + assert chkpnt_tuple.metadata["test_config_2"] == "bar" + + # invoke graph, but interrupt before tool call + await app_w_interrupt.ainvoke( + {"messages": ["what is weather in sf"]}, + { + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + }, + }, + ) + + config = {"configurable": {"thread_id": "2"}} + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_2 = (await async_checkpointer.aget_tuple(config)).metadata + assert chkpnt_metadata_2["test_config_3"] == "foo" + assert chkpnt_metadata_2["test_config_4"] == "bar" + + # resume graph execution + await app_w_interrupt.ainvoke( + input=None, + config={ + "configurable": { + "thread_id": "2", + "test_config_3": "foo", + "test_config_4": "bar", + } + }, + ) + + # assert that checkpoint metadata contains the run's configurable fields + chkpnt_metadata_3 = (await async_checkpointer.aget_tuple(config)).metadata + assert chkpnt_metadata_3["test_config_3"] == "foo" + assert chkpnt_metadata_3["test_config_4"] == "bar" + + # Verify that all checkpoint metadata have the expected keys. This check + # is needed because a run may have an arbitrary number of steps depending + # on how the graph is constructed. + chkpnt_tuples_2 = async_checkpointer.alist(config) + async for chkpnt_tuple in chkpnt_tuples_2: + assert chkpnt_tuple.metadata["test_config_3"] == "foo" + assert chkpnt_tuple.metadata["test_config_4"] == "bar" + + +async def test_checkpointer_null_pending_writes() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return [self.name] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_edge(START, "1") + graph = builder.compile(checkpointer=MemorySaverNoPending()) + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 3 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 4 + + +async def test_store_injected_async( + async_checkpointer: BaseCheckpointSaver, async_store: BaseStore +) -> None: + class State(TypedDict): + count: Annotated[int, operator.add] + + doc_id = str(uuid.uuid4()) + doc = {"some-key": "this-is-a-val"} + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") + thread_1 = str(uuid.uuid4()) + thread_2 = str(uuid.uuid4()) + + class Node: + def __init__(self, i: int | None = None): + self.i = i + + async def __call__( + self, inputs: State, config: RunnableConfig, store: BaseStore + ): + assert isinstance(store, BaseStore) + await store.aput( + ( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar") + ), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} + + def other_node(inputs: State, config: RunnableConfig, store: BaseStore): + assert isinstance(store, BaseStore) + store.put(("not", "interesting"), "key", {"val": "val"}) + item = store.get(("not", "interesting"), "key") + assert item is not None + assert item.value == {"val": "val"} + return {"count": 0} + + builder = StateGraph(State) + builder.add_node("node", Node()) + builder.add_node("other_node", other_node) + builder.add_edge("__start__", "node") + builder.add_edge("node", "other_node") + + N = 50 + M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + + graph = builder.compile(store=async_store, checkpointer=async_checkpointer) + + # Test batch operations with multiple threads + results = await graph.abatch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], + ) + result = results[-1] + assert result == {"count": N + 1} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} + assert len(await async_store.asearch(namespace)) == 1 + + # Check results after another turn of the same thread + result = await graph.ainvoke( + {"count": 0}, {"configurable": {"thread_id": thread_1}} + ) + assert result == {"count": (N + 1) * 2} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len(await async_store.asearch(namespace)) == 1 + + # Test with a different thread + result = await graph.ainvoke( + {"count": 0}, {"configurable": {"thread_id": thread_2}} + ) + assert result == {"count": N + 1} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == { + **doc, + "from_thread": thread_2, + "some_val": 0, + } # Overwrites the whole doc + assert ( + len(await async_store.asearch(namespace)) == 1 + ) # still overwriting the same one + + +async def test_debug_retry(async_checkpointer: BaseCheckpointSaver): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + async def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + builder = StateGraph(State) + builder.add_node("one", node("one")) + builder.add_node("two", node("two")) + builder.add_edge(START, "one") + builder.add_edge("one", "two") + builder.add_edge("two", END) + + graph = builder.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + await graph.ainvoke({"messages": []}, config=config, durability="async") + + # re-run step: 1 + async for c in async_checkpointer.alist(config): + if c.metadata["step"] == 1: + target_config = c.parent_config + break + assert target_config is not None + + update_config = await graph.aupdate_state(target_config, values=None) + + events = [ + c + async for c in graph.astream( + None, config=update_config, stream_mode="debug", durability="async" + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + + checkpoint_history = { + c.config["configurable"]["checkpoint_id"]: c + async for c in graph.aget_state_history(config) + } + + def lax_normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + return config["configurable"] + + for stream in checkpoint_events: + stream_conf = lax_normalize_config(stream["config"]) + stream_parent_conf = lax_normalize_config(stream["parent_config"]) + assert stream_conf != stream_parent_conf + + # ensure the streamed checkpoint == checkpoint from checkpointer.list() + history = checkpoint_history[stream["config"]["configurable"]["checkpoint_id"]] + history_conf = lax_normalize_config(history.config) + assert stream_conf == history_conf + + history_parent_conf = lax_normalize_config(history.parent_config) + assert stream_parent_conf == history_parent_conf + + +async def test_debug_subgraphs( + async_checkpointer: BaseCheckpointSaver, durability: Durability +): + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + async def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + graph = parent.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + events = [ + c + async for c in graph.astream( + {"messages": []}, + config=config, + stream_mode="debug", + durability=durability, + ) + ] + + checkpoint_events = list( + reversed([e["payload"] for e in events if e["type"] == "checkpoint"]) + ) + if durability == "exit": + checkpoint_events = checkpoint_events[:1] + checkpoint_history = [c async for c in graph.aget_state_history(config)] + + assert len(checkpoint_events) == len(checkpoint_history) + + def normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + return config["configurable"] + + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config(history.config) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +async def test_debug_nested_subgraphs( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + from collections import defaultdict + + class State(TypedDict): + messages: Annotated[list[str], operator.add] + + def node(name): + async def _node(state: State): + return {"messages": [f"entered {name} node"]} + + return _node + + grand_parent = StateGraph(State) + parent = StateGraph(State) + child = StateGraph(State) + + child.add_node("c_one", node("c_one")) + child.add_node("c_two", node("c_two")) + child.add_edge(START, "c_one") + child.add_edge("c_one", "c_two") + child.add_edge("c_two", END) + + parent.add_node("p_one", node("p_one")) + parent.add_node("p_two", child.compile()) + parent.add_edge(START, "p_one") + parent.add_edge("p_one", "p_two") + parent.add_edge("p_two", END) + + grand_parent.add_node("gp_one", node("gp_one")) + grand_parent.add_node("gp_two", parent.compile()) + grand_parent.add_edge(START, "gp_one") + grand_parent.add_edge("gp_one", "gp_two") + grand_parent.add_edge("gp_two", END) + + graph = grand_parent.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + events = [ + c + async for c in graph.astream( + {"messages": []}, + config=config, + stream_mode="debug", + subgraphs=True, + durability=durability, + ) + ] + + stream_ns: dict[tuple, dict] = defaultdict(list) + for ns, e in events: + if e["type"] == "checkpoint": + stream_ns[ns].append(e["payload"]) + + assert list(stream_ns.keys()) == [ + (), + (AnyStr("gp_two:"),), + (AnyStr("gp_two:"), AnyStr("p_two:")), + ] + + history_ns = {} + for ns in stream_ns.keys(): + + async def get_history(): + history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": "1", "checkpoint_ns": "|".join(ns)}} + ) + ] + return history[::-1] + + history_ns[ns] = await get_history() + + def normalize_config(config: dict | None) -> dict | None: + if config is None: + return None + + clean_config = {} + clean_config["thread_id"] = config["configurable"]["thread_id"] + clean_config["checkpoint_id"] = config["configurable"]["checkpoint_id"] + clean_config["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] + if "checkpoint_map" in config["configurable"]: + clean_config["checkpoint_map"] = config["configurable"]["checkpoint_map"] + + return clean_config + + for checkpoint_events, checkpoint_history, ns in zip( + stream_ns.values(), history_ns.values(), stream_ns.keys() + ): + if durability == "exit": + checkpoint_events = checkpoint_events[-1:] + if ns: # Save no checkpoints for subgraphs when durability="exit" + assert not checkpoint_history + continue + assert len(checkpoint_events) == len(checkpoint_history) + for stream, history in zip(checkpoint_events, checkpoint_history): + assert stream["values"] == history.values + assert stream["next"] == list(history.next) + assert normalize_config(stream["config"]) == normalize_config( + history.config + ) + assert normalize_config(stream["parent_config"]) == normalize_config( + history.parent_config + ) + + assert len(stream["tasks"]) == len(history.tasks) + for stream_task, history_task in zip(stream["tasks"], history.tasks): + assert stream_task["id"] == history_task.id + assert stream_task["name"] == history_task.name + assert stream_task["interrupts"] == history_task.interrupts + assert stream_task.get("error") == history_task.error + assert stream_task.get("state") == history_task.state + + +@pytest.mark.parametrize("subgraph_persist", [True, False]) +async def test_parent_command( + async_checkpointer: BaseCheckpointSaver, subgraph_persist: bool +) -> None: + from langchain_core.messages import BaseMessage + from langchain_core.tools import tool + + @tool(return_direct=True) + def get_user_name() -> Command: + """Retrieve user name""" + return Command(update={"user_name": "Meow"}, graph=Command.PARENT) + + subgraph_builder = StateGraph(MessagesState) + subgraph_builder.add_node("tool", get_user_name) + subgraph_builder.add_edge(START, "tool") + subgraph = subgraph_builder.compile(checkpointer=subgraph_persist) + + class CustomParentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + # this key is not available to the child graph + user_name: str + + builder = StateGraph(CustomParentState) + builder.add_node("alice", subgraph) + builder.add_edge(START, "alice") + + graph = builder.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke( + {"messages": [("user", "get user name")]}, config, durability="exit" + ) == { + "messages": [ + _AnyIdHumanMessage( + content="get user name", additional_kwargs={}, response_metadata={} + ), + ], + "user_name": "Meow", + } + assert await graph.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage( + content="get user name", + additional_kwargs={}, + response_metadata={}, + ), + ], + "user_name": "Meow", + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(), + interrupts=(), + ) + + +@NEEDS_CONTEXTVARS +async def test_interrupt_subgraph(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + baz: str + + def foo(state): + return {"baz": "foo"} + + def bar(state): + value = interrupt("Please provide baz value:") + return {"baz": value} + + child_builder = StateGraph(State) + child_builder.add_node(bar) + child_builder.add_edge(START, "bar") + + builder = StateGraph(State) + builder.add_node(foo) + builder.add_node("bar", child_builder.compile()) + builder.add_edge(START, "foo") + builder.add_edge("foo", "bar") + + graph = builder.compile(checkpointer=async_checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert await graph.ainvoke({"baz": ""}, thread1) + # Resume with answer + assert await graph.ainvoke(Command(resume="bar"), thread1) + + +@NEEDS_CONTEXTVARS +async def test_interrupt_multiple(async_checkpointer: BaseCheckpointSaver): + class State(TypedDict): + my_key: Annotated[str, operator.add] + + async def node(s: State) -> State: + answer = interrupt({"value": 1}) + answer2 = interrupt({"value": 2}) + return {"my_key": answer + " " + answer2} + + builder = StateGraph(State) + builder.add_node("node", node) + builder.add_edge(START, "node") + + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + + assert [ + e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1) + ] == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 1}, + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + async for event in graph.astream( + Command(resume="answer 1", update={"my_key": "foofoo"}), + thread1, + stream_mode="updates", + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 2}, + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + async for event in graph.astream( + Command(resume="answer 2"), thread1, stream_mode="updates" + ) + ] == [ + {"node": {"my_key": "answer 1 answer 2"}}, + ] + + +@NEEDS_CONTEXTVARS +async def test_interrupt_loop(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + age: int + other: str + + async def ask_age(s: State): + """Ask an expert for help.""" + question = "How old are you?" + value = None + for _ in range(10): + value: str = interrupt(question) + if not value.isdigit() or int(value) < 18: + question = "invalid response" + value = None + else: + break + + return {"age": int(value)} + + builder = StateGraph(State) + builder.add_node("node", ask_age) + builder.add_edge(START, "node") + + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + + assert [e async for e in graph.astream({"other": ""}, thread1)] == [ + { + "__interrupt__": ( + Interrupt( + value="How old are you?", + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + async for event in graph.astream( + Command(resume="13"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + id=AnyStr(), + ), + ) + } + ] + + assert [ + event + async for event in graph.astream( + Command(resume="15"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + id=AnyStr(), + ), + ) + } + ] + + assert [event async for event in graph.astream(Command(resume="19"), thread1)] == [ + {"node": {"age": 19}}, + ] + + +@NEEDS_CONTEXTVARS +async def test_interrupt_functional(async_checkpointer: BaseCheckpointSaver) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + @entrypoint(checkpointer=async_checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**foo_result, "b": value} + bar_result = await bar(bar_input) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@NEEDS_CONTEXTVARS +async def test_interrupt_task_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + @entrypoint(checkpointer=async_checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + bar_result = await bar(foo_result) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert await graph.ainvoke({"a": ""}, config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + id=AnyStr(), + ), + ] + } + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + +async def test_command_with_static_breakpoints( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that we can use Command to resume and update with static breakpoints.""" + + class State(TypedDict): + """The graph state.""" + + foo: str + + def node1(state: State): + return { + "foo": state["foo"] + "|node-1", + } + + def node2(state: State): + return { + "foo": state["foo"] + "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["node1"]) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(update={"foo": "def"}), config) + assert result == {"foo": "def|node-1|node-2"} + + +async def test_multistep_plan(async_checkpointer: BaseCheckpointSaver) -> None: + from langchain_core.messages import AnyMessage + + class State(TypedDict, total=False): + plan: list[str | list[str]] + messages: Annotated[list[AnyMessage], add_messages] + + def planner(state: State): + if state.get("plan") is None: + # create plan somehow + plan = ["step1", ["step2", "step3"], "step4"] + # pick the first step to execute next + first_step, *plan = plan + # put the rest of plan in state + return Command(goto=first_step, update={"plan": plan}) + elif state["plan"]: + # go to the next step of the plan + next_step, *next_plan = state["plan"] + return Command(goto=next_step, update={"plan": next_plan}) + else: + # the end of the plan + pass + + def step1(state: State): + return Command(goto="planner", update={"messages": [("human", "step1")]}) + + def step2(state: State): + return Command(goto="planner", update={"messages": [("human", "step2")]}) + + def step3(state: State): + return Command(goto="planner", update={"messages": [("human", "step3")]}) + + def step4(state: State): + return Command(goto="planner", update={"messages": [("human", "step4")]}) + + builder = StateGraph(State) + builder.add_node(planner) + builder.add_node(step1) + builder.add_node(step2) + builder.add_node(step3) + builder.add_node(step4) + builder.add_edge(START, "planner") + + graph = builder.compile(checkpointer=async_checkpointer) + + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke({"messages": [("human", "start")]}, config) == { + "messages": [ + _AnyIdHumanMessage(content="start"), + _AnyIdHumanMessage(content="step1"), + _AnyIdHumanMessage(content="step2"), + _AnyIdHumanMessage(content="step3"), + _AnyIdHumanMessage(content="step4"), + ], + "plan": [], + } + + +async def test_command_goto_with_static_breakpoints( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Use Command goto with static breakpoints.""" + + class State(TypedDict): + """The graph state.""" + + foo: Annotated[str, operator.add] + + def node1(state: State): + return { + "foo": "|node-1", + } + + def node2(state: State): + return { + "foo": "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["node1"]) + + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(goto=["node2"]), config) + assert result == {"foo": "abc|node-1|node-2|node-2"} + + +async def test_parallel_node_execution(): + """Test that parallel nodes execute concurrently.""" + + class State(TypedDict): + results: Annotated[list[str], operator.add] + + async def slow_node(state: State): + await asyncio.sleep(1) + return {"results": ["slow"]} + + async def fast_node(state: State): + await asyncio.sleep(2) + return {"results": ["fast"]} + + builder = StateGraph(State) + builder.add_node("slow", slow_node) + builder.add_node("fast", fast_node) + builder.add_edge(START, "slow") + builder.add_edge(START, "fast") + + graph = builder.compile() + + start = perf_counter() + result = await graph.ainvoke({"results": []}) + duration = perf_counter() - start + + # Fast node result should be available first + assert "fast" in result["results"][0] + + # Total duration should be less than sum of both nodes + assert duration < 3.0 + + +@NEEDS_CONTEXTVARS +async def test_multiple_interrupt_state_persistence( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that state is preserved correctly across multiple interrupts.""" + + class State(TypedDict): + steps: Annotated[list[str], operator.add] + + def interruptible_node(state: State): + first = interrupt("First interrupt") + second = interrupt("Second interrupt") + return {"steps": [first, second]} + + builder = StateGraph(State) + builder.add_node("node", interruptible_node) + builder.add_edge(START, "node") + + app = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First execution - should hit first interrupt + await app.ainvoke({"steps": []}, config) + + # State should still be empty since node hasn't returned + state = await app.aget_state(config) + assert state.values == {"steps": []} + + # Resume after first interrupt - should hit second interrupt + await app.ainvoke(Command(resume="step1"), config) + + # State should still be empty since node hasn't returned + state = await app.aget_state(config) + assert state.values == {"steps": []} + + # Resume after second interrupt - node should complete + result = await app.ainvoke(Command(resume="step2"), config) + + # Now state should contain both steps since node returned + assert result["steps"] == ["step1", "step2"] + state = await app.aget_state(config) + assert state.values["steps"] == ["step1", "step2"] + + +async def test_concurrent_execution(): + """Test concurrent execution with async nodes.""" + + class State(TypedDict): + counter: Annotated[int, operator.add] + + results = deque() + + async def slow_node(state: State): + await asyncio.sleep(0.1) + return {"counter": 1} + + builder = StateGraph(State) + builder.add_node("node", slow_node) + builder.add_edge(START, "node") + graph = builder.compile() + + async def run_graph(): + result = await graph.ainvoke({"counter": 0}) + results.append(result) + + # Create and gather tasks + tasks = [run_graph() for _ in range(10)] + await asyncio.gather(*tasks) + + # Verify results are independent + assert len(results) == 10 + for result in results: + assert result["counter"] == 1 + + +async def test_checkpoint_recovery_async( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + """Test recovery from checkpoints after failures with async nodes.""" + + class State(TypedDict): + steps: Annotated[list[str], operator.add] + attempt: int # Track number of attempts + + async def failing_node(state: State): + # Fail on first attempt, succeed on retry + if state["attempt"] == 1: + raise RuntimeError("Simulated failure") + await asyncio.sleep(0.1) # Simulate async work + return {"steps": ["node1"]} + + async def second_node(state: State): + await asyncio.sleep(0.1) # Simulate async work + return {"steps": ["node2"]} + + builder = StateGraph(State) + builder.add_node("node1", failing_node) + builder.add_node("node2", second_node) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First attempt should fail + with pytest.raises(RuntimeError): + await graph.ainvoke( + {"steps": ["start"], "attempt": 1}, + config, + durability=durability, + ) + + # Verify checkpoint state + state = await graph.aget_state(config) + assert state is not None + assert state.values == {"steps": ["start"], "attempt": 1} # input state saved + assert state.next == ("node1",) # Should retry failed node + + # Retry with updated attempt count + result = await graph.ainvoke( + {"steps": [], "attempt": 2}, config, durability=durability + ) + assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} + + # Verify checkpoint history shows both attempts + history = [c async for c in graph.aget_state_history(config)] + if durability != "exit": + assert len(history) == 6 # Initial + failed attempt + successful attempt + else: + assert len(history) == 2 # error + success + + # Verify the error was recorded in checkpoint + failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) + assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error + + +async def test_multiple_updates_root() -> None: + def node_a(state): + return [Command(update="a1"), Command(update="a2")] + + def node_b(state): + return "b" + + graph = ( + StateGraph(Annotated[str, operator.add]) + .add_sequence([node_a, node_b]) + .add_edge(START, "node_a") + .compile() + ) + + assert await graph.ainvoke("") == "a1a2b" + + # only streams the last update from node_a + assert [c async for c in graph.astream("", stream_mode="updates")] == [ + {"node_a": ["a1", "a2"]}, + {"node_b": "b"}, + ] + + +async def test_multiple_updates() -> None: + class State(TypedDict): + foo: Annotated[str, operator.add] + + def node_a(state): + return [Command(update={"foo": "a1"}), Command(update={"foo": "a2"})] + + def node_b(state): + return {"foo": "b"} + + graph = ( + StateGraph(State) + .add_sequence([node_a, node_b]) + .add_edge(START, "node_a") + .compile() + ) + + assert await graph.ainvoke({"foo": ""}) == { + "foo": "a1a2b", + } + + # only streams the last update from node_a + assert [c async for c in graph.astream({"foo": ""}, stream_mode="updates")] == [ + {"node_a": [{"foo": "a1"}, {"foo": "a2"}]}, + {"node_b": {"foo": "b"}}, + ] + + +@NEEDS_CONTEXTVARS +async def test_falsy_return_from_task(async_checkpointer: BaseCheckpointSaver) -> None: + """Test with a falsy return from a task.""" + + @task + async def falsy_task() -> bool: + return False + + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + await falsy_task() + interrupt("test") + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({"a": 5}, configurable) + await graph.ainvoke(Command(resume="123"), configurable) + + +@NEEDS_CONTEXTVARS +async def test_multiple_interrupts_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test multiple interrupts with functional API.""" + from langgraph.func import entrypoint, task + + counter = 0 + + @task + async def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 2, 3]: + values.extend([await double(idx), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + result = await graph.ainvoke(Command(resume="c"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 4, "b", 6, "c"], + } + assert counter == 3 + + +@NEEDS_CONTEXTVARS +async def test_multiple_interrupts_functional_cache( + async_checkpointer: BaseCheckpointSaver, cache: BaseCache +): + """Test multiple interrupts with functional API.""" + counter = 0 + + @task(cache_policy=CachePolicy()) + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=async_checkpointer, cache=cache) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 1, 2, 2, 3, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 + + # clear the cache + await double.aclear_cache(cache) + + # now should recompute + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({}, configurable) + await graph.ainvoke(Command(resume="a"), configurable) + await graph.ainvoke(Command(resume="b"), configurable) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 6 + + +@NEEDS_CONTEXTVARS +async def test_double_interrupt_subgraph( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class AgentState(TypedDict): + input: str + + def node_1(state: AgentState): + result = interrupt("interrupt node 1") + return {"input": result} + + def node_2(state: AgentState): + result = interrupt("interrupt node 2") + return {"input": result} + + subgraph_builder = ( + StateGraph(AgentState) + .add_node("node_1", node_1) + .add_node("node_2", node_2) + .add_edge(START, "node_1") + .add_edge("node_1", "node_2") + .add_edge("node_2", END) + ) + + # invoke the sub graph + subgraph = subgraph_builder.compile(checkpointer=async_checkpointer) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + id=AnyStr(), + ), + ) + }, + ] + # resume from the first interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_1": {"input": "123"}, + }, + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + id=AnyStr(), + ), + ) + }, + ] + # resume from the second interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_2": {"input": "123"}, + }, + ] + + subgraph = subgraph_builder.compile() + + def invoke_sub_agent(state: AgentState): + return subgraph.invoke(state) + + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + + parent_agent = ( + StateGraph(AgentState) + .add_node("invoke_sub_agent", invoke_sub_agent) + .add_edge(START, "invoke_sub_agent") + .add_edge("invoke_sub_agent", END) + .compile(checkpointer=async_checkpointer) + ) + + assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + id=AnyStr(), + ), + ) + }, + ] + + # resume from the first interrupt + assert [c async for c in parent_agent.astream(Command(resume=True), thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + id=AnyStr(), + ), + ) + } + ] + + # resume from 2nd interrupt + assert [c async for c in parent_agent.astream(Command(resume=True), thread)] == [ + { + "invoke_sub_agent": {"input": True}, + }, + ] + + +@NEEDS_CONTEXTVARS +async def test_async_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + async def slow() -> dict: + await asyncio.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": asyncio.get_running_loop().time()} + + @entrypoint() + async def graph(inputs: dict) -> list: + first = await slow() + second = await slow() + return [first, second] + + arrival_times = [] + + async for chunk in graph.astream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(asyncio.get_running_loop().time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay + + +@NEEDS_CONTEXTVARS +async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + async def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + async def call_same_subgraph(state): + result = await add_subgraph.ainvoke(state) + another_result = await add_subgraph.ainvoke({"a": result["result"], "b": 10}) + return another_result + + parent_call_same_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { + "result": 15 + } + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + async def call_multiple_subgraphs(state): + add_result = await add_subgraph.ainvoke(state) + multiply_result = await multiply_subgraph.ainvoke(state) + return { + "add_result": add_result["result"], + "multiply_result": multiply_result["result"], + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output_schema=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } + + +@NEEDS_CONTEXTVARS +async def test_multiple_subgraphs_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: + # Define addition subgraph + @entrypoint() + async def add(inputs): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + async def multiply_task(a, b): + return a * b + + @entrypoint() + async def multiply(inputs): + return await multiply_task(*inputs) + + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = await add.ainvoke([a, b]) + another_result = await add.ainvoke([result, 10]) + return another_result + + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = await add.ainvoke([a, b]) + multiply_result = await multiply.ainvoke([a, b]) + return [add_result, multiply_result] + + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) + + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + + +@NEEDS_CONTEXTVARS +async def test_multiple_subgraphs_mixed_entrypoint( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test calling multiple StateGraph subgraphs from an entrypoint.""" + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} + + add_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(add) + .add_edge(START, "add") + .compile() + ) + + async def multiply(state): + return {"result": state["a"] * state["b"]} + + multiply_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) + + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))["result"] + return another_result + + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))["result"] + return [add_result, multiply_result] + + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) + + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + + +@NEEDS_CONTEXTVARS +async def test_multiple_subgraphs_mixed_state_graph( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" + + class State(TypedDict): + a: int + b: int + + class Output(TypedDict): + result: int + + # Define addition subgraph + @entrypoint() + async def add(inputs): + a, b = inputs + return a + b + + # Define multiplication subgraph using tasks + @task + async def multiply_task(a, b): + return a * b + + @entrypoint() + async def multiply(inputs): + return await multiply_task(*inputs) + + # Test calling the same subgraph multiple times + async def call_same_subgraph(state): + result = await add.ainvoke([state["a"], state["b"]]) + another_result = await add.ainvoke([result, 10]) + return {"result": another_result} + + parent_call_same_subgraph = ( + StateGraph(State, output_schema=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { + "result": 15 + } + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + async def call_multiple_subgraphs(state): + add_result = await add.ainvoke([state["a"], state["b"]]) + multiply_result = await multiply.ainvoke([state["a"], state["b"]]) + return { + "add_result": add_result, + "multiply_result": multiply_result, + } + + parent_call_multiple_subgraphs = ( + StateGraph(State, output_schema=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } + + +@NEEDS_CONTEXTVARS +async def test_multiple_subgraphs_checkpointer( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + sub_counter: Annotated[int, operator.add] + + async def subgraph_node(state): + return {"sub_counter": 2} + + sub_graph_1 = ( + StateGraph(SubgraphState) + .add_node(subgraph_node) + .add_edge(START, "subgraph_node") + .compile(checkpointer=True) + ) + + class OtherSubgraphState(TypedDict): + other_sub_counter: Annotated[int, operator.add] + + async def other_subgraph_node(state): + return {"other_sub_counter": 3} + + sub_graph_2 = ( + StateGraph(OtherSubgraphState) + .add_node(other_subgraph_node) + .add_edge(START, "other_subgraph_node") + .compile() + ) + + class ParentState(TypedDict): + parent_counter: int + + async def parent_node(state): + result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]}) + other_result = await sub_graph_2.ainvoke( + {"other_sub_counter": result["sub_counter"]} + ) + return {"parent_counter": other_result["other_sub_counter"]} + + parent_graph = ( + StateGraph(ParentState) + .add_node(parent_node) + .add_edge(START, "parent_node") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 5 + } + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 7 + } + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 5}}), + ] + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" + ) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 7}}), + ] + + +@NEEDS_CONTEXTVARS +async def test_async_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + async def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} + + @entrypoint() + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + + +def test_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert foo.invoke({"a": "1"}, config) == {"a": "1"} + + @entrypoint() + def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + + +async def test_entrypoint_stateful(async_checkpointer: BaseCheckpointSaver) -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=async_checkpointer) + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + assert await foo.ainvoke({"a": "1"}, config) == { + "current": {"a": "1"}, + "previous": None, + } + assert await foo.ainvoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + } + assert await foo.ainvoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + }, + } + assert states == [ + None, + {"current": {"a": "1"}, "previous": None}, + {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}}, + ] + + # Test stream + @entrypoint(checkpointer=async_checkpointer) + async def foo(inputs, *, previous: Any) -> Any: + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "2"}} + items = [item async for item in foo.astream({"a": "1"}, config)] + assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}] + + +async def test_entrypoint_stateful_update_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=async_checkpointer) + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + # assert print(foo.input_channels) + await foo.aupdate_state(config, {"a": "-1"}) + assert await foo.ainvoke({"a": "1"}, config) == { + "current": {"a": "1"}, + "previous": {"a": "-1"}, + } + assert await foo.ainvoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + } + assert await foo.ainvoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + }, + } + + # update state + await foo.aupdate_state(config, {"a": "3"}) + + # Test stream + assert [item async for item in foo.astream({"a": "1"}, config)] == [ + {"foo": {"current": {"a": "1"}, "previous": {"a": "3"}}} + ] + assert states == [ + {"a": "-1"}, + {"current": {"a": "1"}, "previous": {"a": "-1"}}, + { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": {"a": "-1"}}, + }, + {"a": "3"}, + ] + + +async def test_entrypoint_from_async_generator() -> None: + """@entrypoint does not support async generators.""" + with pytest.raises(NotImplementedError): + + @entrypoint() + async def foo(inputs) -> Any: + yield "a" + yield "b" + + +@NEEDS_CONTEXTVARS +async def test_named_tasks_functional() -> None: + class Foo: + async def foo(self, value: str) -> dict: + return value + "foo" + + f = Foo() + + # class method task + foo = task(f.foo, name="custom_foo") + other_foo = task(f.foo, name="other_foo") + + # regular function task + @task(name="custom_bar") + async def bar(value: str) -> dict: + return value + "|bar" + + async def baz(update: str, value: str) -> dict: + return value + f"|{update}" + + # partial function task (unnamed) + baz_task = task(functools.partial(baz, "baz")) + # partial function task (named_) + custom_baz_task = task(functools.partial(baz, "custom_baz"), name="custom_baz") + + class Qux: + def __call__(self, value: str) -> dict: + return value + "|qux" + + qux_task = task(Qux(), name="qux") + + @entrypoint() + async def workflow(inputs: dict) -> dict: + foo_result = await foo(inputs) + await other_foo(inputs) + bar_result = await bar(foo_result) + baz_result = await baz_task(bar_result) + custom_baz_result = await custom_baz_task(baz_result) + qux_result = await qux_task(custom_baz_result) + return qux_result + + assert [c async for c in workflow.astream("", stream_mode="updates")] == [ + {"custom_foo": "foo"}, + {"other_foo": "foo"}, + {"custom_bar": "foo|bar"}, + {"baz": "foo|bar|baz"}, + {"custom_baz": "foo|bar|baz|custom_baz"}, + {"qux": "foo|bar|baz|custom_baz|qux"}, + {"workflow": "foo|bar|baz|custom_baz|qux"}, + ] + + +@NEEDS_CONTEXTVARS +async def test_overriding_injectable_args_with_async_task( + async_store: BaseStore, +) -> None: + """Test overriding injectable args in tasks.""" + + @task + async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: + assert store is value + assert writer is value + + @entrypoint(store=async_store) + async def main(inputs, store: BaseStore) -> str: + assert store is not None + await foo(store=None, writer=None, value=None) + await foo(store="hello", writer="hello", value="hello") + return "OK" + + assert await main.ainvoke({}) == "OK" + + +async def test_tags_stream_mode_messages() -> None: + model = GenericFakeChatModel(messages=iter(["foo"]), tags=["meow"]) + + async def call_model(state, config): + return {"messages": await model.ainvoke(state["messages"], config)} + + graph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_edge(START, "call_model") + .compile() + ) + assert [ + c + async for c in graph.astream( + { + "messages": "hi", + }, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdAIMessageChunk(content="foo", chunk_position="last"), + { + "langgraph_step": 1, + "langgraph_node": "call_model", + "langgraph_triggers": ("branch:to:call_model",), + "langgraph_path": ("__pregel_pull", "call_model"), + "langgraph_checkpoint_ns": AnyStr("call_model:"), + "checkpoint_ns": AnyStr("call_model:"), + "ls_provider": "genericfakechatmodel", + "ls_model_type": "chat", + "tags": ["meow"], + }, + ) + ] + + +async def test_stream_mode_messages_command() -> None: + from langchain_core.messages import HumanMessage + + async def my_node(state): + return {"messages": HumanMessage(content="foo")} + + async def my_other_node(state): + return Command(update={"messages": HumanMessage(content="bar")}) + + graph = ( + StateGraph(MessagesState) + .add_sequence([my_node, my_other_node]) + .add_edge(START, "my_node") + .compile() + ) + assert [ + c + async for c in graph.astream( + { + "messages": [], + }, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdHumanMessage(content="foo"), + { + "langgraph_step": 1, + "langgraph_node": "my_node", + "langgraph_triggers": ("branch:to:my_node",), + "langgraph_path": ("__pregel_pull", "my_node"), + "langgraph_checkpoint_ns": AnyStr("my_node:"), + }, + ), + ( + _AnyIdHumanMessage(content="bar"), + { + "langgraph_step": 2, + "langgraph_node": "my_other_node", + "langgraph_triggers": ("branch:to:my_other_node",), + "langgraph_path": ("__pregel_pull", "my_other_node"), + "langgraph_checkpoint_ns": AnyStr("my_other_node:"), + }, + ), + ] + + +async def test_stream_messages_dedupe_inputs() -> None: + from langchain_core.messages import AIMessage + + async def call_model(state): + return {"messages": AIMessage("hi", id="1")} + + async def route(state): + return Command(goto="node_2", graph=Command.PARENT) + + subgraph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_node(route) + .add_edge(START, "call_model") + .add_edge("call_model", "route") + .compile() + ) + + graph = ( + StateGraph(MessagesState) + .add_node("node_1", subgraph) + .add_node("node_2", lambda state: state) + .add_edge(START, "node_1") + .compile() + ) + + chunks = [ + chunk + async for ns, chunk in graph.astream( + {"messages": "hi"}, stream_mode="messages", subgraphs=True + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("hi", id="1") + assert chunks[0][1]["langgraph_node"] == "call_model" + + +async def test_stream_messages_dedupe_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.messages import AIMessage + + to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] + + async def call_model(state): + return {"messages": to_emit.pop(0)} + + async def route(state): + return Command(goto="node_2", graph=Command.PARENT) + + subgraph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_node(route) + .add_edge(START, "call_model") + .add_edge("call_model", "route") + .compile() + ) + + graph = ( + StateGraph(MessagesState) + .add_node("node_1", subgraph) + .add_node("node_2", lambda state: state) + .add_edge(START, "node_1") + .compile(checkpointer=async_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + + chunks = [ + chunk + async for ns, chunk in graph.astream( + {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye", id="1") + assert chunks[0][1]["langgraph_node"] == "call_model" + + chunks = [ + chunk + async for ns, chunk in graph.astream( + {"messages": "hi again"}, + thread1, + stream_mode="messages", + subgraphs=True, + ) + ] + + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye again", id="2") + assert chunks[0][1]["langgraph_node"] == "call_model" + + +@NEEDS_CONTEXTVARS +async def test_interrupt_subgraph_reenter_checkpointer_true( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + foo: str + bar: str + + class ParentState(TypedDict): + foo: str + counter: int + + called = [] + bar_values = [] + + async def subnode_1(state: SubgraphState): + called.append("subnode_1") + bar_values.append(state.get("bar")) + return {"foo": "subgraph_1"} + + async def subnode_2(state: SubgraphState): + called.append("subnode_2") + value = interrupt("Provide value") + value += "baz" + return {"foo": "subgraph_2", "bar": value} + + subgraph = ( + StateGraph(SubgraphState) + .add_node(subnode_1) + .add_node(subnode_2) + .add_edge(START, "subnode_1") + .add_edge("subnode_1", "subnode_2") + .compile(checkpointer=True) + ) + + async def call_subgraph(state: ParentState): + called.append("call_subgraph") + return await subgraph.ainvoke(state) + + async def node(state: ParentState): + called.append("parent") + if state["counter"] < 1: + return Command( + goto="call_subgraph", update={"counter": state["counter"] + 1} + ) + + return {"foo": state["foo"] + "|" + "parent"} + + parent = ( + StateGraph(ParentState) + .add_node(call_subgraph) + .add_node(node) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "node") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + assert await parent.ainvoke({"foo": "", "counter": 0}, config) == { + "foo": "", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + assert await parent.ainvoke(Command(resume="bar"), config) == { + "foo": "subgraph_2", + "counter": 1, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + assert await parent.ainvoke(Command(resume="qux"), config) == { + "foo": "subgraph_2|parent", + "counter": 1, + } + assert called == [ + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + ] + + # invoke parent again (new turn) + assert await parent.ainvoke({"foo": "meow", "counter": 0}, config) == { + "foo": "meow", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + id=AnyStr(), + ) + ], + } + # confirm that we preserve the state values from the previous invocation + assert bar_values == [None, "barbaz", "quxbaz"] + + +@NEEDS_CONTEXTVARS +async def test_handles_multiple_interrupts_from_tasks( + async_checkpointer: BaseCheckpointSaver, +) -> None: + @task + async def add_participant(name: str) -> str: + feedback = interrupt(f"Hey do you want to add {name}?") + + if feedback is False: + return f"The user changed their mind and doesn't want to add {name}!" + + if feedback is True: + return f"Added {name}!" + + raise ValueError("Invalid feedback") + + @entrypoint(checkpointer=async_checkpointer) + async def program(_state: Any) -> list[str]: + first = await add_participant("James") + second = await add_participant("Will") + return [first, second] + + config = {"configurable": {"thread_id": "1"}} + + result = await program.ainvoke("this is ignored", config=config) + assert result == { + "__interrupt__": [ + Interrupt( + value="Hey do you want to add James?", + id=AnyStr(), + ), + ] + } + + state = await program.aget_state(config=config) + assert len(state.tasks[0].interrupts) == 1 + task_interrupt = state.tasks[0].interrupts[0] + assert task_interrupt.value == "Hey do you want to add James?" + + result = await program.ainvoke(Command(resume=True), config=config) + assert result == { + "__interrupt__": [ + Interrupt( + value="Hey do you want to add Will?", + id=AnyStr(), + ), + ] + } + + state = await program.aget_state(config=config) + assert len(state.tasks[0].interrupts) == 1 + task_interrupt = state.tasks[0].interrupts[0] + assert task_interrupt.value == "Hey do you want to add Will?" + + result = await program.ainvoke(Command(resume=True), config=config) + assert result is not None + assert len(result) == 2 + assert result[0] == "Added James!" + assert result[1] == "Added Will!" + + +@NEEDS_CONTEXTVARS +async def test_interrupts_in_tasks_surfaced_once( + async_checkpointer: BaseCheckpointSaver, +) -> None: + @task + async def add_participant(name: str) -> str: + feedback = interrupt(f"Hey do you want to add {name}?") + + if feedback is False: + return f"The user changed their mind and doesn't want to add {name}!" + + if feedback is True: + return f"Added {name}!" + + raise ValueError("Invalid feedback") + + @entrypoint(checkpointer=async_checkpointer) + async def program(_state: Any) -> list[str]: + first = await add_participant("James") + second = await add_participant("Will") + return [first, second] + + config = {"configurable": {"thread_id": "1"}} + + interrupts = [ + e + async for e in program.astream("this is ignored", config=config) + if "__interrupt__" in e + ] + assert len(interrupts) == 1 + + state = await program.aget_state(config=config) + assert len(state.tasks[0].interrupts) == 1 + task_interrupt = state.tasks[0].interrupts[0] + assert task_interrupt.value == "Hey do you want to add James?" + + interrupts = [ + e + async for e in program.astream(Command(resume=True), config=config) + if "__interrupt__" in e + ] + assert len(interrupts) == 1 + + state = await program.aget_state(config=config) + assert len(state.tasks[0].interrupts) == 1 + task_interrupt = state.tasks[0].interrupts[0] + assert task_interrupt.value == "Hey do you want to add Will?" + + result = await program.ainvoke(Command(resume=True), config=config) + assert result is not None + assert len(result) == 2 + assert result[0] == "Added James!" + assert result[1] == "Added Will!" + + +async def test_pregel_loop_refcount(): + gc.collect() + try: + gc.disable() + + class State(TypedDict): + messages: Annotated[list, add_messages] + + graph_builder = StateGraph(State) + + async def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} + + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [ + obj + for obj in gc.get_objects() + if isinstance(obj, AsyncPregelLoop) + ] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable() + + +async def test_bulk_state_updates(async_checkpointer: BaseCheckpointSaver) -> None: + class State(TypedDict): + foo: str + baz: str + + def node_a(state: State) -> State: + return {"foo": "bar"} + + def node_b(state: State) -> State: + return {"baz": "qux"} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + + # First update with node_a + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ] + ], + ) + + # Then bulk update with both nodes + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ] + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + # Check if there are only two checkpoints + checkpoints = [ + c async for c in async_checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + + # perform multiple steps at the same time + config = {"configurable": {"thread_id": "2"}} + + await graph.abulk_update_state( + config, + [ + [ + StateUpdate({"foo": "bar"}, "node_a"), + ], + [ + StateUpdate({"foo": "updated"}, "node_a"), + StateUpdate({"baz": "new"}, "node_b"), + ], + ], + ) + + state = await graph.aget_state(config) + assert state.values == {"foo": "updated", "baz": "new"} + + checkpoints = [ + c async for c in async_checkpointer.alist({"configurable": {"thread_id": "1"}}) + ] + assert len(checkpoints) == 2 + + # Should raise error if updating without as_node + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state( + config, + [ + [ + StateUpdate(values={"foo": "error"}, as_node=None), + StateUpdate(values={"bar": "error"}, as_node=None), + ] + ], + ) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No supersteps provided"): + await graph.abulk_update_state(config, []) + + # Should raise if no updates are provided + with pytest.raises(ValueError, match="No updates provided"): + await graph.abulk_update_state(config, [[], []]) + + # Should raise if __end__ or __copy__ update is applied in bulk + with pytest.raises(InvalidUpdateError): + await graph.abulk_update_state( + config, + [ + [ + StateUpdate(values=None, as_node="__end__"), + StateUpdate(values=None, as_node="__copy__"), + ], + ], + ) + + +async def test_update_as_input( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + foo: str + + def agent(state: State) -> State: + return {"foo": "agent"} + + def tool(state: State) -> State: + return {"foo": "tool"} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("tool", tool) + .add_edge(START, "agent") + .add_edge("agent", "tool") + .compile(checkpointer=async_checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == {"foo": "tool"} + + assert await graph.ainvoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == {"foo": "tool"} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "1"}}) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + # First turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + # Second turn + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent"}, "agent")], + [StateUpdate({"foo": "tool"}, "tool")], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "tool"} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}}) + ] + + if durability != "exit": + assert new_history == history + else: + assert [new_history[0], new_history[4]] == history + + +async def test_batch_update_as_input( + async_checkpointer: BaseCheckpointSaver, durability: Durability +) -> None: + class State(TypedDict): + foo: str + tasks: Annotated[list[int], operator.add] + + def agent(state: State) -> State: + return {"foo": "agent"} + + def map(state: State) -> Command["task"]: + return Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ) + + def task(state: dict) -> State: + return {"tasks": [state["index"]]} + + graph = ( + StateGraph(State) + .add_node("agent", agent) + .add_node("map", map) + .add_node("task", task) + .add_edge(START, "agent") + .add_edge("agent", "map") + .compile(checkpointer=async_checkpointer) + ) + + assert await graph.ainvoke( + {"foo": "input"}, + {"configurable": {"thread_id": "1"}}, + durability=durability, + ) == {"foo": "map", "tasks": [0, 1, 2]} + + def map_snapshot(i: StateSnapshot) -> dict: + return { + "values": i.values, + "next": i.next, + "step": i.metadata.get("step"), + "tasks": [t.name for t in i.tasks], + } + + history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "1"}}) + ] + + await graph.abulk_update_state( + {"configurable": {"thread_id": "2"}}, + [ + [StateUpdate({"foo": "input"}, "__input__")], + [StateUpdate({"foo": "input"}, "__start__")], + [StateUpdate({"foo": "agent", "tasks": []}, "agent")], + [ + StateUpdate( + Command( + goto=[ + Send("task", {"index": 0}), + Send("task", {"index": 1}), + Send("task", {"index": 2}), + ], + update={"foo": "map"}, + ), + "map", + ) + ], + [ + StateUpdate({"tasks": [0]}, "task"), + StateUpdate({"tasks": [1]}, "task"), + StateUpdate({"tasks": [2]}, "task"), + ], + ], + ) + + state = await graph.aget_state({"configurable": {"thread_id": "2"}}) + assert state.values == {"foo": "map", "tasks": [0, 1, 2]} + + new_history = [ + map_snapshot(s) + async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}}) + ] + + if durability != "exit": + assert new_history == history + else: + assert new_history[:1] == history + + +async def test_draw_invalid(): + from langchain_core.messages import BaseMessage + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + workflow = StateGraph(AgentState) + + async def call_model(state: AgentState) -> AgentState: + return state + + async def call_tool(state: AgentState) -> AgentState: + return state + + async def do_nothing(state: AgentState) -> AgentState: + return state + + def should_continue(state): + messages = state["messages"] + last_message = messages[-1] + if last_message.content.startswith("end"): + return END + else: + return [Send("tool", last_message), Send("nothing", last_message)] + + workflow.add_node("agent", call_model) + workflow.add_node("tool", call_tool) + workflow.add_node("nothing", do_nothing) + workflow.set_entry_point("agent") + workflow.add_conditional_edges( + "agent", + should_continue, + path_map=["tool", "nothing", END], + ) + workflow.add_edge("tool", "agent") + + graph = workflow.compile() + + assert graph.get_graph().to_json() == { + "nodes": [ + { + "id": "__start__", + "type": "runnable", + "data": { + "id": ["langgraph", "_internal", "_runnable", "RunnableCallable"], + "name": "__start__", + }, + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "_internal", "_runnable", "RunnableCallable"], + "name": "agent", + }, + }, + { + "id": "tool", + "type": "runnable", + "data": { + "id": ["langgraph", "_internal", "_runnable", "RunnableCallable"], + "name": "tool", + }, + }, + { + "id": "nothing", + "type": "runnable", + "data": { + "id": ["langgraph", "_internal", "_runnable", "RunnableCallable"], + "name": "nothing", + }, + }, + {"id": "__end__"}, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__", "conditional": True}, + {"source": "agent", "target": "nothing", "conditional": True}, + {"source": "agent", "target": "tool", "conditional": True}, + {"source": "tool", "target": "agent"}, + {"source": "nothing", "target": "__end__"}, + ], + } + + +@NEEDS_CONTEXTVARS +async def test_imp_exception( + async_checkpointer: BaseCheckpointSaver, +) -> None: + @task() + async def my_task(number: int): + await asyncio.sleep(0.1) + return number * 2 + + @task() + async def task_with_exception(number: int): + await asyncio.sleep(0.1) + raise Exception("This is a test exception") + + @entrypoint(checkpointer=async_checkpointer) + async def my_workflow(number: int): + await my_task(number) + try: + await task_with_exception(number) + except Exception as e: + print(f"Exception caught: {e}") + await my_task(number) + return "done" + + thread1 = {"configurable": {"thread_id": "1"}} + assert await my_workflow.ainvoke(1, thread1) == "done" + + assert [c async for c in my_workflow.astream(1, thread1)] == [ + {"my_task": 2}, + {"my_task": 2}, + {"my_workflow": "done"}, + ] + + assert [c async for c in my_workflow.astream_events(1, thread1)] == [ + { + "event": "on_chain_start", + "data": {"input": 1}, + "name": "LangGraph", + "tags": [], + "run_id": AnyStr(), + "metadata": {"thread_id": "1"}, + "parent_ids": [], + }, + { + "event": "on_chain_start", + "data": {"input": 1}, + "name": "my_workflow", + "tags": ["graph:step:4"], + "run_id": AnyStr(), + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_workflow", + "langgraph_triggers": ("__start__",), + "langgraph_path": ("__pregel_pull", "my_workflow"), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [AnyStr()], + }, + { + "event": "on_chain_start", + "data": {"input": {"number": 1}}, + "name": "my_task", + "tags": ["seq:step:1"], + "run_id": AnyStr(), + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "my_task", + "tags": ["seq:step:1"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "data": {"chunk": 2}, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_end", + "data": {"output": 2, "input": {"number": 1}}, + "run_id": AnyStr(), + "name": "my_task", + "tags": ["seq:step:1"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "LangGraph", + "tags": [], + "metadata": {"thread_id": "1"}, + "data": {"chunk": {"my_task": 2}}, + "parent_ids": [], + }, + { + "event": "on_chain_start", + "data": {"input": {"number": 1}}, + "name": "task_with_exception", + "tags": ["seq:step:1"], + "run_id": AnyStr(), + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_start", + "data": {"input": {"number": 1}}, + "name": "my_task", + "tags": ["seq:step:1"], + "run_id": AnyStr(), + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "my_task", + "tags": ["seq:step:1"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "data": {"chunk": 2}, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_end", + "data": {"output": 2, "input": {"number": 1}}, + "run_id": AnyStr(), + "name": "my_task", + "tags": ["seq:step:1"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_task", + "langgraph_triggers": ("__pregel_push",), + "langgraph_path": ( + "__pregel_push", + ("__pregel_pull", "my_workflow"), + 2, + True, + ), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [ + AnyStr(), + AnyStr(), + ], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "my_workflow", + "tags": ["graph:step:4"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_workflow", + "langgraph_triggers": ("__start__",), + "langgraph_path": ("__pregel_pull", "my_workflow"), + "langgraph_checkpoint_ns": AnyStr(), + }, + "data": {"chunk": "done"}, + "parent_ids": [AnyStr()], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "LangGraph", + "tags": [], + "metadata": {"thread_id": "1"}, + "data": {"chunk": {"my_task": 2}}, + "parent_ids": [], + }, + { + "event": "on_chain_end", + "data": {"output": "done", "input": 1}, + "run_id": AnyStr(), + "name": "my_workflow", + "tags": ["graph:step:4"], + "metadata": { + "thread_id": "1", + "langgraph_step": 4, + "langgraph_node": "my_workflow", + "langgraph_triggers": ("__start__",), + "langgraph_path": ("__pregel_pull", "my_workflow"), + "langgraph_checkpoint_ns": AnyStr(), + }, + "parent_ids": [AnyStr()], + }, + { + "event": "on_chain_stream", + "run_id": AnyStr(), + "name": "LangGraph", + "tags": [], + "metadata": {"thread_id": "1"}, + "data": {"chunk": {"my_workflow": "done"}}, + "parent_ids": [], + }, + { + "event": "on_chain_end", + "data": {"output": "done"}, + "run_id": AnyStr(), + "name": "LangGraph", + "tags": [], + "metadata": {"thread_id": "1"}, + "parent_ids": [], + }, + ] + + +@pytest.mark.parametrize("with_timeout", [False, "inner", "outer", "both"]) +@pytest.mark.parametrize("subgraph_persist", [True, False]) +async def test_parent_command_goto( + async_checkpointer: BaseCheckpointSaver, + subgraph_persist: bool, + with_timeout: Literal[False, "inner", "outer", "both"], +) -> None: + class State(TypedDict): + dialog_state: Annotated[list[str], operator.add] + + async def node_a_child(state): + return {"dialog_state": ["a_child_state"]} + + async def node_b_child(state): + return Command( + graph=Command.PARENT, + goto="node_b_parent", + update={"dialog_state": ["b_child_state"]}, + ) + + sub_builder = StateGraph(State) + sub_builder.add_node(node_a_child) + sub_builder.add_node(node_b_child) + sub_builder.add_edge(START, "node_a_child") + sub_builder.add_edge("node_a_child", "node_b_child") + sub_graph = sub_builder.compile(checkpointer=subgraph_persist) + if with_timeout in ("inner", "both"): + sub_graph.step_timeout = 1 + + async def node_b_parent(state): + return {"dialog_state": ["node_b_parent"]} + + main_builder = StateGraph(State) + main_builder.add_node(node_b_parent) + main_builder.add_edge(START, "subgraph_node") + main_builder.add_node("subgraph_node", sub_graph, destinations=("node_b_parent",)) + main_graph = main_builder.compile(async_checkpointer, name="parent") + if with_timeout in ("outer", "both"): + main_graph.step_timeout = 1 + + config = {"configurable": {"thread_id": 1}} + + assert await main_graph.ainvoke( + input={"dialog_state": ["init_state"]}, config=config + ) == {"dialog_state": ["init_state", "b_child_state", "node_b_parent"]} + + +@pytest.mark.parametrize("with_timeout", [True, False]) +async def test_timeout_with_parent_command( + async_checkpointer: BaseCheckpointSaver, with_timeout: bool +) -> None: + """Test that parent commands are properly propagated during timeouts.""" + + class State(TypedDict): + value: str + + async def parent_command_node(state: State) -> State: + await asyncio.sleep(0.1) # Add some delay before raising + return Command(graph=Command.PARENT, goto="test_cmd", update={"key": "value"}) + + builder = StateGraph(State) + builder.add_node("parent_cmd", parent_command_node) + builder.set_entry_point("parent_cmd") + graph = builder.compile(checkpointer=async_checkpointer) + if with_timeout: + graph.step_timeout = 1 + + # Should propagate parent command, not timeout + thread1 = {"configurable": {"thread_id": "1"}} + with pytest.raises(ParentCommand) as exc_info: + await graph.ainvoke({"value": "start"}, thread1) + assert exc_info.value.args[0].goto == "test_cmd" + assert exc_info.value.args[0].update == {"key": "value"} + + +async def test_fork_and_update_task_results( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test forking and updating task results with state history.""" + + def checkpoint(values: dict[str, Any]): + return ("checkpoint", {"values": values}) + + def task(name: str, result: Any): + return ("task", {"name": name, "result": result}) + + def get_tree(history: list[StateSnapshot]) -> list: + """Build a tree structure from state history for comparison.""" + if not history: + return [] + + # Build a tree structure similar to renderForks + node_map: dict[str, dict] = {} + root_nodes: list[dict] = [] + + # Second pass: establish parent-child relationships + for item in reversed(history): + checkpoint_id = item.config["configurable"]["checkpoint_id"] + parent_checkpoint_id = ( + item.parent_config["configurable"]["checkpoint_id"] + if item.parent_config + else None + ) + node_map[checkpoint_id] = {"item": item, "children": []} + + parent = node_map.get(parent_checkpoint_id) + (parent["children"] if parent else root_nodes).append( + node_map[checkpoint_id] + ) + + def node_to_tree(node: dict) -> list: + """Convert a node to tree structure.""" + result = [ + checkpoint(node["item"].values), + ] + [ + task(task_info.name, task_info.result) + for task_info in node["item"].tasks + ] + + if len(node["children"]) > 1: + branches = [node_to_tree(child) for child in node["children"]] + return result + [branches] + elif len(node["children"]) == 1: + return result + node_to_tree(node["children"][0]) + else: + return result + + if len(root_nodes) == 1: + # Process all root nodes + return node_to_tree(root_nodes[0]) + + elif len(root_nodes) > 1: + # Multiple root nodes - treat as branches + branches = [node_to_tree(node) for node in root_nodes] + return branches + else: + return [] + + class State(TypedDict): + name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b] + + # Define the graph with a sequence of nodes + def one(state: State) -> Command: + return Command(goto=[Send("two", {})], update={"name": "one"}) + + def two(state: State) -> State: + return {"name": "two"} + + def three(state: State) -> State: + return {"name": "three"} + + graph = ( + StateGraph(State) + .add_node("one", one) + .add_node("two", two) + .add_node("three", three) + .add_edge(START, "one") + .add_edge("one", "two") + .add_edge("two", "three") + .compile(checkpointer=async_checkpointer) + ) + + config = {"configurable": {"thread_id": "1"}} + history: list[StateSnapshot] = [] + + # Initial run + await graph.ainvoke({"name": "start"}, config) + history = [c async for c in graph.aget_state_history(config)] + + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ] + + # Update the start state + await graph.ainvoke( + None, + await graph.aupdate_state( + history[4].config, + values=[StateUpdate(values={"name": "start*"}, as_node="__start__")], + as_node="__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + ] + + # Fork from task "one" + # Start from the checkpoint that has the task "one" + assert history[3].values == {"name": "start*"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "one" + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[3].config, + [StateUpdate(values={"name": "one*"}, as_node="one")], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": ""}), + task("__start__", {"name": "start*"}), + [ + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one"}), + checkpoint({"name": "start* > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one > two > two > three"}), + ], + [ + checkpoint({"name": "start*"}), + task("one", {"name": "one*"}), + checkpoint({"name": "start* > one*"}), + task("two", {"name": "two"}), + checkpoint({"name": "start* > one* > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start* > one* > two > three"}), + ], + ], + ], + ] + + config = {"configurable": {"thread_id": "2"}} + + # Initialize the thread once again + await graph.ainvoke({"name": "start"}, config) + history = [c async for c in graph.aget_state_history(config)] + + # Fork from task "two" + # Start from the checkpoint that has the task "two" + assert history[2].values == {"name": "start > one"} + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[2].config, + [ + StateUpdate(values={"name": "two"}, as_node="two"), + StateUpdate(values={"name": "two"}, as_node="two"), + ], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ] + + # Fork task three + assert history[1].values == {"name": "start > one > two > two"} + assert len(history[1].tasks) == 1 + assert history[1].tasks[0].name == "three" + + await graph.ainvoke( + None, + await graph.aupdate_state( + history[1].config, + [StateUpdate(values={"name": "three*"}, as_node="three")], + "__copy__", + ), + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + ], + ], + ], + ] + + # Regenerate task three + assert history[3].values == {"name": "start > one > two > two"} + assert len(history[3].tasks) == 1 + assert history[3].tasks[0].name == "three" + + await graph.ainvoke( + None, await graph.aupdate_state(history[3].config, None, "__copy__") + ) + + history = [c async for c in graph.aget_state_history(config)] + assert get_tree(history) == [ + checkpoint({"name": ""}), + task("__start__", {"name": "start"}), + checkpoint({"name": "start"}), + task("one", {"name": "one"}), + [ + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one"}), + task("two", {"name": "two"}), + task("two", {"name": "two"}), + [ + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three*"}), + checkpoint({"name": "start > one > two > two > three*"}), + ], + [ + checkpoint({"name": "start > one > two > two"}), + task("three", {"name": "three"}), + checkpoint({"name": "start > one > two > two > three"}), + ], + ], + ], + ], + ] + + +async def test_subgraph_streaming_async() -> None: + """Test subgraph streaming when used as a node in async version""" + + # Create a fake chat model that returns a simple response + model = GenericFakeChatModel(messages=iter(["The weather is sunny today."])) + + # Create a subgraph that uses the fake chat model + async def call_model_node( + state: MessagesState, config: RunnableConfig + ) -> MessagesState: + """Node that calls the model with the last message.""" + messages = state["messages"] + last_message = messages[-1].content if messages else "" + response = await model.ainvoke([("user", last_message)], config) + return {"messages": [response]} + + # Build the subgraph + subgraph = StateGraph(MessagesState) + subgraph.add_node("call_model", call_model_node) + subgraph.add_edge(START, "call_model") + compiled_subgraph = subgraph.compile() + + class SomeCustomState(TypedDict): + last_chunk: NotRequired[str] + num_chunks: NotRequired[int] + + # Will invoke a subgraph as a function + async def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict: + """Node that runs the subgraph.""" + msgs = {"messages": [("user", "What is the weather in Tokyo?")]} + events = [] + async for event in compiled_subgraph.astream( + msgs, config, stream_mode="messages" + ): + events.append(event) + ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events] + return { + "last_chunk": ai_msg_chunks[-1], + "num_chunks": len(ai_msg_chunks), + } + + # Build the main workflow + workflow = StateGraph(SomeCustomState) + workflow.add_node("subgraph", parent_node) + workflow.add_edge(START, "subgraph") + compiled_workflow = workflow.compile() + + # Test the basic functionality + result = await compiled_workflow.ainvoke({}) + + assert result["last_chunk"].content == "today." + assert result["num_chunks"] == 9 + + +@NEEDS_CONTEXTVARS +async def test_null_resume_disallowed_with_multiple_interrupts( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + text_1: str + text_2: str + + async def human_node_1(state: State): + value = interrupt(state["text_1"]) + return {"text_1": value} + + async def human_node_2(state: State): + value = interrupt(state["text_2"]) + return {"text_2": value} + + graph_builder = StateGraph(State) + graph_builder.add_node("human_node_1", human_node_1) + graph_builder.add_node("human_node_2", human_node_2) + + # Add both nodes in parallel from START + graph_builder.add_edge(START, "human_node_1") + graph_builder.add_edge(START, "human_node_2") + + checkpointer = InMemorySaver() + graph = graph_builder.compile(checkpointer=checkpointer) + + thread_id = str(uuid.uuid4()) + config: RunnableConfig = {"configurable": {"thread_id": thread_id}} + await graph.ainvoke( + {"text_1": "original text 1", "text_2": "original text 2"}, config=config + ) + + resume_map = { + i.id: f"resume for prompt: {i.value}" + for i in (await graph.aget_state(config)).interrupts + } + with pytest.raises( + RuntimeError, + match="When there are multiple pending interrupts, you must specify the interrupt id when resuming.", + ): + await graph.ainvoke(Command(resume="singular resume"), config=config) + + assert await graph.ainvoke(Command(resume=resume_map), config=config) == { + "text_1": "resume for prompt: original text 1", + "text_2": "resume for prompt: original text 2", + } + + +async def test_astream_waiter_cleanup_on_cancel( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that AsyncPregelLoop cleans up waiter tasks after cancellation.""" + + recorded_tasks: list[asyncio.Task[None]] = [] + finished_tasks: list[asyncio.Task[None]] = [] + + original_wait = AsyncQueue.wait + + async def tracked_wait(self: AsyncQueue) -> None: + task = asyncio.current_task() + assert task is not None + recorded_tasks.append(task) + try: + await original_wait(self) + finally: + finished_tasks.append(task) + + monkeypatch.setattr(AsyncQueue, "wait", tracked_wait) + + class State(TypedDict, total=False): + count: int + + async def slow_node(state: State) -> State: + await asyncio.sleep(0.05) + state = dict(state) + state["count"] = state.get("count", 0) + 1 + await asyncio.sleep(0.1) + return state + + builder = StateGraph(State) + builder.add_node("slow", slow_node) + builder.add_edge(START, "slow") + builder.add_edge("slow", END) + graph = builder.compile() + + async def consumer() -> None: + async for _ in graph.astream({"msg": "hi"}, stream_mode="messages"): + await asyncio.sleep(0.01) + + task = asyncio.create_task(consumer()) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + await asyncio.sleep(0.05) + + assert recorded_tasks, "expected stream.wait() task to be created" + assert set(finished_tasks) == set(recorded_tasks) + assert all(t.done() for t in recorded_tasks) + + +@NEEDS_CONTEXTVARS +async def test_interrupt_stream_mode_values(async_checkpointer: BaseCheckpointSaver): + """Test that interrupts are surfaced on 'values' stream mode""" + + class State(TypedDict): + robot_input: str + human_input: str + + def robot_input_node(state: State) -> State: + return {"robot_input": "beep boop i am a robot"} + + def human_input_node(state: State) -> Command: + human_input = interrupt("interrupt") + return Command(update={"human_input": human_input}) + + builder = StateGraph(State) + builder.add_node(robot_input_node) + builder.add_node(human_input_node) + builder.add_edge(START, "robot_input_node") + builder.add_edge("robot_input_node", "human_input_node") + app = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + result = [ + (mode, e) + async for mode, e in app.astream( + State(), config, stream_mode=["updates", "values"] + ) + ] + assert len(result) == 4 + assert result == [ + ("updates", {"robot_input_node": {"robot_input": "beep boop i am a robot"}}), + ("values", {"robot_input": "beep boop i am a robot"}), + ("updates", {"__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),)}), + ( + "values", + { + "robot_input": "beep boop i am a robot", + "__interrupt__": (Interrupt(value="interrupt", id=AnyStr()),), + }, + ), + ] + resume_result = [ + (mode, e) + async for mode, e in app.astream( + Command(resume="i am a human"), config, stream_mode=["updates", "values"] + ) + ] + assert resume_result == [ + ("values", {"robot_input": "beep boop i am a robot"}), + ("updates", {"human_input_node": {"human_input": "i am a human"}}), + ( + "values", + {"robot_input": "beep boop i am a robot", "human_input": "i am a human"}, + ), + ] + + +async def test_supersteps_populate_task_results( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class State(TypedDict): + num: int + text: str + + def double(state: State) -> State: + return {"num": state["num"] * 2, "text": state["text"] * 2} + + graph = ( + StateGraph(State) + .add_node("double", double) + .add_edge(START, "double") + .add_edge("double", END) + .compile(checkpointer=async_checkpointer) + ) + + # reference run with ainvoke + ref_cfg = {"configurable": {"thread_id": "ref"}} + await graph.ainvoke({"num": 1, "text": "one"}, ref_cfg) + ref_history = [h async for h in graph.aget_state_history(ref_cfg)] + + # Helper: pull first task result for a node name from history + def first_task_result(history: list[StateSnapshot], node: str) -> Any: + for s in history: + for t in s.tasks: + if t.name == node: + return t.result + return None + + ref_start_result = first_task_result(ref_history, "__start__") + ref_double_result = first_task_result(ref_history, "double") + assert ref_start_result == {"num": 1, "text": "one"} + assert ref_double_result == {"num": 2, "text": "oneone"} + + # using supersteps + bulk_cfg = {"configurable": {"thread_id": "bulk"}} + await graph.abulk_update_state( + bulk_cfg, + [ + [StateUpdate(values={}, as_node="__input__")], + [StateUpdate(values={"num": 1, "text": "one"}, as_node="__start__")], + [StateUpdate(values={"num": 2, "text": "oneone"}, as_node="double")], + ], + ) + bulk_history = [h async for h in graph.aget_state_history(bulk_cfg)] + + bulk_start_result = first_task_result(bulk_history, "__start__") + bulk_double_result = first_task_result(bulk_history, "double") + + assert bulk_start_result == ref_start_result == {"num": 1, "text": "one"} + assert bulk_double_result == ref_double_result == {"num": 2, "text": "oneone"} + + +async def test_fork_does_not_apply_pending_writes( + async_checkpointer: BaseCheckpointSaver, +) -> None: + """Test that forking with aupdate_state does not apply pending writes from original execution.""" + + class State(TypedDict): + value: Annotated[int, operator.add] + + def node_a(state: State) -> State: + return {"value": 10} + + def node_b(state: State) -> State: + return {"value": 100} + + graph = ( + StateGraph(State) + .add_node("node_a", node_a) + .add_node("node_b", node_b) + .add_edge(START, "node_a") + .add_edge("node_a", "node_b") + .compile(checkpointer=async_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + await graph.ainvoke({"value": 1}, thread1) + + history = [c async for c in graph.aget_state_history(thread1)] + checkpoint_before_a = next(s for s in history if s.next == ("node_a",)) + + fork_config = await graph.aupdate_state( + checkpoint_before_a.config, {"value": 20}, as_node="node_a" + ) + result = await graph.ainvoke(None, fork_config) + + # 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 + assert result == {"value": 121} diff --git a/langgraph/source/libs/langgraph/tests/test_pydantic.py b/langgraph/source/libs/langgraph/tests/test_pydantic.py new file mode 100644 index 0000000000000000000000000000000000000000..87049204fb439e20e73198b9d2c57dad8beb4977 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_pydantic.py @@ -0,0 +1,314 @@ +import datetime +import decimal +import ipaddress +import pathlib +import re +import sys +import uuid +from enum import Enum +from typing import Annotated, Literal, Optional + +from pydantic import ( + BaseModel, + ByteSize, + Field, + SecretStr, + confloat, + conint, + conlist, + constr, + field_validator, + model_validator, +) + +from langgraph._internal._pydantic import is_supported_by_pydantic +from langgraph.constants import END, START +from langgraph.graph.state import StateGraph + + +def test_is_supported_by_pydantic() -> None: + """Test if types are supported by pydantic.""" + import typing + + import pydantic + import typing_extensions + + class TypedDictExtensions(typing_extensions.TypedDict): + x: int + + assert is_supported_by_pydantic(TypedDictExtensions) is True + + class VanillaClass: + x: int + + assert is_supported_by_pydantic(VanillaClass) is False + + class BuiltinTypedDict(typing.TypedDict): # noqa: TID251 + x: int + + if sys.version_info >= (3, 12): + assert is_supported_by_pydantic(BuiltinTypedDict) is True + else: + assert is_supported_by_pydantic(BuiltinTypedDict) is False + + class PydanticModel(pydantic.BaseModel): + x: int + + assert is_supported_by_pydantic(PydanticModel) is True + + +def test_nested_pydantic_models() -> None: + """Test that nested Pydantic models are properly constructed from leaf nodes up.""" + + class NestedModel(BaseModel): + value: int + name: str + + # For constrained types + PositiveInt = Annotated[int, Field(gt=0)] + NonNegativeFloat = Annotated[float, Field(ge=0)] + + # Enum type + class UserRole(Enum): + ADMIN = "admin" + USER = "user" + GUEST = "guest" + + # Forward reference model + class RecursiveModel(BaseModel): + value: str + child: Optional["RecursiveModel"] = None + + # Discriminated union models + class Cat(BaseModel): + pet_type: Literal["cat"] + meow: str + + class Dog(BaseModel): + pet_type: Literal["dog"] + bark: str + + # Cyclic reference model + class Person(BaseModel): + id: str + name: str + friends: list[str] = Field(default_factory=list) # IDs of friends + + conlist_type = conlist(item_type=int, min_length=2, max_length=5) + + class State(BaseModel): + # Basic nested model tests + top_level: str + auuid: uuid.UUID + nested: NestedModel + optional_nested: Annotated[NestedModel | None, lambda x, y: y, "Foo"] + dict_nested: dict[str, NestedModel] + simple_str_list: list[str] + list_nested: Annotated[ + dict | list[dict[str, NestedModel]], lambda x, y: (x or []) + [y] + ] + tuple_nested: tuple[str, NestedModel] + tuple_list_nested: list[tuple[int, NestedModel]] + complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]] + + # Forward reference test + recursive: RecursiveModel + + # Discriminated union test + pet: Cat | Dog + + # Cyclic reference test + people: dict[str, Person] # Map of ID -> Person + + # Rich type adapters + ip_address: ipaddress.IPv4Address + ip_address_v6: ipaddress.IPv6Address + amount: decimal.Decimal + file_path: pathlib.Path + timestamp: datetime.datetime + date_only: datetime.date + time_only: datetime.time + duration: datetime.timedelta + immutable_set: frozenset[int] + binary_data: bytes + pattern: re.Pattern + secret: SecretStr + file_size: ByteSize + + # Constrained types + positive_value: PositiveInt + non_negative: NonNegativeFloat + limited_string: constr(min_length=3, max_length=10) + bounded_int: conint(ge=10, le=100) + restricted_float: confloat(gt=0, lt=1) + required_list: conlist_type + + # Enum & Literal + role: UserRole + status: Literal["active", "inactive", "pending"] + + # Annotated & NewType + validated_age: Annotated[int, Field(gt=0, lt=120)] + + # Generic containers with validators + decimal_list: list[decimal.Decimal] + id_tuple: tuple[uuid.UUID, uuid.UUID] + + inputs = { + # Basic nested models + "top_level": "initial", + "auuid": str(uuid.uuid4()), + "nested": {"value": 42, "name": "test"}, + "optional_nested": {"value": 10, "name": "optional"}, + "dict_nested": {"a": {"value": 5, "name": "a"}}, + "list_nested": [{"a": {"value": 6, "name": "b"}}], + "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}], + "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]], + "simple_str_list": ["siss", "boom", "bah"], + "complex_tuple": [ + "complex", + {"nested": [9, {"value": 10, "name": "deep"}]}, + ], + # Forward reference + "recursive": {"value": "parent", "child": {"value": "child", "child": None}}, + # Discriminated union (using a cat in this case) + "pet": {"pet_type": "cat", "meow": "meow!"}, + # Cyclic references + "people": { + "1": { + "id": "1", + "name": "Alice", + "friends": ["2", "3"], # Alice is friends with Bob and Charlie + }, + "2": { + "id": "2", + "name": "Bob", + "friends": ["1"], # Bob is friends with Alice + }, + "3": { + "id": "3", + "name": "Charlie", + "friends": ["1", "2"], # Charlie is friends with Alice and Bob + }, + }, + # Rich type adapters + "ip_address": "192.168.1.1", + "ip_address_v6": "2001:db8::1", + "amount": "123.45", + "file_path": "/tmp/test.txt", + "timestamp": "2025-04-07T10:58:04", + "date_only": "2025-04-07", + "time_only": "10:58:04", + "duration": 3600, # seconds + "immutable_set": [1, 2, 3, 4], + "binary_data": b"hello world", + "pattern": "^test$", + "secret": "password123", + "file_size": 1024, + # Constrained types + "positive_value": 42, + "non_negative": 0.0, + "limited_string": "test", + "bounded_int": 50, + "restricted_float": 0.5, + "required_list": [10, 20, 30], + # Enum & Literal + "role": "admin", + "status": "active", + # Annotated & NewType + "validated_age": 30, + # Generic containers with validators + "decimal_list": ["10.5", "20.75", "30.25"], + "id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())], + } + + update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}} + + expected = State(**inputs) + + def node_fn(state: State) -> dict: + # Basic assertions + assert isinstance(state.auuid, uuid.UUID) + assert state == expected + + # Rich type assertions + assert isinstance(state.ip_address, ipaddress.IPv4Address) + assert isinstance(state.ip_address_v6, ipaddress.IPv6Address) + assert isinstance(state.amount, decimal.Decimal) + assert isinstance(state.file_path, pathlib.Path) + assert isinstance(state.timestamp, datetime.datetime) + assert isinstance(state.date_only, datetime.date) + assert isinstance(state.time_only, datetime.time) + assert isinstance(state.duration, datetime.timedelta) + assert isinstance(state.immutable_set, frozenset) + assert isinstance(state.binary_data, bytes) + assert isinstance(state.pattern, re.Pattern) + + # Constrained types + assert state.positive_value > 0 + assert state.non_negative >= 0 + assert 3 <= len(state.limited_string) <= 10 + assert 10 <= state.bounded_int <= 100 + assert 0 < state.restricted_float < 1 + assert 2 <= len(state.required_list) <= 5 + + # Enum & Literal + assert state.role == UserRole.ADMIN + assert state.status == "active" + + # Annotated + assert 0 < state.validated_age < 120 + + # Generic containers + assert len(state.decimal_list) == 3 + assert len(state.id_tuple) == 2 + + return update + + builder = StateGraph(State) + builder.add_node("process", node_fn) + builder.set_entry_point("process") + builder.set_finish_point("process") + graph = builder.compile() + + result = graph.invoke(inputs.copy()) + + assert result == {**inputs, **update} + + new_inputs = inputs.copy() + new_inputs["list_nested"] = {"foo": "bar"} + expected = State(**new_inputs) + assert {**new_inputs, **update} == graph.invoke(new_inputs.copy()) + + +def test_pydantic_state_field_validator(): + class State(BaseModel): + name: str + text: str = "" + only_root: int = 13 + + @field_validator("name", mode="after") + @classmethod + def validate_name(cls, value): + if value[0].islower(): + raise ValueError("Name must start with a capital letter") + return "Validated " + value + + @model_validator(mode="before") + @classmethod + def validate_amodel(cls, values: "State"): + return values | {"only_root": 392} + + input_state = {"name": "John"} + + def process_node(state: State): + assert State.model_validate(input_state) == state + return {"text": "Hello, " + state.name + "!"} + + builder = StateGraph(state_schema=State) + builder.add_node("process", process_node) + builder.add_edge(START, "process") + builder.add_edge("process", END) + g = builder.compile() + res = g.invoke(input_state) + assert res["text"] == "Hello, Validated John!" diff --git a/langgraph/source/libs/langgraph/tests/test_remote_graph.py b/langgraph/source/libs/langgraph/tests/test_remote_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..e0eacf00665d32208b87798c52e1dc02001d4ee5 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_remote_graph.py @@ -0,0 +1,1301 @@ +import re +import sys +from typing import Annotated +from unittest.mock import AsyncMock, MagicMock + +import langsmith as ls +import pytest +from langchain_core.messages import AnyMessage, BaseMessage +from langchain_core.runnables import RunnableConfig +from langchain_core.runnables.graph import Edge as DrawableEdge +from langchain_core.runnables.graph import Node as DrawableNode +from langgraph_sdk.schema import StreamPart +from typing_extensions import TypedDict + +from langgraph.errors import GraphInterrupt +from langgraph.graph import StateGraph, add_messages +from langgraph.pregel import Pregel +from langgraph.pregel.remote import RemoteGraph +from langgraph.types import Interrupt, StateSnapshot +from tests.any_str import AnyStr +from tests.conftest import NO_DOCKER +from tests.example_app.example_graph import app + +if NO_DOCKER: + pytest.skip( + "Skipping tests that require Docker. Unset NO_DOCKER to run them.", + allow_module_level=True, + ) + +pytestmark = pytest.mark.anyio + +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + +SKIP_PYTHON_314 = pytest.mark.skipif( + sys.version_info >= (3, 14), + reason="Not yet testing Python 3.14 with the server bc of dependency limits on the api side", +) + + +def test_with_config(): + # set up test + remote_pregel = RemoteGraph( + "test_graph_id", + config={ + "configurable": { + "foo": "bar", + "thread_id": "thread_id_1", + } + }, + ) + + # call method / assertions + config = {"configurable": {"hello": "world"}} + remote_pregel_copy = remote_pregel.with_config(config) + + # assert that a copy was returned + assert remote_pregel_copy != remote_pregel + # assert that configs were merged + assert remote_pregel_copy.config == { + "configurable": { + "foo": "bar", + "thread_id": "thread_id_1", + "hello": "world", + } + } + + +def test_get_graph(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.assistants.get_graph.return_value = { + "nodes": [ + {"id": "__start__", "type": "schema", "data": "__start__"}, + {"id": "__end__", "type": "schema", "data": "__end__"}, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "RunnableCallable"], + "name": "agent_1", + }, + }, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__"}, + ], + } + + remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) + + # call method / assertions + drawable_graph = remote_pregel.get_graph() + + assert drawable_graph.nodes == { + "__start__": DrawableNode( + id="__start__", name="__start__", data="__start__", metadata=None + ), + "__end__": DrawableNode( + id="__end__", name="__end__", data="__end__", metadata=None + ), + "agent": DrawableNode( + id="agent", + name="agent_1", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, + metadata=None, + ), + } + + assert drawable_graph.edges == [ + DrawableEdge(source="__start__", target="agent"), + DrawableEdge(source="agent", target="__end__"), + ] + + +@pytest.mark.anyio +async def test_aget_graph(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.assistants.get_graph.return_value = { + "nodes": [ + {"id": "__start__", "type": "schema", "data": "__start__"}, + {"id": "__end__", "type": "schema", "data": "__end__"}, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "RunnableCallable"], + "name": "agent_1", + }, + }, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__"}, + ], + } + + remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client) + + # call method / assertions + drawable_graph = await remote_pregel.aget_graph() + + assert drawable_graph.nodes == { + "__start__": DrawableNode( + id="__start__", name="__start__", data="__start__", metadata=None + ), + "__end__": DrawableNode( + id="__end__", name="__end__", data="__end__", metadata=None + ), + "agent": DrawableNode( + id="agent", + name="agent_1", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, + metadata=None, + ), + } + + assert drawable_graph.edges == [ + DrawableEdge(source="__start__", target="agent"), + DrawableEdge(source="agent", target="__end__"), + ] + + +def test_get_state(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.threads.get_state.return_value = { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_snapshot = remote_pregel.get_state(config) + + assert state_snapshot == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + interrupts=(), + ) + + +@pytest.mark.anyio +async def test_aget_state(): + mock_async_client = AsyncMock() + mock_async_client.threads.get_state.return_value = { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_2", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "tasks": [], + } + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_snapshot = await remote_pregel.aget_state(config) + + assert state_snapshot == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_2", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + tasks=(), + interrupts=(), + ) + + +def test_get_state_history(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.threads.get_history.return_value = [ + { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + ] + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_history_snapshot = list( + remote_pregel.get_state_history(config, filter=None, before=None, limit=None) + ) + + assert len(state_history_snapshot) == 1 + assert state_history_snapshot[0] == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + interrupts=(), + ) + + +@pytest.mark.anyio +async def test_aget_state_history(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.threads.get_history.return_value = [ + { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + ] + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_history_snapshot = [] + async for state_snapshot in remote_pregel.aget_state_history( + config, filter=None, before=None, limit=None + ): + state_history_snapshot.append(state_snapshot) + + assert len(state_history_snapshot) == 1 + assert state_history_snapshot[0] == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + interrupts=(), + ) + + +def test_update_state(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.threads.update_state.return_value = { + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + response = remote_pregel.update_state(config, {"key": "value"}) + + assert response == { + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + +@pytest.mark.anyio +async def test_aupdate_state(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.threads.update_state.return_value = { + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) + + config = {"configurable": {"thread_id": "thread1"}} + response = await remote_pregel.aupdate_state(config, {"key": "value"}) + + assert response == { + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + +def test_stream(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.runs.stream.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart(event="values", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart( + event="messages", + data=[ + { + "content": [{"text": "Hello", "type": "text", "index": 0}], + "type": "AIMessageChunk", + }, + { + "langgraph_step": 1, + "langgraph_node": "call_llm", + "langgraph_triggers": ["branch:to:call_llm"], + "langgraph_path": ["__pregel_pull", "call_llm"], + }, + ], + ), + StreamPart( + event="updates", + data={ + "__interrupt__": [ + { + "value": {"question": "Does this look good?"}, + "id": AnyStr(), + } + ] + }, + ), + ] + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + # test raising graph interrupt if invoked as a subgraph + with pytest.raises(GraphInterrupt) as exc: + for stream_part in remote_pregel.stream( + {"input": "data"}, + # pretend we invoked this as a subgraph + config={ + "configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"} + }, + stream_mode="values", + ): + pass + + assert exc.value.args[0] == [ + Interrupt( + value={"question": "Does this look good?"}, + id=AnyStr(), + ) + ] + + # stream modes doesn't include 'updates' + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="values", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data1"}, + {"chunk": "data2"}, + {"chunk": "data3"}, + ] + + # stream_mode messages + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="messages", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ( + { + "content": [{"text": "Hello", "type": "text", "index": 0}], + "type": "AIMessageChunk", + }, + { + "langgraph_step": 1, + "langgraph_node": "call_llm", + "langgraph_triggers": ["branch:to:call_llm"], + "langgraph_path": ["__pregel_pull", "call_llm"], + }, + ), + ] + + mock_sync_client.runs.stream.return_value = [ + StreamPart(event="updates", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), + ] + + # default stream_mode is updates + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data3"}, + {"chunk": "data4"}, + {"__interrupt__": ()}, + ] + + # list stream_mode includes mode names + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ("updates", {"chunk": "data3"}, None), + ("updates", {"chunk": "data4"}, None), + ("updates", {"__interrupt__": ()}, None), + ] + + # subgraphs + list modes + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), "updates", {"chunk": "data3"}), + ((), "updates", {"chunk": "data4"}), + ((), "updates", {"__interrupt__": ()}), + ] + + # subgraphs + single mode + stream_parts = [] + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), {"chunk": "data3"}), + ((), {"chunk": "data4"}), + ((), {"__interrupt__": ()}), + ] + + +@pytest.mark.anyio +async def test_astream(): + # set up test + mock_async_client = MagicMock() + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart(event="values", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart( + event="messages", + data=[ + { + "content": [{"text": "Hello", "type": "text", "index": 0}], + "type": "AIMessageChunk", + }, + { + "langgraph_step": 1, + "langgraph_node": "call_llm", + "langgraph_triggers": ["branch:to:call_llm"], + "langgraph_path": ["__pregel_pull", "call_llm"], + }, + ], + ), + StreamPart( + event="updates", + data={ + "__interrupt__": [ + { + "value": {"question": "Does this look good?"}, + "id": AnyStr(), + } + ] + }, + ), + ] + mock_async_client.runs.stream.return_value = async_iter + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) + + # test raising graph interrupt if invoked as a subgraph + with pytest.raises(GraphInterrupt) as exc: + async for stream_part in remote_pregel.astream( + {"input": "data"}, + # pretend we invoked this as a subgraph + config={ + "configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"} + }, + stream_mode="values", + ): + pass + + assert exc.value.args[0] == [ + Interrupt( + value={"question": "Does this look good?"}, + id=AnyStr(), + ) + ] + + # stream modes doesn't include 'updates' + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="values", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data1"}, + {"chunk": "data2"}, + {"chunk": "data3"}, + ] + + # stream_mode messages + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="messages", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ( + { + "content": [{"text": "Hello", "type": "text", "index": 0}], + "type": "AIMessageChunk", + }, + { + "langgraph_step": 1, + "langgraph_node": "call_llm", + "langgraph_triggers": ["branch:to:call_llm"], + "langgraph_path": ["__pregel_pull", "call_llm"], + }, + ), + ] + + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="updates", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), + ] + mock_async_client.runs.stream.return_value = async_iter + + # default stream_mode is updates + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data3"}, + {"chunk": "data4"}, + {"__interrupt__": ()}, + ] + + # list stream_mode includes mode names + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ("updates", {"chunk": "data3"}, None), + ("updates", {"chunk": "data4"}, None), + ("updates", {"__interrupt__": ()}, None), + ] + + # subgraphs + list modes + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), "updates", {"chunk": "data3"}), + ((), "updates", {"chunk": "data4"}), + ((), "updates", {"__interrupt__": ()}), + ] + + # subgraphs + single mode + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), {"chunk": "data3"}), + ((), {"chunk": "data4"}), + ((), {"__interrupt__": ()}), + ] + + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="updates|my|subgraph", data={"chunk": "data3"}), + StreamPart(event="updates|hello|subgraph", data={"chunk": "data4"}), + StreamPart(event="updates|bye|subgraph", data={"__interrupt__": ()}), + ] + mock_async_client.runs.stream.return_value = async_iter + + # subgraphs + list modes + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + (("my", "subgraph"), "updates", {"chunk": "data3"}), + (("hello", "subgraph"), "updates", {"chunk": "data4"}), + (("bye", "subgraph"), "updates", {"__interrupt__": ()}), + ] + + # subgraphs + single mode + stream_parts = [] + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + (("my", "subgraph"), {"chunk": "data3"}), + (("hello", "subgraph"), {"chunk": "data4"}), + (("bye", "subgraph"), {"__interrupt__": ()}), + ] + + +def test_invoke(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.runs.stream.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart( + event="values", data={"messages": [{"type": "human", "content": "world"}]} + ), + ] + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + config = {"configurable": {"thread_id": "thread_1"}} + result = remote_pregel.invoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config + ) + + assert result == {"messages": [{"type": "human", "content": "world"}]} + + +def test_invoke_sanitizes_thread_id(): + # Ensure that invoking with thread_id passes thread_id as a top-level arg + # and removes it from the config body. + mock_sync_client = MagicMock() + mock_sync_client.runs.stream.return_value = [] + remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) + + config = {"configurable": {"thread_id": "thread_1"}} + remote_pregel.invoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config + ) + + assert mock_sync_client.runs.stream.called + _, kwargs = mock_sync_client.runs.stream.call_args + assert kwargs.get("thread_id") == "thread_1" + passed_config = kwargs.get("config") or {} + assert "configurable" in passed_config + assert "thread_id" not in passed_config["configurable"] + assert not passed_config["configurable"] + + +def test_stream_sanitizes_thread_id(): + # Ensure that streaming with thread_id passes thread_id as a top-level arg + # and removes it from the config body. + mock_sync_client = MagicMock() + mock_sync_client.runs.stream.return_value = [] + remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) + + config = {"configurable": {"thread_id": "thread_2"}} + list(remote_pregel.stream({"input": {"messages": []}}, config)) + + assert mock_sync_client.runs.stream.called + _, kwargs = mock_sync_client.runs.stream.call_args + assert kwargs.get("thread_id") == "thread_2" + passed_config = kwargs.get("config") or {} + assert "configurable" in passed_config + assert "thread_id" not in passed_config["configurable"] + assert not passed_config["configurable"] + + +@pytest.mark.anyio +async def test_ainvoke(): + # set up test + mock_async_client = MagicMock() + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart( + event="values", data={"messages": [{"type": "human", "content": "world"}]} + ), + ] + mock_async_client.runs.stream.return_value = async_iter + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) + + config = {"configurable": {"thread_id": "thread_1"}} + result = await remote_pregel.ainvoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config + ) + + assert result == {"messages": [{"type": "human", "content": "world"}]} + + +@pytest.mark.skip( + "Unskip this test to manually test the LangSmith Deployment integration" +) +@pytest.mark.anyio +async def test_langgraph_cloud_integration(): + from langgraph.checkpoint.memory import InMemorySaver + from langgraph_sdk.client import get_client, get_sync_client + + from langgraph.graph import END, START, MessagesState, StateGraph + + # create RemotePregel instance + client = get_client(url="http://localhost:8123") + sync_client = get_sync_client(url="http://localhost:8123") + remote_pregel = RemoteGraph( + "agent", + client=client, + sync_client=sync_client, + ) + + # define graph + workflow = StateGraph(MessagesState) + workflow.add_node("agent", remote_pregel) + workflow.add_edge(START, "agent") + workflow.add_edge("agent", END) + app = workflow.compile(checkpointer=InMemorySaver()) + + # test invocation + input = { + "messages": [ + { + "role": "human", + "content": "What's the weather in SF?", + } + ] + } + + # test invoke + app.invoke( + input, + config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}}, + interrupt_before=["agent"], + ) + + # test stream + async for _ in app.astream( + input, + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + subgraphs=True, + stream_mode=["debug", "messages"], + ): + pass + + # test stream events + async for chunk in remote_pregel.astream_events( + input, + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + version="v2", + subgraphs=True, + stream_mode=[], + ): + pass + + # test get state + await remote_pregel.aget_state( + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + subgraphs=True, + ) + + # test update state + await remote_pregel.aupdate_state( + config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}}, + values={ + "messages": [ + { + "role": "ai", + "content": "Hello world again!", + } + ] + }, + ) + + # test get history + async for state in remote_pregel.aget_state_history( + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + ): + pass + + # test get graph + remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID + await remote_pregel.aget_graph(xray=True) + + +def test_sanitize_config(): + # Create a test instance + remote = RemoteGraph("test-graph") + + # Test 1: Basic config with primitives + basic_config: RunnableConfig = { + "recursion_limit": 10, + "tags": ["tag1", "tag2"], + "metadata": {"str_key": "value", "int_key": 42, "bool_key": True}, + "configurable": {"param1": "value1", "param2": 123}, + } + sanitized = remote._sanitize_config(basic_config) + assert sanitized["recursion_limit"] == 10 + assert sanitized["tags"] == ["tag1", "tag2"] + assert sanitized["metadata"] == { + "str_key": "value", + "int_key": 42, + "bool_key": True, + } + assert sanitized["configurable"] == {"param1": "value1", "param2": 123} + + # Test 2: Config with non-string tags and complex metadata + complex_config: RunnableConfig = { + "tags": ["tag1", 123, {"obj": "tag"}, "tag2"], # Only string tags should remain + "metadata": { + "nested": { + "key": "value", + "num": 42, + "invalid": lambda x: x, + }, # Last item should be removed + "list": [1, 2, "three"], + "invalid": lambda x: x, # Should be removed + "tuple": (1, 2, 3), # Should be converted to list + }, + } + sanitized = remote._sanitize_config(complex_config) + assert sanitized["tags"] == ["tag1", "tag2"] + assert sanitized["metadata"] == { + "nested": {"key": "value", "num": 42}, + "list": [1, 2, "three"], + "tuple": [1, 2, 3], + } + assert "invalid" not in sanitized["metadata"] + + # Test 3: Config with configurable fields that should be dropped + config_with_drops: RunnableConfig = { + "configurable": { + "normal_param": "value", + "checkpoint_map": {"key": "value"}, # Should be dropped + "checkpoint_id": "123", # Should be dropped + "checkpoint_ns": "ns", # Should be dropped + } + } + sanitized = remote._sanitize_config(config_with_drops) + assert sanitized["configurable"] == {"normal_param": "value"} + assert "checkpoint_map" not in sanitized["configurable"] + assert "checkpoint_id" not in sanitized["configurable"] + assert "checkpoint_ns" not in sanitized["configurable"] + + # Test 4: Empty config + empty_config: RunnableConfig = {} + sanitized = remote._sanitize_config(empty_config) + assert sanitized == {} + + # Test 5: Config with non-string keys in configurable + invalid_keys_config: RunnableConfig = { + "configurable": { + "valid": "value", + 123: "invalid", # Should be dropped + ("tuple", "key"): "invalid", # Should be dropped + } + } + sanitized = remote._sanitize_config(invalid_keys_config) + assert sanitized["configurable"] == {"valid": "value"} + + # Test 6: Deeply nested structures + nested_config: RunnableConfig = { + "metadata": { + "level1": { + "level2": { + "level3": { + "str": "value", + "list": [1, [2, [3]]], + "dict": {"a": {"b": {"c": "d"}}}, + } + } + } + } + } + sanitized = remote._sanitize_config(nested_config) + assert sanitized["metadata"]["level1"]["level2"]["level3"]["str"] == "value" + assert sanitized["metadata"]["level1"]["level2"]["level3"]["list"] == [1, [2, [3]]] + assert sanitized["metadata"]["level1"]["level2"]["level3"]["dict"] == { + "a": {"b": {"c": "d"}} + } + + +"""Test RemoteGraph against an actual server.""" + + +@pytest.fixture +def remote_graph() -> RemoteGraph: + return RemoteGraph("app", url="http://localhost:2024") + + +@pytest.fixture +def nested_remote_graph(remote_graph: RemoteGraph) -> Pregel: + class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + return ( + StateGraph(State) + .add_node("nested", remote_graph) + .add_edge("__start__", "nested") + .compile(name="nested_remote_graph") + ) + + +@pytest.fixture +async def nested_graph() -> Pregel: + class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + return ( + StateGraph(State) + .add_node("nested", app) + .add_edge("__start__", "nested") + .compile(name="nested_graph") + ) + + +def get_message_dict(msg: BaseMessage | dict): + # just get the core stuff from within the message + if isinstance(msg, dict): + return { + "content": msg.get("content"), + "type": msg.get("type"), + "name": msg.get("name"), + "tool_calls": msg.get("tool_calls"), + "invalid_tool_calls": msg.get("invalid_tool_calls"), + } + return { + "content": msg.content, + "type": msg.type, + "name": msg.name, + "tool_calls": getattr(msg, "tool_calls", None), + "invalid_tool_calls": getattr(msg, "invalid_tool_calls", None), + } + + +@NEEDS_CONTEXTVARS +@SKIP_PYTHON_314 +async def test_remote_graph_basic_invoke(remote_graph: RemoteGraph) -> None: + # Basic smoke test of the remote graph + response = await remote_graph.ainvoke( + {"messages": [{"role": "user", "content": "hello"}]} + ) + assert response == { + "content": "answer", + "additional_kwargs": {}, + "response_metadata": {}, + "type": "ai", + "name": None, + "id": "ai3", + "tool_calls": [], + "invalid_tool_calls": [], + "usage_metadata": None, + } + + +class monotonic_uid: + def __init__(self): + self._uid = 0 + + def __call__(self, match=None): + val = self._uid + self._uid += 1 + hexval = f"{val:032x}" + uuid_str = f"{hexval[:8]}-{hexval[8:12]}-{hexval[12:16]}-{hexval[16:20]}-{hexval[20:32]}" + return uuid_str + + +uid_pattern = re.compile( + r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}" +) + + +@NEEDS_CONTEXTVARS +@SKIP_PYTHON_314 +async def test_remote_graph_stream_messages_tuple( + nested_graph: Pregel, nested_remote_graph: Pregel +) -> None: + events = [] + namespaces = [] + uid_generator = monotonic_uid() + async for ns, messages in nested_remote_graph.astream( + {"messages": [{"role": "user", "content": "hello"}]}, + stream_mode="messages", + subgraphs=True, + ): + events.extend(messages) + namespaces.append( + tuple(uid_pattern.sub(uid_generator, ns_part) for ns_part in ns) + ) + inmem_events = [] + inmem_namespaces = [] + uid_generator = monotonic_uid() + async for ns, messages in nested_graph.astream( + {"messages": [{"role": "user", "content": "hello"}]}, + stream_mode="messages", + subgraphs=True, + ): + inmem_events.extend(messages) + inmem_namespaces.append( + tuple(uid_pattern.sub(uid_generator, ns_part) for ns_part in ns) + ) + assert len(events) == len(inmem_events) + assert len(namespaces) == len(inmem_namespaces) + + coerced_events = [get_message_dict(e) for e in events] + coerced_inmem_events = [get_message_dict(e) for e in inmem_events] + assert coerced_events == coerced_inmem_events + # TODO: Fix the namespace matching in the next api release. + # assert namespaces == inmem_namespaces + + +@pytest.mark.anyio +@pytest.mark.parametrize("distributed_tracing", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("headers", [None, {"foo": "bar"}]) +async def test_include_headers( + distributed_tracing: bool, stream: bool, headers: dict[str, str] | None +): + mock_async_client = MagicMock() + async_iter = MagicMock() + return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + ] + async_iter.__aiter__.return_value = return_value + astream_mock = mock_async_client.runs.stream + astream_mock.return_value = async_iter + + mock_sync_client = MagicMock() + sync_iter = MagicMock() + sync_iter.__iter__.return_value = return_value + stream_mock = mock_sync_client.runs.stream + stream_mock.return_value = async_iter + + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + sync_client=mock_sync_client, + distributed_tracing=distributed_tracing, + ) + + config = {"configurable": {"thread_id": "thread_1"}} + with ls.tracing_context(enabled=True, client=MagicMock()): + with ls.trace("foo"): + if stream: + async for _ in remote_pregel.astream( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, + config, + headers=headers, + ): + pass + + else: + await remote_pregel.ainvoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, + config, + headers=headers, + ) + expected = headers.copy() if headers else None + if distributed_tracing: + if expected is None: + expected = {} + expected["langsmith-trace"] = AnyStr() + expected["baggage"] = AnyStr("langsmith-metadata=") + + assert astream_mock.call_args.kwargs["headers"] == expected + stream_mock.assert_not_called() + + with ls.tracing_context(enabled=True, client=MagicMock()): + with ls.trace("foo"): + if stream: + for _ in remote_pregel.stream( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, + config, + headers=headers, + ): + pass + + else: + remote_pregel.invoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, + config, + headers=headers, + ) + assert stream_mock.call_args.kwargs["headers"] == expected diff --git a/langgraph/source/libs/langgraph/tests/test_retry.py b/langgraph/source/libs/langgraph/tests/test_retry.py new file mode 100644 index 0000000000000000000000000000000000000000..ac37bea91f7ac49f851e5ea2ead8a27aea31a7d1 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_retry.py @@ -0,0 +1,344 @@ +from unittest.mock import Mock, patch + +import pytest +from typing_extensions import TypedDict + +from langgraph.graph import START, StateGraph +from langgraph.pregel._retry import _should_retry_on +from langgraph.types import RetryPolicy + + +def test_should_retry_on_single_exception(): + """Test retry with a single exception type.""" + policy = RetryPolicy(retry_on=ValueError) + + # Should retry on ValueError + assert _should_retry_on(policy, ValueError("test error")) is True + + # Should not retry on other exceptions + assert _should_retry_on(policy, TypeError("test error")) is False + assert _should_retry_on(policy, Exception("test error")) is False + + +def test_should_retry_on_sequence_of_exceptions(): + """Test retry with a sequence of exception types.""" + policy = RetryPolicy(retry_on=(ValueError, KeyError)) + + # Should retry on listed exceptions + assert _should_retry_on(policy, ValueError("test error")) is True + assert _should_retry_on(policy, KeyError("test error")) is True + + # Should not retry on other exceptions + assert _should_retry_on(policy, TypeError("test error")) is False + assert _should_retry_on(policy, Exception("test error")) is False + + +def test_should_retry_on_subclass_of_exception(): + """Test retry on subclass of specified exception.""" + + class CustomError(ValueError): + pass + + policy = RetryPolicy(retry_on=ValueError) + + # Should retry on subclass of specified exception + assert _should_retry_on(policy, CustomError("test error")) is True + + +def test_should_retry_on_callable(): + """Test retry with a callable predicate.""" + + # Only retry on ValueError with message containing 'retry' + def should_retry(exc: Exception) -> bool: + return isinstance(exc, ValueError) and "retry" in str(exc) + + policy = RetryPolicy(retry_on=should_retry) + + # Should retry when predicate returns True + assert _should_retry_on(policy, ValueError("please retry this")) is True + + # Should not retry when predicate returns False + assert _should_retry_on(policy, ValueError("other error")) is False + assert _should_retry_on(policy, TypeError("please retry this")) is False + + +def test_should_retry_on_invalid_type(): + """Test retry with an invalid retry_on type.""" + policy = RetryPolicy(retry_on=123) # type: ignore + + with pytest.raises(TypeError, match="retry_on must be an Exception class"): + _should_retry_on(policy, ValueError("test error")) + + +def test_should_retry_on_empty_sequence(): + """Test retry with an empty sequence.""" + policy = RetryPolicy(retry_on=()) + + # Should not retry when sequence is empty + assert _should_retry_on(policy, ValueError("test error")) is False + + +def test_should_retry_default_retry_on(): + """Test the default retry_on function.""" + import httpx + import requests + + # Create a RetryPolicy with default_retry_on + policy = RetryPolicy() + + # Should retry on ConnectionError + assert _should_retry_on(policy, ConnectionError("connection refused")) is True + + # Should not retry on common programming errors + assert _should_retry_on(policy, ValueError("invalid value")) is False + assert _should_retry_on(policy, TypeError("invalid type")) is False + assert _should_retry_on(policy, ArithmeticError("division by zero")) is False + assert _should_retry_on(policy, ImportError("module not found")) is False + assert _should_retry_on(policy, LookupError("key not found")) is False + assert _should_retry_on(policy, NameError("name not defined")) is False + assert _should_retry_on(policy, SyntaxError("invalid syntax")) is False + assert _should_retry_on(policy, RuntimeError("runtime error")) is False + assert _should_retry_on(policy, ReferenceError("weak reference")) is False + assert _should_retry_on(policy, StopIteration()) is False + assert _should_retry_on(policy, StopAsyncIteration()) is False + assert _should_retry_on(policy, OSError("file not found")) is False + + # Should retry on httpx.HTTPStatusError with 5xx status code + response_5xx = Mock() + response_5xx.status_code = 503 + http_error_5xx = httpx.HTTPStatusError( + "server error", request=Mock(), response=response_5xx + ) + assert _should_retry_on(policy, http_error_5xx) is True + + # Should not retry on httpx.HTTPStatusError with 4xx status code + response_4xx = Mock() + response_4xx.status_code = 404 + http_error_4xx = httpx.HTTPStatusError( + "not found", request=Mock(), response=response_4xx + ) + assert _should_retry_on(policy, http_error_4xx) is False + + # Should retry on requests.HTTPError with 5xx status code + response_req_5xx = Mock() + response_req_5xx.status_code = 502 + req_error_5xx = requests.HTTPError("bad gateway") + req_error_5xx.response = response_req_5xx + assert _should_retry_on(policy, req_error_5xx) is True + + # Should not retry on requests.HTTPError with 4xx status code + response_req_4xx = Mock() + response_req_4xx.status_code = 400 + req_error_4xx = requests.HTTPError("bad request") + req_error_4xx.response = response_req_4xx + assert _should_retry_on(policy, req_error_4xx) is False + + # Should retry on requests.HTTPError with no response + req_error_no_resp = requests.HTTPError("connection error") + req_error_no_resp.response = None + assert _should_retry_on(policy, req_error_no_resp) is True + + # Should retry on other exceptions by default + class CustomException(Exception): + pass + + assert _should_retry_on(policy, CustomException("custom error")) is True + + +def test_graph_with_single_retry_policy(): + """Test a simple graph with a single RetryPolicy for a node.""" + + class State(TypedDict): + foo: str + + attempt_count = 0 + + def failing_node(state: State): + nonlocal attempt_count + attempt_count += 1 + if attempt_count < 3: # Fail the first two attempts + raise ValueError("Intentional failure") + return {"foo": "success"} + + def other_node(state: State): + return {"foo": "other_node"} + + # Create a retry policy with specific parameters + retry_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.01, # Short interval for tests + backoff_factor=2.0, + jitter=False, # Disable jitter for predictable timing + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, retry_policy=retry_policy) + .add_node("other_node", other_node) + .add_edge(START, "failing_node") + .add_edge("failing_node", "other_node") + .compile() + ) + + with patch("time.sleep") as mock_sleep: + result = graph.invoke({"foo": ""}) + + # Verify retry behavior + assert attempt_count == 3 # The node should have been tried 3 times + assert result["foo"] == "other_node" # Final result should be from other_node + + # Verify the sleep intervals + call_args_list = [args[0][0] for args in mock_sleep.call_args_list] + assert call_args_list == [0.01, 0.02] + + +def test_graph_with_jitter_retry_policy(): + """Test a graph with a RetryPolicy that uses jitter.""" + + class State(TypedDict): + foo: str + + attempt_count = 0 + + def failing_node(state): + nonlocal attempt_count + attempt_count += 1 + if attempt_count < 2: # Fail the first attempt + raise ValueError("Intentional failure") + return {"foo": "success"} + + # Create a retry policy with jitter enabled + retry_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.01, + jitter=True, # Enable jitter for randomized backoff + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("failing_node", failing_node, retry_policy=retry_policy) + .add_edge(START, "failing_node") + .compile() + ) + + # Test graph execution with mocked random and sleep + with ( + patch("random.uniform", return_value=0.05) as mock_random, + patch("time.sleep") as mock_sleep, + ): + result = graph.invoke({"foo": ""}) + + # Verify retry behavior + assert attempt_count == 2 # The node should have been tried twice + assert result["foo"] == "success" + + # Verify jitter was applied + mock_random.assert_called_with(0, 1) # Jitter should use random.uniform(0, 1) + mock_sleep.assert_called_with(0.01 + 0.05) # Sleep should include jitter + + +def test_graph_with_multiple_retry_policies(): + """Test a graph with multiple retry policies for a node.""" + + class State(TypedDict): + foo: str + error_type: str + + attempt_counts = {"value_error": 0, "key_error": 0} + + def failing_node(state): + error_type = state["error_type"] + + if error_type == "value_error": + attempt_counts["value_error"] += 1 + if attempt_counts["value_error"] < 2: + raise ValueError("Value error") + elif error_type == "key_error": + attempt_counts["key_error"] += 1 + if attempt_counts["key_error"] < 3: + raise KeyError("Key error") + + return {"foo": f"recovered_from_{error_type}"} + + # Create multiple retry policies + value_error_policy = RetryPolicy( + max_attempts=2, + initial_interval=0.01, + jitter=False, + retry_on=ValueError, + ) + + key_error_policy = RetryPolicy( + max_attempts=3, + initial_interval=0.02, + jitter=False, + retry_on=KeyError, + ) + + # Create and compile the graph with a list of retry policies + graph = ( + StateGraph(State) + .add_node( + "failing_node", + failing_node, + retry_policy=(value_error_policy, key_error_policy), + ) + .add_edge(START, "failing_node") + .compile() + ) + + # Test ValueError scenario + with patch("time.sleep"): + result_value_error = graph.invoke({"foo": "", "error_type": "value_error"}) + + assert attempt_counts["value_error"] == 2 + assert result_value_error["foo"] == "recovered_from_value_error" + + # Reset attempt counts + attempt_counts = {"value_error": 0, "key_error": 0} + + # Test KeyError scenario + with patch("time.sleep"): + result_key_error = graph.invoke({"foo": "", "error_type": "key_error"}) + + assert attempt_counts["key_error"] == 3 + assert result_key_error["foo"] == "recovered_from_key_error" + + +def test_graph_with_max_attempts_exceeded(): + """Test a graph where max_attempts is exceeded.""" + + class State(TypedDict): + foo: str + + def always_failing_node(state): + raise ValueError("Always fails") + + # Create a retry policy with limited attempts + retry_policy = RetryPolicy( + max_attempts=2, + initial_interval=0.01, + jitter=False, + retry_on=ValueError, + ) + + # Create and compile the graph + graph = ( + StateGraph(State) + .add_node("always_failing", always_failing_node, retry_policy=retry_policy) + .add_edge(START, "always_failing") + .compile() + ) + + # Test graph execution + with ( + patch("time.sleep") as mock_sleep, + pytest.raises(ValueError, match="Always fails"), + ): + graph.invoke({"foo": ""}) + + mock_sleep.assert_called_with(0.01) diff --git a/langgraph/source/libs/langgraph/tests/test_runnable.py b/langgraph/source/libs/langgraph/tests/test_runnable.py new file mode 100644 index 0000000000000000000000000000000000000000..e7d1950ce9d89df4dda04f1b5522722281bb7af0 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_runnable.py @@ -0,0 +1,414 @@ +from __future__ import annotations + +from typing import Any, Optional + +import pytest +from langchain_core.runnables.config import RunnableConfig +from langgraph.store.base import BaseStore + +from langgraph._internal._runnable import RunnableCallable +from langgraph.runtime import Runtime +from langgraph.types import StreamWriter + +pytestmark = pytest.mark.anyio + + +def test_runnable_callable_func_accepts(): + def sync_func(x: Any) -> str: + return f"{x}" + + async def async_func(x: Any) -> str: + return f"{x}" + + def func_with_store(x: Any, store: BaseStore) -> str: + return f"{x}" + + def func_with_writer(x: Any, writer: StreamWriter) -> str: + return f"{x}" + + async def afunc_with_store(x: Any, store: BaseStore) -> str: + return f"{x}" + + async def afunc_with_writer(x: Any, writer: StreamWriter) -> str: + return f"{x}" + + runnables = { + "sync": RunnableCallable(sync_func), + "async": RunnableCallable(func=None, afunc=async_func), + "with_store": RunnableCallable(func_with_store), + "with_writer": RunnableCallable(func_with_writer), + "awith_store": RunnableCallable(afunc_with_store), + "awith_writer": RunnableCallable(afunc_with_writer), + } + + expected_store = {"with_store": True, "awith_store": True} + expected_writer = {"with_writer": True, "awith_writer": True} + + for name, runnable in runnables.items(): + if expected_writer.get(name, False): + assert "writer" in runnable.func_accepts + else: + assert "writer" not in runnable.func_accepts + + if expected_store.get(name, False): + assert "store" in runnable.func_accepts + else: + assert "store" not in runnable.func_accepts + + +async def test_runnable_callable_basic(): + def sync_func(x: Any) -> str: + return f"{x}" + + async def async_func(x: Any) -> str: + return f"{x}" + + runnable_sync = RunnableCallable(sync_func) + runnable_async = RunnableCallable(func=None, afunc=async_func) + + result_sync = runnable_sync.invoke("test") + assert result_sync == "test" + + # Test asynchronous ainvoke + result_async = await runnable_async.ainvoke("test") + assert result_async == "test" + + +def test_runnable_callable_injectable_arguments() -> None: + """Test injectable arguments for RunnableCallable. + + This test verifies that injectable arguments like BaseStore work correctly. + It tests: + - Optional store injection + - Required store injection + - Store injection via config + - Store injection override behavior + - Store value injection and validation + """ + + # Test Optional[BaseStore] annotation. + def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: # noqa: UP045 + """Test function that accepts an optional store parameter.""" + assert store is None + return "success" + + assert ( + RunnableCallable(func_optional_store).invoke( + {"x": "1"}, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store=None, + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Test BaseStore annotation + def func_required_store(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter.""" + assert store is None + return "success" + + with pytest.raises(ValueError): + # Should fail b/c store is not Optional and config is not populated with store. + assert RunnableCallable(func_required_store).invoke({}) == "success" + + # Manually provide store + assert RunnableCallable(func_required_store).invoke({}, store=None) == "success" + + # Specify a value for store in the config + assert ( + RunnableCallable(func_required_store).invoke( + {}, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store=None, + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Specify a value for store in config, but override with None + assert ( + RunnableCallable(func_optional_store).invoke( + {"x": "1"}, + store=None, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", # type: ignore[assignment] + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Set of tests where we verify that 'foobar' is injected as the store value. + def func_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter and validates its value. + + The store value is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + assert ( + RunnableCallable(func_required_store_v2).invoke( + {}, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", # type: ignore[assignment] + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + assert RunnableCallable(func_required_store_v2).invoke( + # And manual override takes precedence. + {}, + store="foobar", + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", # type: ignore[assignment] + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + + +async def test_runnable_callable_injectable_arguments_async() -> None: + """Test injectable arguments for async RunnableCallable. + + This test verifies that injectable arguments like BaseStore work correctly + in the async context. It tests: + - Optional store injection + - Required store injection + - Store injection via config + - Store injection override behavior + """ + + # Test Optional[BaseStore] annotation. + def func_optional_store(inputs: Any, store: Optional[BaseStore]) -> str: # noqa: UP045 + """Test function that accepts an optional store parameter.""" + assert store is None + return "success" + + async def afunc_optional_store(inputs: Any, store: BaseStore | None) -> str: + """Async version of func_optional_store.""" + assert store is None + return "success" + + assert ( + await RunnableCallable( + func=func_optional_store, afunc=afunc_optional_store + ).ainvoke({"x": "1"}) + == "success" + ) + + # Test BaseStore annotation + def func_required_store(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter.""" + assert store is None + return "success" + + async def afunc_required_store(inputs: Any, store: BaseStore) -> str: + """Async version of func_required_store.""" + assert store is None + return "success" + + with pytest.raises(ValueError): + # Should fail b/c store is not Optional and config is not populated with store. + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke( + {}, + ) + == "success" + ) + + # Manually provide store + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke( + {}, + store=None, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store=None, + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Specify a value for store in the config + assert ( + await RunnableCallable( + func=func_required_store, afunc=afunc_required_store + ).ainvoke( + {}, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store=None, + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Specify a value for store in config, but override with None + assert ( + await RunnableCallable( + func=func_optional_store, afunc=afunc_optional_store + ).ainvoke( + {"x": "1"}, + store=None, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + # Set of tests where we verify that 'foobar' is injected as the store value. + def func_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Test function that requires a store parameter with specific value. + + The store parameter is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + async def afunc_required_store_v2(inputs: Any, store: BaseStore) -> str: + """Async version of func_required_store_v2. + + The store parameter is expected to be 'foobar' when injected. + """ + assert store == "foobar" + return "success" + + assert ( + await RunnableCallable( + func=func_required_store_v2, afunc=afunc_required_store_v2 + ).ainvoke( + {}, + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + assert ( + await RunnableCallable( + func=func_required_store_v2, afunc=afunc_required_store_v2 + ).ainvoke( + # And manual override takes precedence. + {}, + store="foobar", + config={ + "configurable": { + "__pregel_runtime": Runtime( + store="foobar", + context=None, + stream_writer=lambda _: None, + previous=None, + ) + } + }, + ) + == "success" + ) + + +def test_config_injection() -> None: + def func(x: Any, config: RunnableConfig) -> list[str]: + return config.get("tags", []) + + assert RunnableCallable(func).invoke( + "test", config={"tags": ["test"], "configurable": {}} + ) == ["test"] + + def func_optional(x: Any, config: Optional[RunnableConfig]) -> list[str]: # noqa: UP045 + return config.get("tags", []) if config else [] + + assert RunnableCallable(func_optional).invoke( + "test", config={"tags": ["test"], "configurable": {}} + ) == ["test"] + + def func_untyped(x: Any, config) -> list[str]: + return config.get("tags", []) + + assert RunnableCallable(func_untyped).invoke( + "test", config={"tags": ["test"], "configurable": {}} + ) == ["test"] + + +def test_config_ensured() -> None: + def func(input: str, config: RunnableConfig) -> None: + assert input == "test" + assert config is not None + assert config.get("configurable") is not None + + RunnableCallable(func).invoke("test") + + +async def test_config_ensured_async() -> None: + async def func(input: str, config: RunnableConfig) -> None: + assert input == "test" + assert config is not None + assert config.get("configurable") is not None + + await RunnableCallable(func).ainvoke("test") diff --git a/langgraph/source/libs/langgraph/tests/test_runtime.py b/langgraph/source/libs/langgraph/tests/test_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..0407b84d2d564937f2eb0ab82c075b430b828eb5 --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_runtime.py @@ -0,0 +1,391 @@ +from dataclasses import dataclass +from typing import Any + +import pytest +from pydantic import BaseModel, ValidationError +from typing_extensions import TypedDict + +from langgraph.graph import END, START, StateGraph +from langgraph.runtime import Runtime, get_runtime + + +def test_injected_runtime() -> None: + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + + def injected_runtime(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return {"message": f"api key: {runtime.context.api_key}"} + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("injected_runtime", injected_runtime) + graph.add_edge(START, "injected_runtime") + graph.add_edge("injected_runtime", END) + compiled = graph.compile() + result = compiled.invoke( + {"message": "hello world"}, context=Context(api_key="sk_123456") + ) + assert result == {"message": "api key: sk_123456"} + + +def test_context_runtime() -> None: + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + + def context_runtime(state: State) -> dict[str, Any]: + runtime = get_runtime(Context) + return {"message": f"api key: {runtime.context.api_key}"} + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("context_runtime", context_runtime) + graph.add_edge(START, "context_runtime") + graph.add_edge("context_runtime", END) + compiled = graph.compile() + result = compiled.invoke( + {"message": "hello world"}, context=Context(api_key="sk_123456") + ) + assert result == {"message": "api key: sk_123456"} + + +def test_override_runtime() -> None: + @dataclass + class Context: + api_key: str + + prev = Runtime(context=Context(api_key="abc")) + new = prev.override(context=Context(api_key="def")) + assert new.override(context=Context(api_key="def")).context.api_key == "def" + + +def test_merge_runtime() -> None: + @dataclass + class Context: + api_key: str + + runtime1 = Runtime(context=Context(api_key="abc")) + runtime2 = Runtime(context=Context(api_key="def")) + runtime3 = Runtime(context=None) + + assert runtime1.merge(runtime2).context.api_key == "def" + # override only applies to non-falsy values + assert runtime1.merge(runtime3).context.api_key == "abc" # type: ignore + + +def test_runtime_propogated_to_subgraph() -> None: + @dataclass + class Context: + username: str + + class State(TypedDict, total=False): + subgraph: str + main: str + + def subgraph_node_1(state: State, runtime: Runtime[Context]): + return {"subgraph": f"{runtime.context.username}!"} + + subgraph_builder = StateGraph(State, context_schema=Context) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.set_entry_point("subgraph_node_1") + subgraph = subgraph_builder.compile() + + def main_node(state: State, runtime: Runtime[Context]): + return {"main": f"{runtime.context.username}!"} + + builder = StateGraph(State, context_schema=Context) + builder.add_node(main_node) + builder.add_node("node_1", subgraph) + builder.set_entry_point("main_node") + builder.add_edge("main_node", "node_1") + graph = builder.compile() + + context = Context(username="Alice") + result = graph.invoke({}, context=context) + assert result == {"subgraph": "Alice!", "main": "Alice!"} + + +def test_context_coercion_dataclass() -> None: + """Test that dict context is coerced to dataclass.""" + + @dataclass + class Context: + api_key: str + timeout: int = 30 + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return { + "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test dict coercion with all fields + result = compiled.invoke( + {"message": "test"}, context={"api_key": "sk_test", "timeout": 60} + ) + assert result == {"message": "api_key: sk_test, timeout: 60"} + + # Test dict coercion with default field + result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"}) + assert result == {"message": "api_key: sk_test2, timeout: 30"} + + # Test with actual dataclass instance (should still work) + result = compiled.invoke( + {"message": "test"}, context=Context(api_key="sk_test3", timeout=90) + ) + assert result == {"message": "api_key: sk_test3, timeout: 90"} + + +def test_context_coercion_pydantic() -> None: + """Test that dict context is coerced to Pydantic model.""" + + class Context(BaseModel): + api_key: str + timeout: int = 30 + tags: list[str] = [] + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return { + "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}, tags: {runtime.context.tags}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test dict coercion with all fields + result = compiled.invoke( + {"message": "test"}, + context={"api_key": "sk_test", "timeout": 60, "tags": ["prod", "v2"]}, + ) + assert result == {"message": "api_key: sk_test, timeout: 60, tags: ['prod', 'v2']"} + + # Test dict coercion with defaults + result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"}) + assert result == {"message": "api_key: sk_test2, timeout: 30, tags: []"} + + # Test with actual Pydantic instance (should still work) + result = compiled.invoke( + {"message": "test"}, + context=Context(api_key="sk_test3", timeout=90, tags=["test"]), + ) + assert result == {"message": "api_key: sk_test3, timeout: 90, tags: ['test']"} + + +def test_context_coercion_typeddict() -> None: + """Test that dict context with TypedDict schema passes through as-is.""" + + class Context(TypedDict): + api_key: str + timeout: int + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + # TypedDict context is just a dict at runtime + return { + "message": f"api_key: {runtime.context['api_key']}, timeout: {runtime.context['timeout']}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test dict passes through for TypedDict + result = compiled.invoke( + {"message": "test"}, context={"api_key": "sk_test", "timeout": 60} + ) + assert result == {"message": "api_key: sk_test, timeout: 60"} + + +def test_context_coercion_none() -> None: + """Test that None context is handled properly.""" + + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + + def node_without_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + # Should be None when no context provided + return {"message": f"context is None: {runtime.context is None}"} + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_without_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test with None context + result = compiled.invoke({"message": "test"}, context=None) + assert result == {"message": "context is None: True"} + + # Test without context parameter (defaults to None) + result = compiled.invoke({"message": "test"}) + assert result == {"message": "context is None: True"} + + +def test_context_coercion_errors() -> None: + """Test error handling for invalid context.""" + + @dataclass + class Context: + api_key: str # Required field + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return {"message": "should not reach here"} + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test missing required field + with pytest.raises(TypeError): + compiled.invoke({"message": "test"}, context={"timeout": 60}) + + # Test invalid dict keys + with pytest.raises(TypeError): + compiled.invoke( + {"message": "test"}, context={"api_key": "test", "invalid_field": "value"} + ) + + +@pytest.mark.anyio +async def test_context_coercion_async() -> None: + """Test context coercion with async methods.""" + + @dataclass + class Context: + api_key: str + async_mode: bool = True + + class State(TypedDict): + message: str + + async def async_node(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return { + "message": f"async api_key: {runtime.context.api_key}, async_mode: {runtime.context.async_mode}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", async_node) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test dict coercion with ainvoke + result = await compiled.ainvoke( + {"message": "test"}, context={"api_key": "sk_async", "async_mode": False} + ) + assert result == {"message": "async api_key: sk_async, async_mode: False"} + + # Test dict coercion with astream + chunks = [] + async for chunk in compiled.astream( + {"message": "test"}, context={"api_key": "sk_stream"} + ): + chunks.append(chunk) + + # Find the chunk with our node output + node_output = None + for chunk in chunks: + if "node" in chunk: + node_output = chunk["node"] + break + + assert node_output == {"message": "async api_key: sk_stream, async_mode: True"} + + +def test_context_coercion_stream() -> None: + """Test context coercion with sync stream method.""" + + @dataclass + class Context: + api_key: str + stream_mode: str = "default" + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return { + "message": f"stream api_key: {runtime.context.api_key}, mode: {runtime.context.stream_mode}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + # Test dict coercion with stream + chunks = [] + for chunk in compiled.stream( + {"message": "test"}, context={"api_key": "sk_stream", "stream_mode": "fast"} + ): + chunks.append(chunk) + + # Find the chunk with our node output + node_output = None + for chunk in chunks: + if "node" in chunk: + node_output = chunk["node"] + break + + assert node_output == {"message": "stream api_key: sk_stream, mode: fast"} + + +def test_context_coercion_pydantic_validation_errors() -> None: + """Test that Pydantic validation errors are raised.""" + + class Context(BaseModel): + api_key: str + timeout: int + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return { + "message": f"api_key: {runtime.context.api_key}, timeout: {runtime.context.timeout}" + } + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + + compiled = graph.compile() + + with pytest.raises(ValidationError): + compiled.invoke( + {"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"} + ) diff --git a/langgraph/source/libs/langgraph/tests/test_state.py b/langgraph/source/libs/langgraph/tests/test_state.py new file mode 100644 index 0000000000000000000000000000000000000000..f1ff21a82ac5d8ace73f2af40fc73238a73e667c --- /dev/null +++ b/langgraph/source/libs/langgraph/tests/test_state.py @@ -0,0 +1,373 @@ +import inspect +import operator +import warnings +from dataclasses import dataclass, field +from typing import Annotated, Any, Union +from typing import Annotated as Annotated2 + +import pytest +from langchain_core.runnables import RunnableConfig +from pydantic import BaseModel +from typing_extensions import NotRequired, Required, TypedDict + +from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.ephemeral_value import EphemeralValue +from langgraph.graph.state import ( + StateGraph, + _get_node_name, + _is_field_channel, + _warn_invalid_state_schema, +) + + +class State(BaseModel): + foo: str + bar: int + + +class State2(TypedDict): + foo: str + bar: int + + +@pytest.mark.parametrize( + "schema", + [ + {"foo": "bar"}, + ["hi", lambda x, y: x + y], + State(foo="bar", bar=1), + State2(foo="bar", bar=1), + ], +) +def test_warns_invalid_schema(schema: Any): + with pytest.warns(UserWarning): + _warn_invalid_state_schema(schema) + + +@pytest.mark.parametrize( + "schema", + [ + Annotated[dict, lambda x, y: y], + Annotated2[list, lambda x, y: y], + dict, + State, + State2, + ], +) +def test_doesnt_warn_valid_schema(schema: Any): + # Assert the function does not raise a warning + with warnings.catch_warnings(): + warnings.simplefilter("error") + _warn_invalid_state_schema(schema) + + +def test_state_schema_with_type_hint(): + class InputState(TypedDict): + question: str + + class OutputState(TypedDict): + input_state: InputState + + class FooState(InputState): + foo: str + + def complete_hint(state: InputState) -> OutputState: + return {"input_state": state} + + def miss_first_hint(state, config: RunnableConfig) -> OutputState: + return {"input_state": state} + + def only_return_hint(state, config) -> OutputState: + return {"input_state": state} + + def miss_all_hint(state, config): + return {"input_state": state} + + def pre_foo(_) -> FooState: + return {"foo": "bar"} + + def pre_bar(_) -> FooState: + return {"foo": "bar"} + + class Foo: + def __call__(self, state: FooState) -> OutputState: + assert state.pop("foo") == "bar" + return {"input_state": state} + + class Bar: + def my_node(self, state: FooState) -> OutputState: + assert state.pop("foo") == "bar" + return {"input_state": state} + + graph = StateGraph(InputState, output_schema=OutputState) + actions = [ + complete_hint, + miss_first_hint, + only_return_hint, + miss_all_hint, + pre_foo, + Foo(), + pre_bar, + Bar().my_node, + ] + + for action in actions: + graph.add_node(action) + + def get_name(action) -> str: + return getattr(action, "__name__", action.__class__.__name__) + + graph.set_entry_point(get_name(actions[0])) + for i in range(len(actions) - 1): + graph.add_edge(get_name(actions[i]), get_name(actions[i + 1])) + graph.set_finish_point(get_name(actions[-1])) + + graph = graph.compile() + + input_state = InputState(question="Hello World!") + output_state = OutputState(input_state=input_state) + foo_state = FooState(foo="bar") + for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): + node_name = get_name(actions[i]) + if node_name in {"pre_foo", "pre_bar"}: + assert c[node_name] == foo_state + else: + assert c[node_name] == output_state + + +@pytest.mark.parametrize("total_", [True, False]) +def test_state_schema_optional_values(total_: bool): + class SomeParentState(TypedDict): + val0a: str + val0b: str | None + + class InputState(SomeParentState, total=total_): # type: ignore + val1: str + val2: str | None + val3: Required[Annotated[dict, operator.or_]] + val4: NotRequired[dict] + val5: Annotated[Required[str], "foo"] + val6: Annotated[NotRequired[str], "bar"] + + class OutputState(SomeParentState, total=total_): # type: ignore + out_val1: str + out_val2: str | None + out_val3: Required[str] + out_val4: NotRequired[dict] + out_val5: Annotated[Required[str], "foo"] + out_val6: Annotated[NotRequired[str], "bar"] + + class State(InputState): # this would be ignored + val4: dict + + builder = StateGraph(State, input_schema=InputState, output_schema=OutputState) + builder.add_node("n", lambda x: x) + builder.add_edge("__start__", "n") + graph = builder.compile() + json_schema = graph.get_input_jsonschema() + + assert isinstance(graph.channels["val3"], BinaryOperatorAggregate) + + if total_ is False: + expected_required = set() + expected_optional = {"val2", "val1"} + else: + expected_required = {"val1", "val2"} + expected_optional = set() + + # The others should always have precedence based on the required annotation + expected_required |= {"val0a", "val0b", "val3", "val5"} + expected_optional |= {"val4", "val6"} + + assert set(json_schema.get("required", set())) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) + + # Check output schema. Should be the same process + output_schema = graph.get_output_jsonschema() + if total_ is False: + expected_required = set() + expected_optional = {"out_val2", "out_val1"} + else: + expected_required = {"out_val1", "out_val2"} + expected_optional = set() + + expected_required |= {"val0a", "val0b", "out_val3", "out_val5"} + expected_optional |= {"out_val4", "out_val6"} + + assert set(output_schema.get("required", set())) == expected_required + assert ( + set(output_schema["properties"].keys()) == expected_required | expected_optional + ) + + +@pytest.mark.parametrize("kw_only_", [False, True]) +def test_state_schema_default_values(kw_only_: bool): + kwargs = {} + if "kw_only" in inspect.signature(dataclass).parameters: + kwargs = {"kw_only": kw_only_} + + @dataclass(**kwargs) + class InputState: + val1: str + val2: int | None + val3: Annotated[float | None, "optional annotated"] + val4: str | None = None + val5: list[int] = field(default_factory=lambda: [1, 2, 3]) + val6: dict[str, int] = field(default_factory=lambda: {"a": 1}) + val7: str = field(default=...) + val8: Annotated[int, "some metadata"] = 42 + val9: Annotated[str, "more metadata"] = field(default="some foo") + val10: str = "default" + val11: Annotated[list[str], "annotated list"] = field( + default_factory=lambda: ["a", "b"] + ) + + builder = StateGraph(InputState) + builder.add_node("n", lambda x: x) + builder.add_edge("__start__", "n") + graph = builder.compile() + for json_schema in [graph.get_input_jsonschema(), graph.get_output_jsonschema()]: + expected_required = {"val1", "val7"} + expected_optional = { + "val2", + "val3", + "val4", + "val5", + "val6", + "val8", + "val9", + "val10", + "val11", + } + + assert set(json_schema.get("required", set())) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) + + +def test__get_node_name() -> None: + # lambda + assert _get_node_name(lambda x: x) == "