From 45e3d1a3f168e997baaf820af545c1f9ed172433 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 13:51:38 -0700 Subject: [PATCH] Test all checkpointers everywhere we test 1 of them --- .../tests/__snapshots__/test_pregel.ambr | 2900 +++++++++++++++++ libs/langgraph/tests/memory_assert.py | 9 - libs/langgraph/tests/test_pregel.py | 1904 ++++++----- libs/langgraph/tests/test_pregel_async.py | 1790 +++++----- 4 files changed, 4797 insertions(+), 1806 deletions(-) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index c92bc4688..4835ec611 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -33,6 +33,142 @@ ''' # --- +# 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[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"}' # --- @@ -601,6 +737,1386 @@ ''' # --- +# 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", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "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
version = 2 + variant = b) + __end__([__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": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[memory].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  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', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': 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
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__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", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "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
version = 2 + variant = b) + __end__([__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": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  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', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': 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
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__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", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "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
version = 2 + variant = b) + __end__([__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": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[postgres_pipe].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  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', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': 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
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__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", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "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
version = 2 + variant = b) + __end__([__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": 1, + "type": "schema", + "data": "ParallelInput" + }, + { + "id": 2, + "type": "schema", + "data": "ParallelOutput" + }, + { + "id": 3, + "type": "runnable", + "data": { + "id": [ + "langchain", + "prompts", + "prompt", + "PromptTemplate" + ], + "name": "PromptTemplate" + } + }, + { + "id": 4, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "language_models", + "fake", + "FakeStreamingListLLM" + ], + "name": "FakeStreamingListLLM" + } + }, + { + "id": 5, + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "agent_parser" + } + }, + { + "id": 6, + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnablePassthrough" + ], + "name": "Passthrough" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": 3, + "target": 4 + }, + { + "source": 4, + "target": 5 + }, + { + "source": 1, + "target": 3 + }, + { + "source": 5, + "target": 2 + }, + { + "source": 1, + "target": 6 + }, + { + "source": 6, + "target": 2 + }, + { + "source": "__start__", + "target": 1 + }, + { + "source": "tools", + "target": 1 + }, + { + "source": 2, + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": 2, + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[sqlite].4 + ''' + graph TD; + PromptTemplate --> FakeStreamingListLLM; + FakeStreamingListLLM --> agent_parser; + Parallel_agent_outcome_Input --> PromptTemplate; + agent_parser --> Parallel_agent_outcome_Output; + Parallel_agent_outcome_Input --> Passthrough; + Passthrough --> Parallel_agent_outcome_Output; + __start__ --> Parallel_agent_outcome_Input; + tools --> Parallel_agent_outcome_Input; + Parallel_agent_outcome_Output -.  continue  .-> tools; + Parallel_agent_outcome_Output -.  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', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': 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
__interrupt = after) + tools(tools
version = 2 + variant = b) + __end__([__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 '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- @@ -682,6 +2198,330 @@ ''' # --- +# name: test_conditional_state_graph[memory] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[memory].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_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", + "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] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_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", + "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] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[postgres_pipe].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_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", + "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[sqlite] + '{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_state_graph[sqlite].1 + '{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "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.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' +# --- +# name: test_conditional_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", + "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_dynamic_interrupt ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% @@ -713,6 +2553,58 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query --> retriever_two; + + ''' +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class ''' graph TD; @@ -796,6 +2688,286 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].1 + dict({ + 'definitions': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/definitions/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2 ''' graph TD; @@ -866,6 +3038,286 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + ]), + 'title': 'Input', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].2 + dict({ + 'properties': dict({ + 'answer': dict({ + 'title': 'Answer', + 'type': 'string', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + }), + 'required': list([ + 'answer', + 'docs', + ]), + 'title': 'Output', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' graph TD; @@ -879,6 +3331,58 @@ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_pipe] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- # name: test_message_graph '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' # --- @@ -960,6 +3464,330 @@ ''' # --- +# name: test_message_graph[memory] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[memory].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[memory].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "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] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "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] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pipe].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[postgres_pipe].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "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[sqlite] + '{"title": "LangGraphInput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[sqlite].1 + '{"title": "LangGraphOutput", "type": "array", "items": {"anyOf": [{"$ref": "#/definitions/AIMessage"}, {"$ref": "#/definitions/HumanMessage"}, {"$ref": "#/definitions/ChatMessage"}, {"$ref": "#/definitions/SystemMessage"}, {"$ref": "#/definitions/FunctionMessage"}, {"$ref": "#/definitions/ToolMessage"}]}, "definitions": {"ToolCall": {"title": "ToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "object"}, "id": {"title": "Id", "type": "string"}, "type": {"title": "Type", "enum": ["tool_call"], "type": "string"}}, "required": ["name", "args", "id"]}, "InvalidToolCall": {"title": "InvalidToolCall", "type": "object", "properties": {"name": {"title": "Name", "type": "string"}, "args": {"title": "Args", "type": "string"}, "id": {"title": "Id", "type": "string"}, "error": {"title": "Error", "type": "string"}, "type": {"title": "Type", "enum": ["invalid_tool_call"], "type": "string"}}, "required": ["name", "args", "id", "error"]}, "UsageMetadata": {"title": "UsageMetadata", "type": "object", "properties": {"input_tokens": {"title": "Input Tokens", "type": "integer"}, "output_tokens": {"title": "Output Tokens", "type": "integer"}, "total_tokens": {"title": "Total Tokens", "type": "integer"}}, "required": ["input_tokens", "output_tokens", "total_tokens"]}, "AIMessage": {"title": "AIMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "ai", "enum": ["ai"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}, "tool_calls": {"title": "Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/ToolCall"}}, "invalid_tool_calls": {"title": "Invalid Tool Calls", "default": [], "type": "array", "items": {"$ref": "#/definitions/InvalidToolCall"}}, "usage_metadata": {"$ref": "#/definitions/UsageMetadata"}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "human", "enum": ["human"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "example": {"title": "Example", "default": false, "type": "boolean"}}, "required": ["content"]}, "ChatMessage": {"title": "ChatMessage", "description": "Message that can be assigned an arbitrary speaker (i.e. role).", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "chat", "enum": ["chat"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "role": {"title": "Role", "type": "string"}}, "required": ["content", "role"]}, "SystemMessage": {"title": "SystemMessage", "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))", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "system", "enum": ["system"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content"]}, "FunctionMessage": {"title": "FunctionMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "function", "enum": ["function"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "name"]}, "ToolMessage": {"title": "ToolMessage", "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.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "default": "tool", "enum": ["tool"], "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}, "tool_call_id": {"title": "Tool Call Id", "type": "string"}, "artifact": {"title": "Artifact"}, "status": {"title": "Status", "default": "success", "enum": ["success", "error"], "type": "string"}}, "required": ["content", "tool_call_id"]}}}' +# --- +# name: test_message_graph[sqlite].2 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "tests", + "test_pregel", + "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_nested_graph ''' graph TD; @@ -1347,6 +4175,78 @@ ''' # --- +# 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[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_state_graph_w_config '{"title": "LangGraphConfig", "type": "object", "properties": {"configurable": {"$ref": "#/definitions/Configurable"}}, "definitions": {"Configurable": {"title": "Configurable", "type": "object", "properties": {"tools": {"title": "Tools", "type": "array", "items": {"type": "string"}}}}}}' # --- diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index b02cbb65b..0b0bcf62f 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -74,15 +74,6 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): should produce a side effect that can be asserted. """ - serde = NoopSerializer() - - def __init__( - self, - *, - serde: Optional[SerializerProtocol] = None, - ) -> None: - super().__init__(serde=serde) - def put( self, config: RunnableConfig, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0f06fec2a..06a747a17 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -49,9 +49,6 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.serde.base import SerializerProtocol -from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer -from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.constants import ERROR, Interrupt, Send from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph @@ -71,9 +68,7 @@ from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, - MemorySaverAssertImmutable, MemorySaverNoPending, - NoopSerializer, ) from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -1248,7 +1243,16 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -def test_invoke_checkpoint(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_invoke_checkpoint_two( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1271,8 +1275,6 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaverAssertImmutable() - app = Pregel( nodes={"one": one}, channels={ @@ -1282,26 +1284,26 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: }, input_channels="input", output_channels="output", - checkpointer=memory, + checkpointer=checkpointer, retry_policy=RetryPolicy(), ) # total starts out as 0, so output is 0+2=2 assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) + checkpoint = checkpointer.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 assert errored_once, "errored and retried" - checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): app.invoke(4, {"configurable": {"thread_id": "1"}}) # checkpoint is not updated, error is recorded - checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 assert checkpoint_tup.pending_writes == [ @@ -1309,10 +1311,10 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: ] # on a new thread, total starts out as 0, so output is 0+5=5 assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) + checkpoint = checkpointer.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 - checkpoint = memory.get({"configurable": {"thread_id": "2"}}) + checkpoint = checkpointer.get({"configurable": {"thread_id": "2"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 5 @@ -1599,7 +1601,14 @@ async def test_checkpointer_null_pending_writes() -> None: ] * 4 -def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_invoke_checkpoint_three( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1614,121 +1623,119 @@ def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - with SqliteSaver.from_conn_string(":memory:") as memory: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + ) - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert app.invoke(2, thread_1, debug=1) == 2 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert state.next == () - assert ( - state.config["configurable"]["checkpoint_id"] == memory.get(thread_1)["id"] - ) - # total is now 2, so output is 2+3=5 - assert app.invoke(3, thread_1) == 5 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] == memory.get(thread_1)["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - app.invoke(4, thread_1) - # checkpoint is updated with new input - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert app.invoke(2, thread_1) == 9 - state = app.get_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert app.invoke(2, thread_1, debug=1) == 2 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert state.next == () + assert ( + state.config["configurable"]["checkpoint_id"] + == checkpointer.get(thread_1)["id"] + ) + # total is now 2, so output is 2+3=5 + assert app.invoke(3, thread_1) == 5 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == checkpointer.get(thread_1)["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + app.invoke(4, thread_1) + # checkpoint is updated with new input + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert app.invoke(2, thread_1) == 9 + state = app.get_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert app.invoke(5, thread_2, debug=True) == 5 - state = app.get_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == (), "checkpoint of other thread not touched" - state = app.get_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert app.invoke(5, thread_2, debug=True) == 5 + state = app.get_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == (), "checkpoint of other thread not touched" + state = app.get_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () - assert len(list(app.get_state_history(thread_1, limit=1))) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c for c in app.get_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] - ) - # cursor pagination - cursored = list( - app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config) - ) - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert ( - memory.get(thread_1_history[0].config)["id"] - == thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - assert ( - memory.get(thread_1_history[1].config)["id"] - == thread_1_history[1].config["configurable"]["checkpoint_id"] - ) + assert len(list(app.get_state_history(thread_1, limit=1))) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c for c in app.get_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = list( + app.get_state_history(thread_1, limit=1, before=thread_1_history[0].config) + ) + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert ( + checkpointer.get(thread_1_history[0].config)["id"] + == thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + assert ( + checkpointer.get(thread_1_history[1].config)["id"] + == thread_1_history[1].config["configurable"]["checkpoint_id"] + ) - thread_1_next_config = app.update_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # update makes new checkpoint child of the previous one - assert ( - app.get_state(thread_1_next_config).parent_config - == thread_1_history[1].config - ) - # 1 more checkpoint in history - assert len(list(app.get_state_history(thread_1))) == 8 - assert Counter( - c.metadata["source"] for c in app.get_state_history(thread_1) - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert app.get_state(thread_1) == app.get_state(thread_1_next_config) + thread_1_next_config = app.update_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # update makes new checkpoint child of the previous one + assert ( + app.get_state(thread_1_next_config).parent_config == thread_1_history[1].config + ) + # 1 more checkpoint in history + assert len(list(app.get_state_history(thread_1))) == 8 + assert Counter(c.metadata["source"] for c in app.get_state_history(thread_1)) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert app.get_state(thread_1) == app.get_state(thread_1_next_config) def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -1927,7 +1934,13 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -def test_conditional_graph(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_conditional_graph( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -1936,6 +1949,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: 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: @@ -1973,7 +1990,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ) if data.get("intermediate_steps") is None: data["intermediate_steps"] = [] - data["intermediate_steps"].append((agent_action, observation)) + data["intermediate_steps"].append([agent_action, observation]) return data # Define decision-making logic @@ -2009,22 +2026,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: 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" @@ -2045,14 +2062,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2060,14 +2077,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2080,22 +2097,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ), + ], ], } }, @@ -2103,22 +2120,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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" @@ -2130,7 +2147,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2238,14 +2255,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ) + ] ], } }, @@ -2253,14 +2270,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2276,14 +2293,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: { "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"}, @@ -2297,14 +2314,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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"}, @@ -2323,14 +2340,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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"}, @@ -2345,7 +2362,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2445,14 +2462,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ) + ] ], } }, @@ -2460,14 +2477,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2483,14 +2500,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: { "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"}, @@ -2504,14 +2521,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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"}, @@ -2530,14 +2547,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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"}, @@ -2552,10 +2569,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) - config = {"configurable": {"thread_id": "2"}} + config = {"configurable": {"thread_id": "3"}} llm.i = 0 # reset the llm assert [ @@ -2608,14 +2625,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2623,14 +2640,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", @@ -2646,22 +2663,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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", - ), + ], ], } }, @@ -2669,22 +2686,22 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "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" @@ -2780,14 +2797,24 @@ def test_conditional_entrypoint_to_multiple_state_graph( } +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_conditional_state_graph( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + 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 + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) setup = mocker.Mock() teardown = mocker.Mock() @@ -2870,7 +2897,7 @@ def test_conditional_state_graph( observation = {t.name: t for t in tools}[agent_action.tool].invoke( agent_action.tool_input ) - return {"intermediate_steps": [(agent_action, observation)]} + return {"intermediate_steps": [[agent_action, observation]]} # Define decision-making logic def should_continue(data: AgentState) -> str: @@ -2907,22 +2934,22 @@ def test_conditional_state_graph( 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" @@ -2943,14 +2970,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2966,14 +2993,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="another", log="tool:search_api:another", ), "result for another", - ), + ], ], } }, @@ -2989,7 +3016,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -3083,14 +3110,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3123,14 +3150,14 @@ def test_conditional_state_graph( 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=(), @@ -3155,7 +3182,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], debug=True, ) @@ -3245,14 +3272,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3284,14 +3311,14 @@ def test_conditional_state_graph( 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=(), @@ -3315,7 +3342,7 @@ def test_conditional_state_graph( # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before="*", debug=True, ) @@ -3379,14 +3406,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3398,14 +3425,14 @@ def test_conditional_state_graph( 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"),), @@ -3418,14 +3445,14 @@ def test_conditional_state_graph( "writes": { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3447,7 +3474,7 @@ def test_conditional_state_graph( # test w interrupt after all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after="*", ) config = {"configurable": {"thread_id": "4"}} @@ -3496,14 +3523,14 @@ def test_conditional_state_graph( { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3515,14 +3542,14 @@ def test_conditional_state_graph( 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"),), @@ -3535,14 +3562,14 @@ def test_conditional_state_graph( "writes": { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -4010,8 +4037,13 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize("serde", [NoopSerializer(), JsonPlusSerializer()]) -def test_state_graph_packets(serde: SerializerProtocol) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_state_graph_packets( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -4024,6 +4056,10 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ) from langchain_core.tools import tool + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] @@ -4244,7 +4280,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(serde=serde), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -4521,9 +4557,15 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from copy import deepcopy @@ -4540,6 +4582,10 @@ def test_message_graph( 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 @@ -4743,7 +4789,7 @@ def test_message_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -4973,7 +5019,7 @@ def test_message_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -5236,15 +5282,20 @@ def test_message_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_root_graph( - snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from copy import deepcopy @@ -5261,6 +5312,10 @@ def test_root_graph( 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 @@ -5462,7 +5517,7 @@ def test_root_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -5692,7 +5747,7 @@ def test_root_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -5955,7 +6010,7 @@ def test_root_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5991,7 +6046,7 @@ def test_root_graph( }, ) new_workflow.add_edge("tools", "agent") - new_app = new_workflow.compile(checkpointer=app_w_interrupt.checkpointer) + new_app = new_workflow.compile(checkpointer=checkpointer) model.i = 0 # reset the llm # previous state is converted to new schema @@ -6027,7 +6082,7 @@ def test_root_graph( metadata={ "source": "update", "step": 6, - "writes": {"tools": ("ai", "an extra message")}, + "writes": {"tools": UnsortedSequence("ai", "an extra message")}, }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6310,7 +6365,15 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +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 @@ -6348,49 +6411,56 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile(checkpointer=saver) + 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"}) + # 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 ⛰️", - "market": "DE", - } - assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value ⛰️", "market": "DE"}, - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - interrupts=(Interrupt("Just because..."),), - ), + 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)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("Just because..."),), ), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) -def test_start_branch_then(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +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 @@ -6435,148 +6505,155 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=saver, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + tool_two = tool_two_graph.compile( + store=MemoryStore(), + 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"}) + # 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)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value ⛰️", "market": "DE"}, - }, - ] - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "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={"source": "loop", "step": 0, "writes": None}, - 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={ - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + 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)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "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={"source": "loop", "step": 0, "writes": None}, + 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={ + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + 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"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - 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={ - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread2, 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"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + 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={ + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + 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"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={"source": "loop", "step": 0, "writes": None}, - 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"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - }, - 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={ - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread3, 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"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={"source": "loop", "step": 0, "writes": None}, + 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"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + }, + 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={ + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) -def test_branch_then(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +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 @@ -6606,504 +6683,509 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "market": "US", } - with SqliteSaver.from_conn_string(":memory:") as saver: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=saver) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - *tool_two.stream( - {"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": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], - }, - }, - { - "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": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "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": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "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": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "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": { - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "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": { - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) - - # 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"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + ] == [ + { + "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": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "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": { + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - 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"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) + }, + { + "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": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "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": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "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": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "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": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], + }, + }, + ] - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["finish"] - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] + ) - thread1 = {"configurable": {"thread_id": "1"}} + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { + 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"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + 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"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + 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", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={ - "my_key": "value prepared slow", - "market": "DE", - }, - tasks=(PregelTask(AnyStr(), "finish"),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + }, + tasks=(PregelTask(AnyStr(), "finish"),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + 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"),), - next=("finish",), - config=tool_two.checkpointer.get_tuple(thread1).config, - created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={ - "source": "update", - "step": 3, - "writes": {"tool_two_slow": {"my_key": "er"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) - - with SqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, 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": "1"}} - # stop when about to enter node - assert tool_two.invoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", + # 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", - } - assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, - ) + }, + tasks=(PregelTask(AnyStr(), "finish"),), + next=("finish",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "update", + "step": 3, + "writes": {"tool_two_slow": {"my_key": "er"}}, + }, + 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"),), - next=("tool_two_fast",), - config=tool_two.checkpointer.get_tuple(thread2).config, - created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, - ) + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) - thread3 = {"configurable": {"thread_id": "3"}} - # 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"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - }, - 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"),), - next=("tool_two_slow",), - config=tool_two.checkpointer.get_tuple(thread3).config, - created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, - ) + # 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"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + 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"),), + next=("tool_two_fast",), + config=tool_two.checkpointer.get_tuple(thread2).config, + created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + 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"),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + }, + 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"),), + next=("tool_two_slow",), + config=tool_two.checkpointer.get_tuple(thread3).config, + created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, + ) -def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_in_one_fan_out_state_graph_waiting_edge( + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7168,7 +7250,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7187,10 +7269,10 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["qa"], ) - config = {"configurable": {"thread_id": "1"}} + config = {"configurable": {"thread_id": "2"}} assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) @@ -7224,9 +7306,17 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str ) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7294,7 +7384,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7313,11 +7403,19 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() teardown = mocker.Mock() @@ -7442,7 +7540,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7477,11 +7575,19 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( } +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, + mocker: MockerFixture, + request: pytest.FixtureRequest, + checkpointer_name: str, ) -> None: from pydantic import BaseModel, ConfigDict, ValidationError + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") setup = mocker.Mock() teardown = mocker.Mock() @@ -7604,7 +7710,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -7639,7 +7745,17 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( } -def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -7708,7 +7824,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d24a3d4f2..4eed26017 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -46,7 +46,6 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.constants import ERROR, Interrupt, Send from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph, StateGraph @@ -1803,7 +1802,14 @@ async def test_cond_edge_after_send() -> None: assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] -async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_invoke_checkpoint_three( + mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1818,121 +1824,118 @@ async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: | raise_if_above_10 ) - async with AsyncSqliteSaver.from_conn_string(":memory:") as memory: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - debug=True, - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + debug=True, + ) - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, thread_1) == 2 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await memory.aget(thread_1))["id"] - ) - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, thread_1) == 5 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await memory.aget(thread_1))["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, thread_1) - # checkpoint is not updated - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert await app.ainvoke(2, thread_1) == 9 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1) == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await checkpointer.aget(thread_1))["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1) + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1) == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, thread_2) == 5 - state = await app.aget_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == () - state = await app.aget_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == () + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config ) - # cursor pagination - cursored = [ - c - async for c in app.aget_state_history( - thread_1, limit=1, before=thread_1_history[0].config - ) - ] - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert (await memory.aget(thread_1_history[0].config))[ - "id" - ] == thread_1_history[0].config["configurable"]["checkpoint_id"] - assert (await memory.aget(thread_1_history[1].config))[ - "id" - ] == thread_1_history[1].config["configurable"]["checkpoint_id"] + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await checkpointer.aget(thread_1_history[0].config))[ + "id" + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] + assert (await checkpointer.aget(thread_1_history[1].config))[ + "id" + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] - thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # 1 more checkpoint in history - assert len([c async for c in app.aget_state_history(thread_1)]) == 8 - assert Counter( - [c.metadata["source"] async for c in app.aget_state_history(thread_1)] - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert await app.aget_state(thread_1) == await app.aget_state( - thread_1_next_config - ) + thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # 1 more checkpoint in history + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state(thread_1_next_config) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -2144,7 +2147,13 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -async def test_conditional_graph() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_conditional_graph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -2153,6 +2162,8 @@ async def test_conditional_graph() -> None: from langchain_core.runnables import RunnablePassthrough from langchain_core.tools import tool + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + # Assemble the tools @tool() def search_api(query: str) -> str: @@ -2190,7 +2201,7 @@ async def test_conditional_graph() -> None: ) if data.get("intermediate_steps") is None: data["intermediate_steps"] = [] - data["intermediate_steps"].append((agent_action, observation)) + data["intermediate_steps"].append([agent_action, observation]) return data # Define decision-making logic @@ -2220,22 +2231,22 @@ async def test_conditional_graph() -> None: 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" @@ -2258,14 +2269,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2273,14 +2284,14 @@ async def test_conditional_graph() -> None: "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", @@ -2293,22 +2304,22 @@ async def test_conditional_graph() -> None: "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", - ), + ], ], } }, @@ -2316,22 +2327,22 @@ async def test_conditional_graph() -> None: "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" @@ -2368,14 +2379,14 @@ async def test_conditional_graph() -> None: { "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", @@ -2386,22 +2397,22 @@ async def test_conditional_graph() -> None: { "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" @@ -2412,7 +2423,7 @@ async def test_conditional_graph() -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2522,14 +2533,14 @@ async def test_conditional_graph() -> None: "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", - ) + ] ], } }, @@ -2537,14 +2548,14 @@ async def test_conditional_graph() -> None: "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", @@ -2560,14 +2571,14 @@ async def test_conditional_graph() -> None: { "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"}, @@ -2581,14 +2592,14 @@ async def test_conditional_graph() -> None: "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"}, @@ -2609,14 +2620,14 @@ async def test_conditional_graph() -> None: "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"}, @@ -2633,7 +2644,7 @@ async def test_conditional_graph() -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2744,14 +2755,14 @@ async def test_conditional_graph() -> None: "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", - ) + ] ], } }, @@ -2759,14 +2770,14 @@ async def test_conditional_graph() -> None: "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", @@ -2782,14 +2793,14 @@ async def test_conditional_graph() -> None: { "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"}, @@ -2803,14 +2814,14 @@ async def test_conditional_graph() -> None: "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"}, @@ -2831,14 +2842,14 @@ async def test_conditional_graph() -> None: "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"}, @@ -2855,10 +2866,10 @@ async def test_conditional_graph() -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), + checkpointer=checkpointer, interrupt_before=["tools"], ) - config = {"configurable": {"thread_id": "2"}} + config = {"configurable": {"thread_id": "3"}} llm.i = 0 # reset the llm assert [ @@ -2918,14 +2929,14 @@ async def test_conditional_graph() -> None: "tools": { "input": "what is weather in sf", "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -2933,14 +2944,14 @@ async def test_conditional_graph() -> None: "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", @@ -2956,22 +2967,22 @@ async def test_conditional_graph() -> None: "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", - ), + ], ], } }, @@ -2979,22 +2990,22 @@ async def test_conditional_graph() -> None: "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" @@ -5005,7 +5016,15 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -async def test_start_branch_then() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_start_branch_then( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5050,176 +5069,169 @@ async def test_start_branch_then() -> None: "market": "US", } - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=saver, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + tool_two = tool_two_graph.compile( + store=MemoryStore(), + 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"}) + # 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)] == [ - { - "source": "loop", - "step": 0, - "writes": None, - }, - { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "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={"source": "loop", "step": 0, "writes": None}, - 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={ - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + 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)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "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={"source": "loop", "step": 0, "writes": None}, + 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={ + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + 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"),), - 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={"source": "loop", "step": 0, "writes": None}, - 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={ - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, 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"),), + 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={"source": "loop", "step": 0, "writes": None}, + 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={ + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + 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"),), - 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={"source": "loop", "step": 0, "writes": None}, - 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"),), - 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={ - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - }, - 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={ - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, 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"),), + 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={"source": "loop", "step": 0, "writes": None}, + 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"),), + 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={ + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + }, + 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={ + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) -async def test_branch_then() -> None: +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_branch_then( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5247,606 +5259,578 @@ async def test_branch_then() -> None: "market": "US", } - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=saver) - 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(), - }, + # 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": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - }, - { - "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": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", - "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(), - }, + }, + { + "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 prepared", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:prepare:condition:tool_two_slow"], + "values": { + "my_key": "value", + "market": "DE", }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], + "metadata": { + "source": "loop", + "step": 0, + "writes": None, }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "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(), - }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "7b7b0713-e958-5d07-803c-c9910a7cc162", + "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 slow", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - }, - "next": ["finish"], - "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "values": { + "my_key": "value prepared", + "market": "DE", }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], + "metadata": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "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(), - }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "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": "dd9f2fa5-ccfa-5d12-81ec-942563056a08", + "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 finished", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + }, + "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}], }, - ] + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition::then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52", + "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": { + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], + }, + }, + ] - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] + 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" ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + ] == [ + { + "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": { - "source": "input", - "step": -1, - "writes": {"my_key": "value", "market": "DE"}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + "values": {"my_key": ""}, + "metadata": { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70", - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "1"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, + }, + { + "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 prepared", - "market": "DE", - }, - "metadata": { - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [ - {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()} - ], }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + }, + { + "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": { + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + }, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "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={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) - thread2 = {"configurable": {"thread_id": "2"}} - # 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"),), - 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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, 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"),), + 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={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) - async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - tool_two = tool_two_graph.compile( - checkpointer=saver, interrupt_after=["prepare"] - ) + 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"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1"}} - # 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"),), - 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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - ) + 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"),), + 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={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + ) - thread2 = {"configurable": {"thread_id": "2"}} - # 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"),), - 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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread2, 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"),), + 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={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ + -1 + ].config, + ) - thread3 = {"configurable": {"thread_id": "3"}} - # 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"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, - }, - 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"),), - 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={ - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - 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={ - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread3, 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"),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + }, + 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"),), + 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={ + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + 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={ + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ + -1 + ].config, + ) async def test_in_one_fan_out_state_graph_waiting_edge() -> None: