__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[postgres] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[postgres_pipe] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres_pipe].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[postgres_pool] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[postgres_pool].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_branch_then[sqlite] + ''' + graph TD; + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + + ''' +# --- +# name: test_branch_then[sqlite].1 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[memory] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[memory].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[memory].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[memory].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[memory].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[memory].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[memory].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pipe] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pipe].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pipe].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pipe].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres_pipe].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pool] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pool].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pool].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[postgres_pool].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pool].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[postgres_pool].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[postgres_pool].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[sqlite] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[sqlite].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[sqlite].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[sqlite].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[sqlite].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[sqlite].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[sqlite].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_state_graph[memory] + '{"$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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_conditional_state_graph[memory].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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_conditional_state_graph[memory].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[memory].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[postgres] + '{"$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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres].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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[postgres_pipe] + '{"$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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres_pipe].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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres_pipe].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres_pipe].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[postgres_pool] + '{"$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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres_pool].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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_conditional_state_graph[postgres_pool].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[postgres_pool].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_state_graph[sqlite] + '{"$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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_conditional_state_graph[sqlite].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"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["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": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"default": null, "title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_conditional_state_graph[sqlite].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_state_graph[sqlite].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_message_graph[memory] + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphInput", "type": "array"}' +# --- +# name: test_message_graph[memory].1 + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphOutput", "type": "array"}' +# --- +# name: test_message_graph[memory].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_large_cases", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[memory].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. end .-> __end__; + + ''' +# --- +# name: test_message_graph[postgres] + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphInput", "type": "array"}' +# --- +# name: test_message_graph[postgres].1 + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphOutput", "type": "array"}' +# --- +# name: test_message_graph[postgres].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_large_cases", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. end .-> __end__; + + ''' +# --- +# name: test_message_graph[postgres_pipe] + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphInput", "type": "array"}' +# --- +# name: test_message_graph[postgres_pipe].1 + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphOutput", "type": "array"}' +# --- +# name: test_message_graph[postgres_pipe].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_large_cases", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres_pipe].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. end .-> __end__; + + ''' +# --- +# name: test_message_graph[postgres_pool] + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphInput", "type": "array"}' +# --- +# name: test_message_graph[postgres_pool].1 + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphOutput", "type": "array"}' +# --- +# name: test_message_graph[postgres_pool].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_large_cases", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[postgres_pool].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. end .-> __end__; + + ''' +# --- +# name: test_message_graph[sqlite] + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphInput", "type": "array"}' +# --- +# name: test_message_graph[sqlite].1 + '{"$defs": {"AIMessage": {"additionalProperties": true, "description": "Message from an AI.\\n\\nAIMessage is returned from a chat model as a response to a prompt.\\n\\nThis message represents the output of the model and consists of both\\nthe raw output as returned by the model together standardized fields\\n(e.g., tool calls, usage metadata) added by the LangChain framework.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ai", "default": "ai", "enum": ["ai"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}}, "required": ["content"], "title": "AIMessage", "type": "object"}, "AIMessageChunk": {"additionalProperties": true, "description": "Message chunk from an AI.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "AIMessageChunk", "default": "AIMessageChunk", "enum": ["AIMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}, "tool_calls": {"default": [], "items": {"$ref": "#/$defs/ToolCall"}, "title": "Tool Calls", "type": "array"}, "invalid_tool_calls": {"default": [], "items": {"$ref": "#/$defs/InvalidToolCall"}, "title": "Invalid Tool Calls", "type": "array"}, "usage_metadata": {"anyOf": [{"$ref": "#/$defs/UsageMetadata"}, {"type": "null"}], "default": null}, "tool_call_chunks": {"default": [], "items": {"$ref": "#/$defs/ToolCallChunk"}, "title": "Tool Call Chunks", "type": "array"}}, "required": ["content"], "title": "AIMessageChunk", "type": "object"}, "ChatMessage": {"additionalProperties": true, "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "chat", "default": "chat", "enum": ["chat"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessage", "type": "object"}, "ChatMessageChunk": {"additionalProperties": true, "description": "Chat Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ChatMessageChunk", "default": "ChatMessageChunk", "enum": ["ChatMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"], "title": "ChatMessageChunk", "type": "object"}, "FunctionMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nFunctionMessage are an older version of the ToolMessage schema, and\\ndo not contain the tool_call_id field.\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "function", "default": "function", "enum": ["function"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessage", "type": "object"}, "FunctionMessageChunk": {"additionalProperties": true, "description": "Function Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "FunctionMessageChunk", "default": "FunctionMessageChunk", "enum": ["FunctionMessageChunk"], "title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "name"], "title": "FunctionMessageChunk", "type": "object"}, "HumanMessage": {"additionalProperties": true, "description": "Message from a human.\\n\\nHumanMessages are messages that are passed in from a human to the model.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Instantiate a chat model and invoke it with the messages\\n model = ...\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "human", "default": "human", "enum": ["human"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessage", "type": "object"}, "HumanMessageChunk": {"additionalProperties": true, "description": "Human Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "HumanMessageChunk", "default": "HumanMessageChunk", "enum": ["HumanMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "example": {"default": false, "title": "Example", "type": "boolean"}}, "required": ["content"], "title": "HumanMessageChunk", "type": "object"}, "InputTokenDetails": {"description": "Breakdown of input token counts.\\n\\nDoes *not* need to sum to full input token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "cache_creation": {"title": "Cache Creation", "type": "integer"}, "cache_read": {"title": "Cache Read", "type": "integer"}}, "title": "InputTokenDetails", "type": "object"}, "InvalidToolCall": {"description": "Allowance for errors made by LLM.\\n\\nHere we add an `error` key to surface errors made during generation\\n(e.g., invalid JSON arguments.)", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "error": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Error"}, "type": {"const": "invalid_tool_call", "enum": ["invalid_tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "error"], "title": "InvalidToolCall", "type": "object"}, "OutputTokenDetails": {"description": "Breakdown of output token counts.\\n\\nDoes *not* need to sum to full output token count. Does *not* need to have all keys.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n\\n.. versionadded:: 0.3.9", "properties": {"audio": {"title": "Audio", "type": "integer"}, "reasoning": {"title": "Reasoning", "type": "integer"}}, "title": "OutputTokenDetails", "type": "object"}, "SystemMessage": {"additionalProperties": true, "description": "Message for priming AI behavior.\\n\\nThe system message is usually passed in as the first of a sequence\\nof input messages.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import HumanMessage, SystemMessage\\n\\n messages = [\\n SystemMessage(\\n content=\\"You are a helpful assistant! Your name is Bob.\\"\\n ),\\n HumanMessage(\\n content=\\"What is your name?\\"\\n )\\n ]\\n\\n # Define a chat model and invoke it with the messages\\n print(model.invoke(messages))", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "system", "default": "system", "enum": ["system"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessage", "type": "object"}, "SystemMessageChunk": {"additionalProperties": true, "description": "System Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "SystemMessageChunk", "default": "SystemMessageChunk", "enum": ["SystemMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content"], "title": "SystemMessageChunk", "type": "object"}, "ToolCall": {"description": "Represents a request to call a tool.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"name\\": \\"foo\\",\\n \\"args\\": {\\"a\\": 1},\\n \\"id\\": \\"123\\"\\n }\\n\\n This represents a request to call the tool named \\"foo\\" with arguments {\\"a\\": 1}\\n and an identifier of \\"123\\".", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "type": {"const": "tool_call", "enum": ["tool_call"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id"], "title": "ToolCall", "type": "object"}, "ToolCallChunk": {"description": "A chunk of a tool call (e.g., as part of a stream).\\n\\nWhen merging ToolCallChunks (e.g., via AIMessageChunk.__add__),\\nall string attributes are concatenated. Chunks are only merged if their\\nvalues of `index` are equal and not None.\\n\\nExample:\\n\\n.. code-block:: python\\n\\n left_chunks = [ToolCallChunk(name=\\"foo\\", args=\'{\\"a\\":\', index=0)]\\n right_chunks = [ToolCallChunk(name=None, args=\'1}\', index=0)]\\n\\n (\\n AIMessageChunk(content=\\"\\", tool_call_chunks=left_chunks)\\n + AIMessageChunk(content=\\"\\", tool_call_chunks=right_chunks)\\n ).tool_call_chunks == [ToolCallChunk(name=\'foo\', args=\'{\\"a\\":1}\', index=0)]", "properties": {"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"}, "args": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Args"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Id"}, "index": {"anyOf": [{"type": "integer"}, {"type": "null"}], "title": "Index"}, "type": {"const": "tool_call_chunk", "enum": ["tool_call_chunk"], "title": "Type", "type": "string"}}, "required": ["name", "args", "id", "index"], "title": "ToolCallChunk", "type": "object"}, "ToolMessage": {"additionalProperties": true, "description": "Message for passing the result of executing a tool back to a model.\\n\\nToolMessages contain the result of a tool invocation. Typically, the result\\nis encoded inside the `content` field.\\n\\nExample: A ToolMessage representing a result of 42 from a tool call with id\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n ToolMessage(content=\'42\', tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\')\\n\\n\\nExample: A ToolMessage where only part of the tool output is sent to the model\\n and the full output is passed in to artifact.\\n\\n .. versionadded:: 0.2.17\\n\\n .. code-block:: python\\n\\n from langchain_core.messages import ToolMessage\\n\\n tool_output = {\\n \\"stdout\\": \\"From the graph we can see that the correlation between x and y is ...\\",\\n \\"stderr\\": None,\\n \\"artifacts\\": {\\"type\\": \\"image\\", \\"base64_data\\": \\"/9j/4gIcSU...\\"},\\n }\\n\\n ToolMessage(\\n content=tool_output[\\"stdout\\"],\\n artifact=tool_output,\\n tool_call_id=\'call_Jja7J89XsjrOLA5r!MEOW!SL\',\\n )\\n\\nThe tool_call_id field is used to associate the tool call request with the\\ntool call response. This is useful in situations where a chat model is able\\nto request multiple tool calls in parallel.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "tool", "default": "tool", "enum": ["tool"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessage", "type": "object"}, "ToolMessageChunk": {"additionalProperties": true, "description": "Tool Message chunk.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"const": "ToolMessageChunk", "default": "ToolMessageChunk", "enum": ["ToolMessageChunk"], "title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"default": null, "title": "Artifact"}, "status": {"default": "success", "enum": ["success", "error"], "title": "Status", "type": "string"}}, "required": ["content", "tool_call_id"], "title": "ToolMessageChunk", "type": "object"}, "UsageMetadata": {"description": "Usage metadata for a message, such as token counts.\\n\\nThis is a standard representation of token usage that is consistent across models.\\n\\nExample:\\n\\n .. code-block:: python\\n\\n {\\n \\"input_tokens\\": 350,\\n \\"output_tokens\\": 240,\\n \\"total_tokens\\": 590,\\n \\"input_token_details\\": {\\n \\"audio\\": 10,\\n \\"cache_creation\\": 200,\\n \\"cache_read\\": 100,\\n },\\n \\"output_token_details\\": {\\n \\"audio\\": 10,\\n \\"reasoning\\": 200,\\n }\\n }\\n\\n.. versionchanged:: 0.3.9\\n\\n Added ``input_token_details`` and ``output_token_details``.", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}, "input_token_details": {"$ref": "#/$defs/InputTokenDetails"}, "output_token_details": {"$ref": "#/$defs/OutputTokenDetails"}}, "required": ["input_tokens", "output_tokens", "total_tokens"], "title": "UsageMetadata", "type": "object"}}, "default": null, "items": {"oneOf": [{"$ref": "#/$defs/AIMessage"}, {"$ref": "#/$defs/HumanMessage"}, {"$ref": "#/$defs/ChatMessage"}, {"$ref": "#/$defs/SystemMessage"}, {"$ref": "#/$defs/FunctionMessage"}, {"$ref": "#/$defs/ToolMessage"}, {"$ref": "#/$defs/AIMessageChunk"}, {"$ref": "#/$defs/HumanMessageChunk"}, {"$ref": "#/$defs/ChatMessageChunk"}, {"$ref": "#/$defs/SystemMessageChunk"}, {"$ref": "#/$defs/FunctionMessageChunk"}, {"$ref": "#/$defs/ToolMessageChunk"}]}, "title": "LangGraphOutput", "type": "array"}' +# --- +# name: test_message_graph[sqlite].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_large_cases", + "FakeFuntionChatModel" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "end", + "conditional": true + } + ] + } + ''' +# --- +# name: test_message_graph[sqlite].3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. end .-> __end__; + + ''' +# --- +# name: test_prebuilt_tool_chat + '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}}, "required": ["messages"], "title": "LangGraphInput", "type": "object"}' +# --- +# name: test_prebuilt_tool_chat.1 + '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}}, "required": ["messages"], "title": "LangGraphOutput", "type": "object"}' +# --- +# name: test_prebuilt_tool_chat.2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "prebuilt", + "tool_node", + "ToolNode" + ], + "name": "tools" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "conditional": true + } + ] + } + ''' +# --- +# name: test_prebuilt_tool_chat.3 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -.-> tools; + agent -.-> __end__; + + ''' +# --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_start_branch_then[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 1cfdd9b76..a44711ab3 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -1,230 +1,4 @@ # serializer version: 1 -# name: test_branch_then.1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[duckdb] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[duckdb].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[memory] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[memory].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[postgres] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[postgres].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[postgres_pipe] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[postgres_pipe].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[postgres_pool] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[postgres_pool].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- -# name: test_branch_then[sqlite] - ''' - graph TD; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - - ''' -# --- -# name: test_branch_then[sqlite].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__([__start__
]):::first - prepare(prepare) - tool_two_slow(tool_two_slow) - tool_two_fast(tool_two_fast) - finish(finish) - __end__([__end__
]):::last - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef default fill:#f2f0ff,line-height:1.2 - classDef first fill-opacity:0 - classDef last fill:#bfb6fc - - ''' -# --- # name: test_conditional_entrypoint_graph '{"title": "LangGraphInput"}' # --- diff --git a/libs/langgraph/tests/agents.py b/libs/langgraph/tests/agents.py new file mode 100644 index 000000000..6e5ce595a --- /dev/null +++ b/libs/langgraph/tests/agents.py @@ -0,0 +1,43 @@ +from typing import Literal, Union + +from pydantic import BaseModel + + +# define these objects to avoid importing langchain_core.agents +# and therefore avoid relying on core Pydantic version +class AgentAction(BaseModel): + tool: str + tool_input: Union[str, dict] + log: str + type: Literal["AgentAction"] = "AgentAction" + + model_config = { + "json_schema_extra": { + "description": ( + """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.""" + ) + } + } + + +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" + model_config = { + "json_schema_extra": { + "description": ( + """Final return value of an ActionAgent. + +Agents return an AgentFinish when they have reached a stopping condition.""" + ) + } + } diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py new file mode 100644 index 000000000..ef90f0b3c --- /dev/null +++ b/libs/langgraph/tests/test_large_cases.py @@ -0,0 +1,10019 @@ +import json +import operator +import re +import time +from contextlib import contextmanager +from dataclasses import replace +from typing import Annotated, Any, Iterator, Literal, Optional, TypedDict, Union, cast + +import httpx +import pytest +from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick +from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion + +from langgraph.channels.context import Context +from langgraph.channels.last_value import LastValue +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START +from langgraph.errors import NodeInterrupt +from langgraph.graph import StateGraph +from langgraph.graph.graph import Graph +from langgraph.graph.message import MessageGraph, MessagesState, add_messages +from langgraph.managed.shared_value import SharedValue +from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode +from langgraph.pregel import Channel, Pregel +from langgraph.store.memory import InMemoryStore +from langgraph.types import ( + Command, + Interrupt, + PregelTask, + RetryPolicy, + Send, + StateSnapshot, + StreamWriter, + interrupt, +) +from tests.agents import AgentAction, AgentFinish +from tests.any_str import AnyDict, AnyStr, UnsortedSequence +from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS +from tests.fake_chat import FakeChatModel +from tests.fake_tracer import FakeTracer +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_invoke_two_processes_in_out_interrupt( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.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=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) is None + + # inbox == 3 + checkpoint = checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert app.invoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert app.invoke(20, thread1) is None + + # inbox == 21 + checkpoint = 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) is None + assert app.invoke(None, thread1) == 5 + + # start execution again, stopping at inbox + assert app.invoke(20, thread2) 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 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, + "writes": {"two": 5}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + 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, + "writes": {"input": 3}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + 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, + "writes": {"input": 20}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + 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, + "writes": {"two": 4}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[7].config, + ), + 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, + "writes": {"input": 2}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # 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__": ()}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_fork_always_re_runs_nodes( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + 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=checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + + # start execution, stop at inbox + assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ + ("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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + 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, + "writes": None, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + 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, + "writes": {"__start__": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # 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}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_conditional_graph( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + from langchain_core.language_models.fake import FakeStreamingListLLM + from langchain_core.prompts import PromptTemplate + from langchain_core.runnables import RunnablePassthrough + from langchain_core.tools import tool + + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + # 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) -> Union[AgentAction, AgentFinish]: + if input.startswith("finish"): + _, answer = input.split(":") + return AgentFinish(return_values={"answer": answer}, log=input) + else: + _, tool_name, tool_input = input.split(":") + return AgentAction(tool=tool_name, tool_input=tool_input, log=input) + + agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) + + # Define tool execution logic + def execute_tools(data: dict) -> dict: + data = data.copy() + agent_action: AgentAction = data.pop("agent_outcome") + observation = {t.name: t for t in tools}[agent_action.tool].invoke( + agent_action.tool_input + ) + if data.get("intermediate_steps") is None: + data["intermediate_steps"] = [] + else: + data["intermediate_steps"] = data["intermediate_steps"].copy() + data["intermediate_steps"].append([agent_action, observation]) + return data + + # Define decision-making logic + def should_continue(data: dict) -> 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 = Graph() + + workflow.add_node("agent", agent) + workflow.add_node( + "tools", + execute_tools, + metadata={"parents": {}, "version": 2, "variant": "b"}, + ) + + 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 SHOULD_CHECK_SNAPSHOTS: + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert app.get_graph().draw_mermaid() == snapshot + assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot + assert app.get_graph(xray=True).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 [c for c in app.stream({"input": "what is weather in sf"})] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "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": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] + + # test state get/update methods with interrupt_after + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + if SHOULD_CHECK_SNAPSHOTS: + assert app_w_interrupt.get_graph().to_json() == snapshot + assert app_w_interrupt.get_graph().draw_mermaid() == snapshot + + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + config=app_w_interrupt.checkpointer.get_tuple(config).config, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + assert ( + app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ + "checkpoint_id" + ] + is not None + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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:a different query", + ), + "input": "what is weather in sf", + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "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": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + }, + tasks=(), + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + 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) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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:a different query", + ), + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "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": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + }, + tasks=(), + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # test re-invoke to continue with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + 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) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + "thread_id": "3", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "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": { + "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" + ), + } + }, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_conditional_state_graph( + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, +) -> None: + from langchain_core.language_models.fake import FakeStreamingListLLM + from langchain_core.prompts import PromptTemplate + from langchain_core.tools import tool + + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + setup = mocker.Mock() + teardown = mocker.Mock() + + @contextmanager + def assert_ctx_once() -> Iterator[None]: + assert setup.call_count == 0 + assert teardown.call_count == 0 + try: + yield + finally: + assert setup.call_count == 1 + assert teardown.call_count == 1 + setup.reset_mock() + teardown.reset_mock() + + @contextmanager + def make_httpx_client() -> Iterator[httpx.Client]: + setup() + with httpx.Client() as client: + try: + yield client + finally: + teardown() + + class AgentState(TypedDict, total=False): + input: Annotated[str, UntrackedValue] + agent_outcome: Optional[Union[AgentAction, AgentFinish]] + intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] + session: Annotated[httpx.Client, Context(make_httpx_client)] + + class ToolState(TypedDict, total=False): + agent_outcome: Union[AgentAction, AgentFinish] + session: Annotated[httpx.Client, Context(make_httpx_client)] + + # 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, Union[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 isinstance(data["session"], httpx.Client) + 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 + assert isinstance(data["session"], httpx.Client) + # 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=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 SHOULD_CHECK_SNAPSHOTS: + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + + with assert_ctx_once(): + 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" + ), + } + + with assert_ctx_once(): + 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=checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + with assert_ctx_once(): + assert [ + c + for c in app_w_interrupt.stream({"input": "what is weather in sf"}, 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + with assert_ctx_once(): + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + with assert_ctx_once(): + 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__": ()}, + ] + + with assert_ctx_once(): + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=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) + ] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # test w interrupt before all + app_w_interrupt = workflow.compile( + checkpointer=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) + ] == [ + {"__interrupt__": ()}, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), + next=("agent",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "3", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "3", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + "thread_id": "3", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=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) + ] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "4", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + "thread_id": "4", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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) + + if SHOULD_CHECK_SNAPSHOTS: + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == 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"), + ] + } + + assert [ + c + for c in app.stream( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="messages", + ) + ] == [ + ( + _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", + } + ], + ), + { + "langgraph_step": 1, + "langgraph_node": "agent", + "langgraph_triggers": ["start: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": ["branch:agent:should_continue:tools"], + "langgraph_path": (PULL, "tools"), + "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", + }, + ], + ), + { + "langgraph_step": 3, + "langgraph_node": "agent", + "langgraph_triggers": ["tools"], + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ( + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_path": (PULL, "tools"), + "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": ["branch:agent:should_continue:tools"], + "langgraph_path": (PULL, "tools"), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdAIMessageChunk( + content="answer", + ), + { + "langgraph_step": 5, + "langgraph_node": "agent", + "langgraph_triggers": ["tools"], + "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 + + assert ( + app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="updates", + )[0]["agent"]["messages"] + == [ + { + "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"}, + }, + ], + ) + ] + } + }, + { + "tools": { + "messages": [ + _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", + ), + ] + } + }, + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, + ][0]["agent"]["messages"] + ) + + assert [ + *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) + ] == [ + { + "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"}, + }, + ], + ) + ] + } + }, + { + "tools": { + "messages": [ + _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", + ), + ] + } + }, + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_state_graph_packets( + request: pytest.FixtureRequest, checkpointer_name: str, 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 + + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + session: Annotated[httpx.Client, Context(httpx.Client)] + + @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: + assert isinstance(data["session"], httpx.Client) + return { + "messages": model.invoke(data["messages"]), + "something_extra": "hi there", + } + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + assert isinstance(data["session"], httpx.Client) + 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=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 + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + if not FF_SEND_V2: + return + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + ), + next=("tools",), + config=(app_w_interrupt.checkpointer.get_tuple(config)).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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)),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + "something_extra": "hi there", + } + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + ), + next=("tools", "tools"), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + }, + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 3, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + } + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + ), + next=("tools",), + config=(app_w_interrupt.checkpointer.get_tuple(config)).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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)),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + "something_extra": "hi there", + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + ), + next=("tools", "tools"), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + }, + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 3, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_message_graph( + snapshot: SnapshotAssertion, + deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, +) -> 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 + + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + class FakeFuntionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: Optional[list[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = 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 = FakeFuntionChatModel( + 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 = MessageGraph() + + # 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 SHOULD_CHECK_SNAPSHOTS: + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == 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=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) + ] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + app_w_interrupt = workflow.compile( + checkpointer=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)] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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 BaseMessge 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_root_graph( + deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, +) -> 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 + + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + class FakeFuntionChatModel(FakeMessagesListChatModel): + def bind_functions(self, functions: list): + return self + + def _generate( + self, + messages: list[BaseMessage], + stop: Optional[list[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = 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 = FakeFuntionChatModel( + 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-000000000037", + ) + ] + }, + { + "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-000000000045", + ) + ] + }, + {"agent": AIMessage(content="answer", id="ai3")}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=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) + ] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + app_w_interrupt = workflow.compile( + checkpointer=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)] == [ + { + "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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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 BaseMessge 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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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=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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 6, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # 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-000000000078", + ), + 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-000000000100" + ), + 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": ["start: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": ["rewrite_query"], + }, + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ["rewrite_query"], + }, + }, + ), + ( + "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": ["retriever_one", "retriever_two"], + }, + }, + ), + ("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"], + }, + ), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_dynamic_interrupt( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=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", + } + 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=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...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, + ] + # 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) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + 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...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # 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=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + +@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_copy_checkpoint( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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[Union[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=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", + } + 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 one all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=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...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, + ] + # 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) == { + "my_key": "value ⛰️ one", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"tool_one": {"my_key": " one"}}, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"tool_one": {"my_key": " one"}}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None) + # interrupt is cleared, next task is kept + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=(), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_dynamic_interrupt_subgraph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=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", + } + 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=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...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ) + }, + ] + # 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) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [ + c.metadata + for c in tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + ) + ] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + 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...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[ + *tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 + ) + ][-1].config, + ) + # 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=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[ + *tool_two.checkpointer.list( + {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 + ) + ][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_start_branch_then( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] + + def assert_shared_value(data: State, config: RunnableConfig) -> State: + assert "shared" in data + if thread_id := config["configurable"].get("thread_id"): + if thread_id == "1": + # this is the first thread, so should not see a value + assert data["shared"] == {} + return {"shared": {"1": {"hello": "world"}}} + elif thread_id == "2": + # this should get value saved by thread 1 + assert data["shared"] == {"1": {"hello": "world"}} + elif thread_id == "3": + # this is a different assistant, so should not see previous value + assert data["shared"] == {} + return {} + + def tool_two_slow(data: State, config: RunnableConfig) -> State: + return {"my_key": " slow", **assert_shared_value(data, config)} + + def tool_two_fast(data: State, config: RunnableConfig) -> State: + return {"my_key": " fast", **assert_shared_value(data, config)} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", tool_two_slow) + tool_two_graph.add_node("tool_two_fast", tool_two_fast) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END + ) + tool_two = tool_two_graph.compile() + assert tool_two.get_graph().draw_mermaid() == snapshot + + assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { + "my_key": "value slow", + "market": "DE", + } + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value fast", + "market": "US", + } + + tool_two = tool_two_graph.compile( + store=InMemoryStore(), + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "assistant_id": "a", + "thread_id": "1", + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value ⛰️ slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ slow", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # update state + tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_branch_then( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.set_entry_point("prepare") + tool_two_graph.set_finish_point("finish") + tool_two_graph.add_conditional_edges( + source="prepare", + path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + then="finish", + ) + tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) + tool_two = tool_two_graph.compile() + assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot + assert tool_two.get_graph().draw_mermaid() == snapshot + + assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value prepared fast finished", + "market": "US", + } + + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + + res = [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] + + assert res == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "10", + }, + "parent_config": None, + "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["prepare"], + "tasks": [ + {"id": AnyStr(), "name": "prepare", "interrupts": (), "state": None} + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["finish"], + "tasks": [ + {"id": AnyStr(), "name": "finish", "interrupts": (), "state": None} + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": [], + "tasks": [], + }, + }, + ] + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "2", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "2", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["finish"] + ) + + thread1 = {"configurable": {"thread_id": "11"}} + + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared slow", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={ + "my_key": "value prepared slow", + "market": "DE", + }, + tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "11", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + # update state + tool_two.update_state(thread1, {"my_key": "er"}) + assert tool_two.get_state(thread1) == StateSnapshot( + values={ + "my_key": "value prepared slower", + "market": "DE", + }, + tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 3, + "writes": {"tool_two_slow": {"my_key": "er"}}, + "thread_id": "11", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "21", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "21", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "22", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert tool_two.get_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "22", + }, + parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = tool_two.update_state(thread3, {"my_key": "key", "market": "DE"}) + # check current state + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + "thread_id": "23", + }, + parent_config=None, + ) + # run from this point + assert tool_two.invoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "23", + }, + parent_config=uconfig, + ) + # resume, for same result as above + assert tool_two.invoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert tool_two.get_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "23", + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_dedupe_on_resume( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("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=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke(["0"], thread1, debug=1) == [ + "0", + "1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "2|3", + ] + 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) == 2 + # resume execution + assert graph.invoke(None, thread1, debug=1) == [ + "0", + "1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + # 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 history == [ + StateSnapshot( + values=[ + "0", + "1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"], "3.1": ["3.1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + ], + next=("3", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "1": ["1"], + "2": [ + ["2|Command(goto=Send(node='2', arg=3))"], + ["2|Command(goto=Send(node='flaky', arg=4))"], + ["2|3"], + ], + "flaky": ["flaky|4"], + }, + "thread_id": "1", + "step": 1, + "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"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + ), + StateSnapshot( + values=["0"], + next=("1", "2", "2", "2", "flaky"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "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"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(Interrupt(value="Bahh", when="during"),), + state=None, + result=["flaky|4"], + ), + ), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "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"], + ), + ), + ), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_nested_graph_state( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + 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=checkpointer) + + config = {"configurable": {"thread_id": "1"}} + app.invoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + assert app.get_state(config) == 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", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # 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", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + ), + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # get_state_history returns outer graph checkpoints + history = list(app.get_state_history(config)) + assert history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # get_state_history for a subgraph returns its checkpoints + child_history = [*app.get_state_history(history[0].tasks[0].state)] + assert 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", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + AnyStr(), + "inner_1", + (PULL, "inner_1"), + result={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + + # resume + app.invoke(None, config, debug=True) + # 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", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # 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", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=( + PregelTask( + AnyStr(), + "outer_2", + (PULL, "outer_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} + }, + result={"my_key": "hi my value here and there"}, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + 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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_doubly_nested_graph_state( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + 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=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)] == [ + ((), {"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", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + child_state = app.get_state(outer_state.tasks[0].state) + assert ( + child_state.tasks[0] + == 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(), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ).tasks[0] + ) + 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", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [PULL, AnyStr("child_1")], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + ) + # 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", + "writes": { + "grandchild_1": {"my_key": "hi my value here"} + }, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + ), + ), + ), + 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", + "writes": None, + "step": 0, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # # resume + assert [c for c in app.stream(None, config, subgraphs=True)] == [ + ( + (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", + "writes": { + "parent_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + ) + # 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", + "writes": { + "parent_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("parent_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="parent_2", + path=(PULL, "parent_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + result={"my_key": "hi my value here and there"}, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("parent_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="parent_1", + path=(PULL, "parent_1"), + result={"my_key": "hi my value"}, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + ), + ] + # get child graph history + child_history = list(app.get_state_history(outer_history[2].tasks[0].state)) + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=(), + ), + 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", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="child_1", + path=(PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + } + }, + result={"my_key": "hi my value here and there"}, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + # get grandchild graph history + grandchild_history = list(app.get_state_history(child_history[1].tasks[0].state)) + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + 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", + "writes": {"grandchild_2": {"my_key": "hi my value here and there"}}, + "step": 2, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=(), + ), + 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", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result={"my_key": "hi my value here and there"}, + ), + ), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("grandchild_1",), + 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", + "writes": None, + "step": 0, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_1", + path=(PULL, "grandchild_1"), + result={"my_key": "hi my value here"}, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + 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": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + + # replay grandchild checkpoint + assert [ + c for c in app.stream(None, grandchild_history[2].config, subgraphs=True) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_to_nested_graphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + 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=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=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) + + if not FF_SEND_V2: + # 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"]}}, + ] + return + + assert outer_state == StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + ), + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ) + # check state of each of the inner tasks + assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( + values={"subject": "cats - hohoho", "jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), + } + }, + metadata={ + "step": 1, + "source": "loop", + "writes": None, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_node": "generate_joke", + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1], + "langgraph_step": 0, + "langgraph_triggers": [PUSH], + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), + } + }, + tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), + ) + assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot( + values={"subject": "dogs - hohoho", "jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), + } + }, + metadata={ + "step": 1, + "source": "loop", + "writes": None, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_node": "generate_joke", + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2], + "langgraph_step": 0, + "langgraph_triggers": [PUSH], + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("generate_joke:"): AnyStr(), + } + ), + } + }, + tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), + ) + # update state of dogs joke graph + graph.update_state( + outer_state.tasks[2 if FF_SEND_V2 else 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"]}}, + ] + + actual_snapshot = graph.get_state(config) + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] + }, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = list(graph.get_state_history(config)) + + # get subgraph node state for expected history + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] + }, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"jokes": []}, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about cats - hohoho"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about turtles - hohoho"]}, + ), + ), + next=("__start__", "generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_react_interrupt( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=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=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + 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 + + if not FF_SEND_V2: + return + + # 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": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2), + 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": 1, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ) + } + }, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "3"}} + 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 + + # 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": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2), + 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": 1, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ) + } + }, + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", (), 0), + 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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_react_interrupt_control( + request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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() + 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=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 + + if not FF_SEND_V2: + return + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + 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 + + # 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": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2), + 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": 1, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ) + } + }, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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()) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_weather_subgraph( + request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion +) -> 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 + + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + # 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) + graph.add_edge("normal_llm_node", END) + graph.add_edge("weather_graph", END) + graph = graph.compile(checkpointer=checkpointer) + + 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")] == [ + "I'm", + " very", + ] + assert [c for c in graph.stream(None, thread2, stream_mode="custom")] == [ + " good", + ] + + # run until interrupt + assert [ + c + for c in graph.stream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + ) + + # 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 + ) + ] == [ + ((), {"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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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", + "writes": {"model_node": {"city": "San Francisco"}}, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + 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=( + PregelTask( + id=AnyStr(), + name="weather_node", + path=(PULL, "weather_node"), + ), + ), + ), + ), + ), + ) + 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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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", + "writes": { + "weather_node": { + "messages": [{"role": "assistant", "content": "rainy"}] + } + }, + "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + 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=(), + ), + ), + ), + ) + 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"), + ] + } + }, + ), + ] diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py new file mode 100644 index 000000000..7be4f6943 --- /dev/null +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -0,0 +1,7275 @@ +import asyncio +import operator +import re +import sys +from contextlib import asynccontextmanager +from typing import ( + Annotated, + Any, + AnyStr, + AsyncIterator, + Literal, + Optional, + TypedDict, + Union, + cast, +) + +import httpx +import pytest +from langchain_core.messages import ToolCall +from langchain_core.runnables import RunnableConfig, RunnablePick +from pydantic import BaseModel +from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion + +from langgraph.channels.last_value import LastValue +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START +from langgraph.graph.graph import Graph +from langgraph.graph.message import MessageGraph, add_messages +from langgraph.graph.state import StateGraph +from langgraph.managed.context import Context +from langgraph.managed.shared_value import SharedValue +from langgraph.prebuilt.chat_agent_executor import create_react_agent +from langgraph.prebuilt.tool_node import ToolNode +from langgraph.pregel import Channel, Pregel +from langgraph.store.memory import InMemoryStore +from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter +from tests.any_str import AnyDict +from tests.conftest import ALL_CHECKPOINTERS_ASYNC, awith_checkpointer +from tests.fake_chat import FakeChatModel +from tests.fake_tracer import FakeTracer +from tests.messages import ( + _AnyIdAIMessage, + _AnyIdAIMessageChunk, + _AnyIdHumanMessage, + _AnyIdToolMessage, +) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_invoke_two_processes_in_out_interrupt( + checkpointer_name: str, mocker: MockerFixture +) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=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) is None + + # inbox == 3 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None + + # inbox == 21 + checkpoint = await 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) is None + assert await app.ainvoke(None, thread1) == 5 + + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) 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 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, + "writes": {"two": 5}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + 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, + "writes": {"input": 3}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + 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, + "writes": {"input": 20}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + 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, + "writes": {"two": 4}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + 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, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[7].config, + ), + 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, + "writes": {"input": 2}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # 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__": ()}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_fork_always_re_runs_nodes( + checkpointer_name: str, 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) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=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"]) + ] == [ + ("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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + 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, + "writes": {"add_one": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + 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, + "writes": None, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + 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, + "writes": {"__start__": 1}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # 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}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_conditional_graph(checkpointer_name: str) -> 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.runnables import RunnablePassthrough + from langchain_core.tools import tool + + # 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", + ] + ) + + async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]: + if input.startswith("finish"): + _, answer = input.split(":") + return AgentFinish(return_values={"answer": answer}, log=input) + else: + _, tool_name, tool_input = input.split(":") + return AgentAction(tool=tool_name, tool_input=tool_input, log=input) + + agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) + + # Define tool execution logic + async def execute_tools(data: dict) -> dict: + data = data.copy() + agent_action: AgentAction = data.pop("agent_outcome") + observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke( + agent_action.tool_input + ) + if data.get("intermediate_steps") is None: + data["intermediate_steps"] = [] + else: + data["intermediate_steps"] = data["intermediate_steps"].copy() + data["intermediate_steps"].append([agent_action, observation]) + return data + + # Define decision-making logic + async def should_continue(data: dict, config: RunnableConfig) -> 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 = Graph() + + 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": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "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": { + "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" + ), + } + }, + ] + + 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 + assert "/logs/agent:2/streamed_output/-" in patch_paths + assert "/logs/agent:3/streamed_output/-" in patch_paths + # Check that agent (one of the nodes) has its final output set in the logs + assert "/logs/agent/final_output" in patch_paths + assert "/logs/agent:2/final_output" in patch_paths + assert "/logs/agent:3/final_output" in patch_paths + assert [ + p["value"] + for log in patches + for p in log.ops + if p["path"] == "/logs/agent/final_output" + or p["path"] == "/logs/agent:2/final_output" + or p["path"] == "/logs/agent:3/final_output" + ] == [ + { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + }, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + }, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after + + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "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": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + }, + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "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": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + }, + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test re-invoke to continue with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "3"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + "thread_id": "3", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "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": { + "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" + ), + } + }, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_conditional_graph_state( + mocker: MockerFixture, checkpointer_name: str +) -> 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 + + setup = mocker.Mock() + teardown = mocker.Mock() + + @asynccontextmanager + async def assert_ctx_once() -> AsyncIterator[None]: + assert setup.call_count == 0 + assert teardown.call_count == 0 + try: + yield + finally: + assert setup.call_count == 1 + assert teardown.call_count == 1 + setup.reset_mock() + teardown.reset_mock() + + class MyPydanticContextModel(BaseModel, arbitrary_types_allowed=True): + session: httpx.AsyncClient + something_else: str + + @asynccontextmanager + async def make_context( + config: RunnableConfig, + ) -> AsyncIterator[MyPydanticContextModel]: + assert isinstance(config, dict) + setup() + session = httpx.AsyncClient() + try: + yield MyPydanticContextModel(session=session, something_else="hello") + finally: + await session.aclose() + teardown() + + class AgentState(TypedDict): + input: Annotated[str, UntrackedValue] + agent_outcome: Optional[Union[AgentAction, AgentFinish]] + intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] + context: Annotated[MyPydanticContextModel, Context(make_context)] + + # 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, Union[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: + # check we have httpx session in AgentState + assert isinstance(data["context"], MyPydanticContextModel) + # 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 we have httpx session in AgentState + assert isinstance(data["context"], MyPydanticContextModel) + # 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() + + async with assert_ctx_once(): + 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" + ), + } + + async with assert_ctx_once(): + 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" + ), + } + }, + ] + + async with assert_ctx_once(): + 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" + ), + }, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + async with assert_ctx_once(): + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + 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__": ()}, + ] + + async with assert_ctx_once(): + 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + +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"), + ] + } + + assert [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]}, + stream_mode="messages", + ) + ] == [ + ( + _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", + } + ], + ), + { + "langgraph_step": 1, + "langgraph_node": "agent", + "langgraph_triggers": ["start:agent"], + "langgraph_path": ("__pregel_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": ["branch:agent:should_continue:tools"], + "langgraph_path": ("__pregel_pull", "tools"), + "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", + }, + ], + ), + { + "langgraph_step": 3, + "langgraph_node": "agent", + "langgraph_triggers": ["tools"], + "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ( + _AnyIdToolMessage( + content="result for another", + name="search_api", + tool_call_id="tool_call234", + ), + { + "langgraph_step": 4, + "langgraph_node": "tools", + "langgraph_triggers": ["branch:agent:should_continue:tools"], + "langgraph_path": ("__pregel_pull", "tools"), + "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": ["branch:agent:should_continue:tools"], + "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_checkpoint_ns": AnyStr("tools:"), + }, + ), + ( + _AnyIdAIMessageChunk( + content="answer", + ), + { + "langgraph_step": 5, + "langgraph_node": "agent", + "langgraph_triggers": ["tools"], + "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), + "ls_provider": "fakechatmodel", + "ls_model_type": "chat", + }, + ), + ] + + assert [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "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"}, + }, + ], + ) + ] + } + }, + { + "tools": { + "messages": [ + _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", + ), + ] + } + }, + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_state_graph_packets(checkpointer_name: str) -> 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] + session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] + + @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: + assert isinstance(data["session"], httpx.AsyncClient) + # 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")}}, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # interrupt after agent + + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + if not FF_SEND_V2: + return + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + ), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # 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)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + }, + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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": 3, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + } + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + ), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # 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)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + }, + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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": 3, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_message_graph(checkpointer_name: str) -> 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 FakeFuntionChatModel(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 = FakeFuntionChatModel( + 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 = MessageGraph() + + # 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")}, + ] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=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 + ) + ] == [ + { + "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, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # 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, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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, + "writes": { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + 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, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + +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": ["start: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": ["rewrite_query"], + }, + }, + ), + ( + "debug", + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "retriever_two", + "input": {"query": "query: what is weather in sf", "docs": []}, + "triggers": ["rewrite_query"], + }, + }, + ), + ( + "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": ["retriever_one", "retriever_two"], + }, + }, + ), + ("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"], + }, + ), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_start_branch_then(checkpointer_name: str) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] + other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] + + def assert_shared_value(data: State, config: RunnableConfig) -> State: + assert "shared" in data + if thread_id := config["configurable"].get("thread_id"): + if thread_id == "1": + # this is the first thread, so should not see a value + assert data["shared"] == {} + return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}} + elif thread_id == "2": + # this should get value saved by thread 1 + assert data["shared"] == {"1": {"hello": "world"}} + elif thread_id == "3": + # this is a different assistant, so should not see previous value + assert data["shared"] == {} + return {} + + def tool_two_slow(data: State, config: RunnableConfig) -> State: + return {"my_key": " slow", **assert_shared_value(data, config)} + + def tool_two_fast(data: State, config: RunnableConfig) -> State: + return {"my_key": " fast", **assert_shared_value(data, config)} + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two_slow", tool_two_slow) + tool_two_graph.add_node("tool_two_fast", tool_two_fast) + tool_two_graph.set_conditional_entry_point( + lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END + ) + tool_two = tool_two_graph.compile() + + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value fast", + "market": "US", + } + + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile( + store=InMemoryStore(), + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "assistant_id": "a", + "thread_id": "1", + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + # update state + await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_branch_then(checkpointer_name: str) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + tool_two_graph = StateGraph(State) + tool_two_graph.set_entry_point("prepare") + tool_two_graph.set_finish_point("finish") + tool_two_graph.add_conditional_edges( + source="prepare", + path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", + then="finish", + ) + tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) + tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) + tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) + tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) + tool_two = tool_two_graph.compile() + + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value prepared fast finished", + "market": "US", + } + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "10", + }, + "parent_config": None, + "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["prepare"], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ["branch:prepare:condition:tool_two_slow"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["finish"], + "tasks": [ + { + "id": AnyStr(), + "name": "finish", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": [], + "tasks": [], + }, + }, + ] + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "11"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "11", + }, + "parent_config": None, + "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "11", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["prepare"], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "11", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], + }, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "11", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "11", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "12"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "12", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "12", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "21", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "21", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "22", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "22", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread2, limit=2) + ][-1].config, + ) + + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = await tool_two.aupdate_state( + thread3, {"my_key": "key", "market": "DE"} + ) + # check current state + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + "thread_id": "23", + }, + parent_config=None, + ) + # run from this point + assert await tool_two.ainvoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "23", + }, + parent_config=uconfig, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread3)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "23", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread3, limit=2) + ][-1].config, + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_nested_graph_state(checkpointer_name: str) -> 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") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + assert await app.aget_state(config) == 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", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # 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", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + ), + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # get_state_history returns outer graph checkpoints + history = [c async for c in app.aget_state_history(config)] + assert history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # get_state_history for a subgraph returns its checkpoints + child_history = [ + c async for c in app.aget_state_history(history[0].tasks[0].state) + ] + assert 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", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + AnyStr(), + "inner_1", + (PULL, "inner_1"), + result={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + + # resume + await app.ainvoke(None, config, debug=True) + # 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", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # 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", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=( + PregelTask( + AnyStr(), + "outer_2", + (PULL, "outer_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + result={"my_key": "hi my value here and there"}, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + 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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_doubly_nested_graph_state(checkpointer_name: str) -> 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") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=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) + ] == [ + ((), {"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", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + child_state = await app.aget_state(outer_state.tasks[0].state) + assert ( + child_state.tasks[0] + == 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(), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ).tasks[0] + ) + 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", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [PULL, AnyStr("child_1")], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + ) + # 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", + "writes": { + "grandchild_1": { + "my_key": "hi my value here" + } + }, + "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + ), + ), + ), + 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", + "writes": None, + "step": 0, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # resume + assert [c async for c in app.astream(None, config, subgraphs=True)] == [ + ( + (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", + "writes": { + "parent_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + ) + # get outer graph history + outer_history = [c async for c in app.aget_state_history(config)] + assert ( + outer_history[0] + == [ + 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", + "writes": { + "parent_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("parent_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), name="parent_2", path=(PULL, "parent_2") + ), + ), + ), + 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", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("parent_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), name="parent_1", path=(PULL, "parent_1") + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), name="__start__", path=(PULL, "__start__") + ), + ), + ), + ][0] + ) + # get child graph history + child_history = [ + c async for c in app.aget_state_history(outer_history[2].tasks[0].state) + ] + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=(), + ), + 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", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="child_1", + path=(PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + } + }, + result={"my_key": "hi my value here and there"}, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + # get grandchild graph history + grandchild_history = [ + c async for c in app.aget_state_history(child_history[1].tasks[0].state) + ] + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + 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", + "writes": { + "grandchild_2": {"my_key": "hi my value here and there"} + }, + "step": 2, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=(), + ), + 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", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result={"my_key": "hi my value here and there"}, + ), + ), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("grandchild_1",), + 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", + "writes": None, + "step": 0, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_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(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_1", + path=(PULL, "grandchild_1"), + result={"my_key": "hi my value here"}, + ), + ), + ), + StateSnapshot( + values={}, + next=("__start__",), + 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": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + ), + ] + + # replay grandchild checkpoint + assert [ + c + async for c in app.astream( + None, grandchild_history[2].config, subgraphs=True + ) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_to_nested_graphs(checkpointer_name: str) -> 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=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) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=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) + + if not FF_SEND_V2: + # 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"], + } + return + + assert outer_state == StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + ), + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": { + "__start__": { + "subjects": [ + "cats", + "dogs", + ], + } + }, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ) + + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[2].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"], + } + + actual_snapshot = await graph.aget_state(config) + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] + }, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [c async for c in graph.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": [ + "Joke about cats - hohoho", + "Joke about turtles - hohoho", + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] + }, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"jokes": []}, + next=("__start__", "generate_joke", "generate_joke"), + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about cats - hohoho"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about turtles - hohoho"]}, + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_weather_subgraph( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> 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"] + + class Router(TypedDict): + 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) + 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] + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = graph.compile(checkpointer=checkpointer) + + 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 async for c in graph.astream(inputs, thread2, stream_mode="custom") + ] == [ + "I'm", + " very", + ] + assert [ + c async for c in graph.astream(None, thread2, stream_mode="custom") + ] == [ + " good", + ] + + # run until interrupt + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + ) + # 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 + ) + ] == [ + ((), {"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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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", + "writes": {"model_node": {"city": "San Francisco"}}, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + 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=( + PregelTask( + id=AnyStr(), + name="weather_node", + path=(PULL, "weather_node"), + ), + ), + ), + ), + ), + ) + 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", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + 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", + "writes": { + "weather_node": { + "messages": [ + {"role": "assistant", "content": "rainy"} + ] + } + }, + "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + 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=(), + ), + ), + ), + ) + 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"), + ] + } + }, + ), + ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ad48108a9..b9b858c77 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2,14 +2,13 @@ import enum import json import logging import operator -import re +import threading import time import uuid import warnings -from collections import Counter +from collections import Counter, deque from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from dataclasses import replace from random import randrange from typing import ( Annotated, @@ -24,7 +23,6 @@ from typing import ( Tuple, TypedDict, Union, - cast, get_type_hints, ) @@ -33,12 +31,9 @@ import pytest from langchain_core.runnables import ( RunnableConfig, RunnableLambda, - RunnableMap, RunnablePassthrough, - RunnablePick, ) from langsmith import traceable -from pydantic import BaseModel from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -48,7 +43,6 @@ from langgraph.channels.context import Context 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.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -61,20 +55,16 @@ from langgraph.constants import ( ERROR, FF_SEND_V2, PULL, - PUSH, START, ) -from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt +from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore -from langgraph.store.memory import InMemoryStore from langgraph.types import ( Command, Interrupt, @@ -83,18 +73,16 @@ from langgraph.types import ( StreamWriter, interrupt, ) -from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.agents import AgentAction, AgentFinish +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, ALL_STORES_SYNC, SHOULD_CHECK_SNAPSHOTS, ) -from tests.fake_chat import FakeChatModel -from tests.fake_tracer import FakeTracer from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import ( _AnyIdAIMessage, - _AnyIdAIMessageChunk, _AnyIdHumanMessage, _AnyIdToolMessage, ) @@ -102,46 +90,6 @@ from tests.messages import ( logger = logging.getLogger(__name__) -# define these objects to avoid importing langchain_core.agents -# and therefore avoid relying on core Pydantic version -class AgentAction(BaseModel): - tool: str - tool_input: Union[str, dict] - log: str - type: Literal["AgentAction"] = "AgentAction" - - model_config = { - "json_schema_extra": { - "description": ( - """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.""" - ) - } - } - - -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" - model_config = { - "json_schema_extra": { - "description": ( - """Final return value of an ActionAgent. - -Agents return an AgentFinish when they have reached a stopping condition.""" - ) - } - } - - def test_graph_validation() -> None: def logic(inp: str) -> str: return "" @@ -676,452 +624,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_invoke_two_processes_in_out_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.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=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) is None - - # inbox == 3 - checkpoint = checkpointer.get(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert app.invoke(None, thread1) == 4 - - # start execution again, stop at inbox - assert app.invoke(20, thread1) is None - - # inbox == 21 - checkpoint = 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) is None - assert app.invoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert app.invoke(20, thread2) 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 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, - "writes": {"two": 5}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - 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, - "writes": {"input": 3}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - 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, - "writes": {"input": 20}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - 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, - "writes": {"two": 4}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[7].config, - ), - 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, - "writes": {"input": 2}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # 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__": ()}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_fork_always_re_runs_nodes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - 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=checkpointer) - - thread1 = {"configurable": {"thread_id": "1"}} - - # start execution, stop at inbox - assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ - ("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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - 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, - "writes": None, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - 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, - "writes": {"__start__": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # 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}, - ] - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_run_from_checkpoint_id_retains_previous_writes( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture @@ -2037,1074 +1539,6 @@ def test_imp_stream_order( assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_send_dedupe_on_resume( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - if not FF_SEND_V2: - pytest.skip("Send deduplication is only available in Send V2") - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class InterruptOnce: - ticks: int = 0 - - def __call__(self, state): - self.ticks += 1 - if self.ticks == 1: - raise NodeInterrupt("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=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - assert graph.invoke(["0"], thread1, debug=1) == [ - "0", - "1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - ] - 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) == 2 - # resume execution - assert graph.invoke(None, thread1, debug=1) == [ - "0", - "1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - "flaky|4", - "3", - "3.1", - ] - # 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 ( - history[1] - == [ - StateSnapshot( - values=[ - "0", - "1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - "flaky|4", - "3", - "3.1", - ], - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"3": ["3"], "3.1": ["3.1"]}, - "thread_id": "1", - "step": 2, - "parents": {}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values=[ - "0", - "1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - "flaky|4", - ], - next=("3", "3.1"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": { - "1": ["1"], - "2": [ - ["2|Command(goto=Send(node='2', arg=3))"], - ["2|Command(goto=Send(node='flaky', arg=4))"], - ["2|3"], - ], - "flaky": ["flaky|4"], - }, - "thread_id": "1", - "step": 1, - "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"], - ), - PregelTask( - id=AnyStr(), - name="3.1", - path=("__pregel_pull", "3.1"), - error=None, - interrupts=(), - state=None, - result=["3.1"], - ), - ), - ), - StateSnapshot( - values=["0"], - next=("1", "2", "2", "2", "flaky"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": None, - "thread_id": "1", - "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"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='2', arg=3))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='flaky', arg=4))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), - 2, - AnyStr(), - ), - error=None, - interrupts=(), - state=None, - result=["2|3"], - ), - PregelTask( - id=AnyStr(), - name="flaky", - path=( - "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), - 2, - AnyStr(), - ), - error=None, - interrupts=(Interrupt(value="Bahh", when="during"),), - state=None, - result=["flaky|4"], - ), - ), - ), - StateSnapshot( - values=[], - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": ["0"]}, - "thread_id": "1", - "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"], - ), - ), - ), - ][1] - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_send_react_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - 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=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=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - 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 - - if not FF_SEND_V2: - return - - # 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": 0, - "source": "loop", - "writes": None, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), - 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": 1, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "3"}} - 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 - - # 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": 0, - "source": "loop", - "writes": None, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), - 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": 1, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", (), 0), - 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 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_send_react_interrupt_control( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> None: - from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - 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() - 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=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 - - if not FF_SEND_V2: - return - - # interrupt-update-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - 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 - - # 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": 0, - "source": "loop", - "writes": None, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), - 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": 1, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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()) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str @@ -3449,847 +1883,6 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_conditional_graph( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - from langchain_core.language_models.fake import FakeStreamingListLLM - from langchain_core.prompts import PromptTemplate - from langchain_core.runnables import RunnablePassthrough - from langchain_core.tools import tool - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - # 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) -> Union[AgentAction, AgentFinish]: - if input.startswith("finish"): - _, answer = input.split(":") - return AgentFinish(return_values={"answer": answer}, log=input) - else: - _, tool_name, tool_input = input.split(":") - return AgentAction(tool=tool_name, tool_input=tool_input, log=input) - - agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) - - # Define tool execution logic - def execute_tools(data: dict) -> dict: - data = data.copy() - agent_action: AgentAction = data.pop("agent_outcome") - observation = {t.name: t for t in tools}[agent_action.tool].invoke( - agent_action.tool_input - ) - if data.get("intermediate_steps") is None: - data["intermediate_steps"] = [] - else: - data["intermediate_steps"] = data["intermediate_steps"].copy() - data["intermediate_steps"].append([agent_action, observation]) - return data - - # Define decision-making logic - def should_continue(data: dict) -> 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 = Graph() - - workflow.add_node("agent", agent) - workflow.add_node( - "tools", - execute_tools, - metadata={"parents": {}, "version": 2, "variant": "b"}, - ) - - 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 SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_graph().draw_mermaid() == snapshot - assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot - assert app.get_graph(xray=True).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 [c for c in app.stream({"input": "what is weather in sf"})] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "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": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] - - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - if SHOULD_CHECK_SNAPSHOTS: - assert app_w_interrupt.get_graph().to_json() == snapshot - assert app_w_interrupt.get_graph().draw_mermaid() == snapshot - - assert [ - c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - config=app_w_interrupt.checkpointer.get_tuple(config).config, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - assert ( - app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][ - "checkpoint_id" - ] - is not None - ) - - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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:a different query", - ), - "input": "what is weather in sf", - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - app_w_interrupt.update_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "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": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - 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) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - app_w_interrupt.update_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "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": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - 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) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "3", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "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": { - "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" - ), - } - }, - ] - - def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: def left(data: str) -> str: return data + "->left" @@ -4378,834 +1971,6 @@ def test_conditional_entrypoint_to_multiple_state_graph( } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_conditional_state_graph( - snapshot: SnapshotAssertion, - mocker: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - from langchain_core.language_models.fake import FakeStreamingListLLM - from langchain_core.prompts import PromptTemplate - from langchain_core.tools import tool - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() - - class AgentState(TypedDict, total=False): - input: Annotated[str, UntrackedValue] - agent_outcome: Optional[Union[AgentAction, AgentFinish]] - intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - session: Annotated[httpx.Client, Context(make_httpx_client)] - - class ToolState(TypedDict, total=False): - agent_outcome: Union[AgentAction, AgentFinish] - session: Annotated[httpx.Client, Context(make_httpx_client)] - - # 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, Union[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 isinstance(data["session"], httpx.Client) - 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 - assert isinstance(data["session"], httpx.Client) - # 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=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 SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - - with assert_ctx_once(): - 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" - ), - } - - with assert_ctx_once(): - 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=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream({"input": "what is weather in sf"}, 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - with assert_ctx_once(): - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - with assert_ctx_once(): - 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__": ()}, - ] - - with assert_ctx_once(): - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=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) - ] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # test w interrupt before all - app_w_interrupt = workflow.compile( - checkpointer=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) - ] == [ - {"__interrupt__": ()}, - ] - - assert app_w_interrupt.get_state(config) == StateSnapshot( - values={ - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "agent", (PULL, "agent")),), - next=("agent",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "3", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "3", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - "thread_id": "3", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=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) - ] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "4", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - "thread_id": "4", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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_conditional_state_graph_with_list_edge_inputs(snapshot: SnapshotAssertion): class State(TypedDict): foo: Annotated[list[str], operator.add] @@ -5397,4428 +2162,6 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None ] -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_tool_calling_executor(model, tools) - - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == 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"), - ] - } - - assert [ - c - for c in app.stream( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="messages", - ) - ] == [ - ( - _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", - } - ], - ), - { - "langgraph_step": 1, - "langgraph_node": "agent", - "langgraph_triggers": ["start: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": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "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", - }, - ], - ), - { - "langgraph_step": 3, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": (PULL, "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "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": ["branch:agent:should_continue:tools"], - "langgraph_path": (PULL, "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="answer", - ), - { - "langgraph_step": 5, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "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 - - assert ( - app.invoke( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="updates", - )[0]["agent"]["messages"] - == [ - { - "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"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _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", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ][0]["agent"]["messages"] - ) - - assert [ - *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) - ] == [ - { - "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"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _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", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_state_graph_packets( - request: pytest.FixtureRequest, checkpointer_name: str, 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 - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - class AgentState(TypedDict): - messages: Annotated[list[BaseMessage], add_messages] - session: Annotated[httpx.Client, Context(httpx.Client)] - - @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: - assert isinstance(data["session"], httpx.Client) - return { - "messages": model.invoke(data["messages"]), - "something_extra": "hi there", - } - - # Define decision-making logic - def should_continue(data: AgentState) -> str: - assert isinstance(data["session"], httpx.Client) - 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=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 - ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - {"__interrupt__": ()}, - ] - - if not FF_SEND_V2: - return - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), - next=("tools",), - config=(app_w_interrupt.checkpointer.get_tuple(config)).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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)),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - "something_extra": "hi there", - } - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), - ), - next=("tools", "tools"), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - }, - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 3, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - "something_extra": "hi there", - } - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # interrupt before tools - - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), - next=("tools",), - config=(app_w_interrupt.checkpointer.get_tuple(config)).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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)),), - next=("tools",), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - "something_extra": "hi there", - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), - ), - next=("tools", "tools"), - config=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - }, - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 3, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - "something_extra": "hi there", - } - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_message_graph( - snapshot: SnapshotAssertion, - deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> 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 - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - def _generate( - self, - messages: list[BaseMessage], - stop: Optional[list[str]] = None, - run_manager: Optional[CallbackManagerForLLMRun] = 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 = FakeFuntionChatModel( - 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 = MessageGraph() - - # 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 SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == 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=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) - ] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - app_w_interrupt = workflow.compile( - checkpointer=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)] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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 BaseMessge 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_root_graph( - deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> 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 - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - - class FakeFuntionChatModel(FakeMessagesListChatModel): - def bind_functions(self, functions: list): - return self - - def _generate( - self, - messages: list[BaseMessage], - stop: Optional[list[str]] = None, - run_manager: Optional[CallbackManagerForLLMRun] = 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 = FakeFuntionChatModel( - 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-000000000037", - ) - ] - }, - { - "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-000000000045", - ) - ] - }, - {"agent": AIMessage(content="answer", id="ai3")}, - ] - - app_w_interrupt = workflow.compile( - checkpointer=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) - ] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - app_w_interrupt = workflow.compile( - checkpointer=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)] == [ - { - "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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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 BaseMessge 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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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=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=app_w_interrupt.checkpointer.get_tuple(config).config, - created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 6, - "writes": {"tools": UnsortedSequence("ai", "an extra message")}, - "thread_id": "2", - }, - parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, - ) - - # 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-000000000078", - ), - 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-000000000100" - ), - 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": ["start: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": ["rewrite_query"], - }, - }, - ), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_two", - "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], - }, - }, - ), - ( - "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": ["retriever_one", "retriever_two"], - }, - }, - ), - ("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"], - }, - ), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_dynamic_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - 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=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", - } - 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=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...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - # 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) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - 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...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), - ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # 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=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - -@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_copy_checkpoint( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - 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[Union[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=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", - } - 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 one all good", - "market": "US", - } - - tool_two = tool_two_graph.compile(checkpointer=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...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - # 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) == { - "my_key": "value ⛰️ one", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), - ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # clear the interrupt and next tasks - tool_two.update_state(thread1, None) - # interrupt is cleared, next task is kept - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), - ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_dynamic_interrupt_subgraph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - 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=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", - } - 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=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...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ) - }, - ] - # 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) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [ - c.metadata - for c in tool_two.checkpointer.list( - {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - ) - ] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - 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...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("tool_two:"), - } - }, - ), - ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=[ - *tool_two.checkpointer.list( - {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 - ) - ][-1].config, - ) - # 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=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=[ - *tool_two.checkpointer.list( - {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 - ) - ][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_start_branch_then( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} - - def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} - - def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two_slow", tool_two_slow) - tool_two_graph.add_node("tool_two_fast", tool_two_fast) - tool_two_graph.set_conditional_entry_point( - lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END - ) - tool_two = tool_two_graph.compile() - assert tool_two.get_graph().draw_mermaid() == snapshot - - assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { - "my_key": "value slow", - "market": "DE", - } - assert tool_two.invoke({"my_key": "value", "market": "US"}) == { - "my_key": "value fast", - "market": "US", - } - - tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { - "my_key": "value ⛰️", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value ⛰️ slow", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ slow", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - # update state - tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_branch_then( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - tool_two_graph = StateGraph(State) - tool_two_graph.set_entry_point("prepare") - tool_two_graph.set_finish_point("finish") - tool_two_graph.add_conditional_edges( - source="prepare", - path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", - then="finish", - ) - tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) - tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) - tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) - tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) - tool_two = tool_two_graph.compile() - assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot - assert tool_two.get_graph().draw_mermaid() == snapshot - - assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert tool_two.invoke({"my_key": "value", "market": "US"}) == { - "my_key": "value prepared fast finished", - "market": "US", - } - - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - thread10 = {"configurable": {"thread_id": "10"}} - - res = [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] - - assert res == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "10", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - {"id": AnyStr(), "name": "prepare", "interrupts": (), "state": None} - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["finish"], - "tasks": [ - {"id": AnyStr(), "name": "finish", "interrupts": (), "state": None} - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "1", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "2", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "2", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["finish"] - ) - - thread1 = {"configurable": {"thread_id": "11"}} - - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared slow", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={ - "my_key": "value prepared slow", - "market": "DE", - }, - tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "thread_id": "11", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - # update state - tool_two.update_state(thread1, {"my_key": "er"}) - assert tool_two.get_state(thread1) == StateSnapshot( - values={ - "my_key": "value prepared slower", - "market": "DE", - }, - tasks=(PregelTask(AnyStr(), "finish", (PULL, "finish")),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 3, - "writes": {"tool_two_slow": {"my_key": "er"}}, - "thread_id": "11", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - tool_two.invoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "21", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "21", - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "22", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert tool_two.get_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "22", - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) - - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = tool_two.update_state(thread3, {"my_key": "key", "market": "DE"}) - # check current state - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - "thread_id": "23", - }, - parent_config=None, - ) - # run from this point - assert tool_two.invoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "23", - }, - parent_config=uconfig, - ) - # resume, for same result as above - assert tool_two.invoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert tool_two.get_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "23", - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge( snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str @@ -11191,2374 +3534,6 @@ def test_doubly_nested_graph_interrupts( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - 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=checkpointer) - - config = {"configurable": {"thread_id": "1"}} - app.invoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert app.get_state(config) == 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", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # 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", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # get_state_history returns outer graph checkpoints - history = list(app.get_state_history(config)) - assert history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # get_state_history for a subgraph returns its checkpoints - child_history = [*app.get_state_history(history[0].tasks[0].state)] - assert 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", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # resume - app.invoke(None, config, debug=True) - # 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", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # 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", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} - }, - result={"my_key": "hi my value here and there"}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - 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 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_doubly_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - 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=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)] == [ - ((), {"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", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - child_state = app.get_state(outer_state.tasks[0].state) - assert ( - child_state.tasks[0] - == 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(), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ).tasks[0] - ) - 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", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [PULL, AnyStr("child_1")], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - ) - # 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", - "writes": { - "grandchild_1": {"my_key": "hi my value here"} - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - ), - ), - ), - 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", - "writes": None, - "step": 0, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - ), - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # # resume - assert [c for c in app.stream(None, config, subgraphs=True)] == [ - ( - (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", - "writes": { - "parent_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - ) - # 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", - "writes": { - "parent_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="parent_2", - path=(PULL, "parent_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="parent_1", - path=(PULL, "parent_1"), - result={"my_key": "hi my value"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - ), - ] - # get child graph history - child_history = list(app.get_state_history(outer_history[2].tasks[0].state)) - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(), - ), - 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", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="child_1", - path=(PULL, "child_1"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - # get grandchild graph history - grandchild_history = list(app.get_state_history(child_history[1].tasks[0].state)) - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - 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", - "writes": {"grandchild_2": {"my_key": "hi my value here and there"}}, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=(), - ), - 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", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_2", - path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - 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", - "writes": None, - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - 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": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # replay grandchild checkpoint - assert [ - c for c in app.stream(None, grandchild_history[2].config, subgraphs=True) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_send_to_nested_graphs( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - 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=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=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) - - if not FF_SEND_V2: - # 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"]}}, - ] - return - - assert outer_state == StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ) - # check state of each of the inner tasks - assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( - values={"subject": "cats - hohoho", "jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": None, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1], - "langgraph_step": 0, - "langgraph_triggers": [PUSH], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), - ) - assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot( - values={"subject": "dogs - hohoho", "jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": None, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2], - "langgraph_step": 0, - "langgraph_triggers": [PUSH], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), - ) - # update state of dogs joke graph - graph.update_state( - outer_state.tasks[2 if FF_SEND_V2 else 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"]}}, - ] - - actual_snapshot = graph.get_state(config) - expected_snapshot = StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert actual_snapshot == expected_snapshot - - # test full history - actual_history = list(graph.get_state_history(config)) - - # get subgraph node state for expected history - expected_history = [ - StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - next=("__start__", "generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_weather_subgraph( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion -) -> 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 - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - # 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) - graph.add_edge("normal_llm_node", END) - graph.add_edge("weather_graph", END) - graph = graph.compile(checkpointer=checkpointer) - - 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")] == [ - "I'm", - " very", - ] - assert [c for c in graph.stream(None, thread2, stream_mode="custom")] == [ - " good", - ] - - # run until interrupt - assert [ - c - for c in graph.stream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("weather_graph:"), - } - }, - ), - ), - ) - - # 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 - ) - ] == [ - ((), {"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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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", - "writes": {"model_node": {"city": "San Francisco"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - 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=( - PregelTask( - id=AnyStr(), - name="weather_node", - path=(PULL, "weather_node"), - ), - ), - ), - ), - ), - ) - 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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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", - "writes": { - "weather_node": { - "messages": [{"role": "assistant", "content": "rainy"}] - } - }, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_id": AnyStr(), - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - 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=(), - ), - ), - ), - ) - 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"), - ] - } - }, - ), - ] - - def test_repeat_condition(snapshot: SnapshotAssertion) -> None: class AgentState(TypedDict): hello: str @@ -15042,3 +5017,201 @@ def test_command_goto_with_static_breakpoints( graph.invoke({"foo": "abc"}, config) result = graph.invoke(Command(goto=["node2"]), config) assert result == {"foo": "abc|node-1|node-2|node-2"} + + +def test_nested_graph_state_error_handling(): + """Test error handling when updating state in nested graphs.""" + + class State(TypedDict): + count: int + + def child_node(state: State): + return {"count": state["count"] + 1} + + child = StateGraph(State) + child.add_node("child", child_node) + child.add_edge(START, "child") + + parent = StateGraph(State) + parent.add_node("child_graph", child.compile()) + parent.add_edge(START, "child_graph") + + app = parent.compile(checkpointer=MemorySaver()) + + # Test invalid state update on parent + with pytest.raises(InvalidUpdateError): + app.update_state({"configurable": {"thread_id": "1"}}, {"invalid_key": "value"}) + + # Test invalid state update on child + with pytest.raises(InvalidUpdateError): + app.update_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, + {"invalid_key": "value"}, + ) + + +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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_interrupt_state_persistence( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + """Test that state is preserved correctly across multiple interrupts.""" + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_checkpoint_recovery(request: pytest.FixtureRequest, checkpointer_name: str): + """Test recovery from checkpoints after failures.""" + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + 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=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First attempt should fail + with pytest.raises(RuntimeError): + graph.invoke({"steps": ["start"], "attempt": 1}, config) + + # 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 + + # Retry with updated attempt count + result = graph.invoke({"steps": [], "attempt": 2}, config) + assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} + + # Verify checkpoint history shows both attempts + history = list(graph.get_state_history(config)) + assert len(history) == 6 # Initial + failed attempt + successful attempt + + # 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 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index cde54ae7c..6ddd22f1d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2,10 +2,9 @@ import asyncio import logging import operator import random -import re import sys import uuid -from collections import Counter +from collections import Counter, deque from contextlib import asynccontextmanager, contextmanager from dataclasses import replace from time import perf_counter @@ -22,21 +21,17 @@ from typing import ( Tuple, TypedDict, Union, - cast, ) from uuid import UUID import httpx import pytest -from langchain_core.messages import ToolCall from langchain_core.runnables import ( RunnableConfig, RunnableLambda, RunnablePassthrough, - RunnablePick, ) from langchain_core.utils.aiter import aclosing -from pydantic import BaseModel from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -45,7 +40,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic -from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( ChannelVersions, Checkpoint, @@ -58,20 +52,16 @@ from langgraph.constants import ( ERROR, FF_SEND_V2, PULL, - PUSH, START, ) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph -from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor +from langgraph.graph.message import MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore -from langgraph.store.memory import InMemoryStore from langgraph.types import ( Command, Interrupt, @@ -80,7 +70,7 @@ from langgraph.types import ( StreamWriter, interrupt, ) -from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence +from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, @@ -89,7 +79,6 @@ from tests.conftest import ( awith_checkpointer, awith_store, ) -from tests.fake_chat import FakeChatModel from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -97,7 +86,6 @@ from tests.memory_assert import ( ) from tests.messages import ( _AnyIdAIMessage, - _AnyIdAIMessageChunk, _AnyIdHumanMessage, _AnyIdToolMessage, ) @@ -1326,475 +1314,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_two_processes_in_out_interrupt( - checkpointer_name: str, mocker: MockerFixture -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=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) is None - - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 - - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None - - # inbox == 21 - checkpoint = await 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) is None - assert await app.ainvoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) 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 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, - "writes": {"two": 5}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - 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, - "writes": {"input": 3}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - 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, - "writes": {"input": 20}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - 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, - "writes": {"two": 4}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - 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, - "writes": {"one": None}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[7].config, - ), - 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, - "writes": {"input": 2}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # 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__": ()}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_fork_always_re_runs_nodes( - checkpointer_name: str, 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) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=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"]) - ] == [ - ("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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - 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, - "writes": {"add_one": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - 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, - "writes": None, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - ), - 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, - "writes": {"__start__": 1}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # 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_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") @@ -4331,1556 +3850,6 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph(checkpointer_name: str) -> 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.runnables import RunnablePassthrough - from langchain_core.tools import tool - - # 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", - ] - ) - - async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]: - if input.startswith("finish"): - _, answer = input.split(":") - return AgentFinish(return_values={"answer": answer}, log=input) - else: - _, tool_name, tool_input = input.split(":") - return AgentAction(tool=tool_name, tool_input=tool_input, log=input) - - agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser) - - # Define tool execution logic - async def execute_tools(data: dict) -> dict: - data = data.copy() - agent_action: AgentAction = data.pop("agent_outcome") - observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke( - agent_action.tool_input - ) - if data.get("intermediate_steps") is None: - data["intermediate_steps"] = [] - else: - data["intermediate_steps"] = data["intermediate_steps"].copy() - data["intermediate_steps"].append([agent_action, observation]) - return data - - # Define decision-making logic - async def should_continue(data: dict, config: RunnableConfig) -> 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 = Graph() - - 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": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "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": { - "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" - ), - } - }, - ] - - 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 - assert "/logs/agent:2/streamed_output/-" in patch_paths - assert "/logs/agent:3/streamed_output/-" in patch_paths - # Check that agent (one of the nodes) has its final output set in the logs - assert "/logs/agent/final_output" in patch_paths - assert "/logs/agent:2/final_output" in patch_paths - assert "/logs/agent:3/final_output" in patch_paths - assert [ - p["value"] - for log in patches - for p in log.ops - if p["path"] == "/logs/agent/final_output" - or p["path"] == "/logs/agent:2/final_output" - or p["path"] == "/logs/agent:3/final_output" - ] == [ - { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - }, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - }, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "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": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "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": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "3", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "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": { - "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" - ), - } - }, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph_state( - mocker: MockerFixture, checkpointer_name: str -) -> 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 - - setup = mocker.Mock() - teardown = mocker.Mock() - - @asynccontextmanager - async def assert_ctx_once() -> AsyncIterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - class MyPydanticContextModel(BaseModel, arbitrary_types_allowed=True): - session: httpx.AsyncClient - something_else: str - - @asynccontextmanager - async def make_context( - config: RunnableConfig, - ) -> AsyncIterator[MyPydanticContextModel]: - assert isinstance(config, dict) - setup() - session = httpx.AsyncClient() - try: - yield MyPydanticContextModel(session=session, something_else="hello") - finally: - await session.aclose() - teardown() - - class AgentState(TypedDict): - input: Annotated[str, UntrackedValue] - agent_outcome: Optional[Union[AgentAction, AgentFinish]] - intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - context: Annotated[MyPydanticContextModel, Context(make_context)] - - # 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, Union[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: - # check we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) - # 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 we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) - # 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() - - async with assert_ctx_once(): - 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" - ), - } - - async with assert_ctx_once(): - 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" - ), - } - }, - ] - - async with assert_ctx_once(): - 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" - ), - }, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - async with assert_ctx_once(): - 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - async with assert_ctx_once(): - 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__": ()}, - ] - - async with assert_ctx_once(): - 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - async def test_conditional_entrypoint_graph() -> None: async def left(data: str) -> str: return data + "->left" @@ -5963,2895 +3932,6 @@ async def test_conditional_entrypoint_graph_state() -> None: ] -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_tool_calling_executor(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"), - ] - } - - assert [ - c - async for c in app.astream( - {"messages": [HumanMessage(content="what is weather in sf")]}, - stream_mode="messages", - ) - ] == [ - ( - _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", - } - ], - ), - { - "langgraph_step": 1, - "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], - "langgraph_path": ("__pregel_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": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "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", - }, - ], - ), - { - "langgraph_step": 3, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ( - _AnyIdToolMessage( - content="result for another", - name="search_api", - tool_call_id="tool_call234", - ), - { - "langgraph_step": 4, - "langgraph_node": "tools", - "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "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": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), - "langgraph_checkpoint_ns": AnyStr("tools:"), - }, - ), - ( - _AnyIdAIMessageChunk( - content="answer", - ), - { - "langgraph_step": 5, - "langgraph_node": "agent", - "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), - "langgraph_checkpoint_ns": AnyStr("agent:"), - "checkpoint_ns": AnyStr("agent:"), - "ls_provider": "fakechatmodel", - "ls_model_type": "chat", - }, - ), - ] - - assert [ - c - async for c in app.astream( - {"messages": [HumanMessage(content="what is weather in sf")]} - ) - ] == [ - { - "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"}, - }, - ], - ) - ] - } - }, - { - "tools": { - "messages": [ - _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", - ), - ] - } - }, - {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}}, - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_state_graph_packets(checkpointer_name: str) -> 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] - session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] - - @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: - assert isinstance(data["session"], httpx.AsyncClient) - # 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")}}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # interrupt after agent - - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - {"__interrupt__": ()}, - ] - - if not FF_SEND_V2: - return - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # 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)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - }, - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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": 3, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # interrupt before tools - - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # 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)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - }, - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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": 3, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_message_graph(checkpointer_name: str) -> 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 FakeFuntionChatModel(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 = FakeFuntionChatModel( - 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 = MessageGraph() - - # 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")}, - ] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=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 - ) - ] == [ - { - "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, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # 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, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - 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, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - -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": ["start: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": ["rewrite_query"], - }, - }, - ), - ( - "debug", - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "retriever_two", - "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], - }, - }, - ), - ( - "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": ["retriever_one", "retriever_two"], - }, - }, - ), - ("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"], - }, - ), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_start_branch_then(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} - - def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} - - def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} - - tool_two_graph = StateGraph(State) - tool_two_graph.add_node("tool_two_slow", tool_two_slow) - tool_two_graph.add_node("tool_two_fast", tool_two_fast) - tool_two_graph.set_conditional_entry_point( - lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END - ) - tool_two = tool_two_graph.compile() - - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value fast", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value slow", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "assistant_id": "a", - "thread_id": "1", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "a", - "thread_id": "2", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - # update state - await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "b", - "thread_id": "3", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_branch_then(checkpointer_name: str) -> None: - class State(TypedDict): - my_key: Annotated[str, operator.add] - market: str - - tool_two_graph = StateGraph(State) - tool_two_graph.set_entry_point("prepare") - tool_two_graph.set_finish_point("finish") - tool_two_graph.add_conditional_edges( - source="prepare", - path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", - then="finish", - ) - tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"}) - tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"}) - tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"}) - tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) - tool_two = tool_two_graph.compile() - - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value prepared fast finished", - "market": "US", - } - - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "10", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["finish"], - "tasks": [ - { - "id": AnyStr(), - "name": "finish", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "11"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "11", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "11", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "12"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "12", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "12", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "21", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "21", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) - - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "22", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "22", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, limit=2) - ][-1].config, - ) - - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state( - thread3, {"my_key": "key", "market": "DE"} - ) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - "thread_id": "23", - }, - parent_config=None, - ) - # run from this point - assert await tool_two.ainvoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "23", - }, - parent_config=uconfig, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread3)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "23", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, limit=2) - ][-1].config, - ) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None: def sorted_add( @@ -10086,2325 +5166,6 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_state(checkpointer_name: str) -> 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") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) - - config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert await app.aget_state(config) == 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", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # 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", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - assert history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # get_state_history for a subgraph returns its checkpoints - child_history = [ - c async for c in app.aget_state_history(history[0].tasks[0].state) - ] - assert 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", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["outer_1"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # resume - await app.ainvoke(None, config, debug=True) - # 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", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # 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", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - 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 - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_state(checkpointer_name: str) -> 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") - - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=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) - ] == [ - ((), {"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", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - child_state = await app.aget_state(outer_state.tasks[0].state) - assert ( - child_state.tasks[0] - == 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(), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ).tasks[0] - ) - 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", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [PULL, AnyStr("child_1")], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - ) - # 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", - "writes": { - "grandchild_1": { - "my_key": "hi my value here" - } - }, - "step": 1, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - ), - ), - ), - 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", - "writes": None, - "step": 0, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - ), - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # resume - assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ( - (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", - "writes": { - "parent_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - ) - # get outer graph history - outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - 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", - "writes": { - "parent_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_2", path=(PULL, "parent_2") - ), - ), - ), - 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", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_1", path=(PULL, "parent_1") - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), name="__start__", path=(PULL, "__start__") - ), - ), - ), - ][0] - ) - # get child graph history - child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) - ] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=(), - ), - 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", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="child_1", - path=(PULL, "child_1"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [AnyStr("parent_1")], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - # get grandchild graph history - grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) - ] - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - 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", - "writes": { - "grandchild_2": {"my_key": "hi my value here and there"} - }, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=(), - ), - 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", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_2", - path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, - ), - ), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - 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", - "writes": None, - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_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(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), - ), - ), - StateSnapshot( - values={}, - next=("__start__",), - 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": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [AnyStr("start:child_1")], - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # replay grandchild checkpoint - assert [ - c - async for c in app.astream( - None, grandchild_history[2].config, subgraphs=True - ) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), - ] - - -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_to_nested_graphs(checkpointer_name: str) -> 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=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) - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=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) - - if not FF_SEND_V2: - # 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"], - } - return - - assert outer_state == StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": { - "__start__": { - "subjects": [ - "cats", - "dogs", - ], - } - }, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ) - - # update state of dogs joke graph - await graph.aupdate_state( - outer_state.tasks[2].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"], - } - - actual_snapshot = await graph.aget_state(config) - expected_snapshot = StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert actual_snapshot == expected_snapshot - - # test full history - actual_history = [c async for c in graph.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": [ - "Joke about cats - hohoho", - "Joke about turtles - hohoho", - ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - next=("__start__", "generate_joke", "generate_joke"), - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - - -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_weather_subgraph( - checkpointer_name: str, snapshot: SnapshotAssertion -) -> 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"] - - class Router(TypedDict): - 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) - 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] - - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = graph.compile(checkpointer=checkpointer) - - 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 async for c in graph.astream(inputs, thread2, stream_mode="custom") - ] == [ - "I'm", - " very", - ] - assert [ - c async for c in graph.astream(None, thread2, stream_mode="custom") - ] == [ - " good", - ] - - # run until interrupt - assert [ - c - async for c in graph.astream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("weather_graph:"), - } - }, - ), - ), - ) - # 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 - ) - ] == [ - ((), {"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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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", - "writes": {"model_node": {"city": "San Francisco"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - 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=( - PregelTask( - id=AnyStr(), - name="weather_node", - path=(PULL, "weather_node"), - ), - ), - ), - ), - ), - ) - 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", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - 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", - "writes": { - "weather_node": { - "messages": [ - {"role": "assistant", "content": "rainy"} - ] - } - }, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_id": AnyStr(), - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:router_node:route_after_prediction:weather_graph" - ], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - 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=(), - ), - ), - ), - ) - 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"), - ] - } - }, - ), - ] - - async def test_checkpoint_metadata() -> None: """This test verifies that a run's configurable fields are merged with the previous checkpoint config for each step in the run. @@ -13331,3 +6092,195 @@ async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> N 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_nested_graph_state_error_handling(): + """Test error handling when updating state in nested graphs.""" + + class State(TypedDict): + count: int + + def child_node(state: State): + return {"count": state["count"] + 1} + + child = StateGraph(State) + child.add_node("child", child_node) + child.add_edge(START, "child") + + parent = StateGraph(State) + parent.add_node("child_graph", child.compile()) + parent.add_edge(START, "child_graph") + + app = parent.compile(checkpointer=MemorySaver()) + + # Test invalid state update on parent + with pytest.raises(InvalidUpdateError): + await app.aupdate_state( + {"configurable": {"thread_id": "1"}}, {"invalid_key": "value"} + ) + + # Test invalid state update on child + with pytest.raises(InvalidUpdateError): + await app.aupdate_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}}, + {"invalid_key": "value"}, + ) + + +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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> 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") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = builder.compile(checkpointer=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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_checkpoint_recovery_async(checkpointer_name: str): + """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") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # First attempt should fail + with pytest.raises(RuntimeError): + await graph.ainvoke({"steps": ["start"], "attempt": 1}, config) + + # 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) + 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)] + assert len(history) == 6 # Initial + failed attempt + successful attempt + + # 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