From 3b56cdf5248bceb1c44b2bc67702fa5e7fc83afc Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 13 Aug 2024 17:35:31 -0700 Subject: [PATCH 01/16] Add Runs.join_stream endpoint --- libs/sdk-py/langgraph_sdk/client.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 13d5560db..9163522b5 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -1449,6 +1449,28 @@ class RunsClient: """ # noqa: E501 return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join") + def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]: + """Stream output from a run in real-time, until the run is done. + Output is not buffered, so any output produced before this call will + not be received here. + + Args: + thread_id: The thread ID to join. + run_id: The run ID to join. + + Returns: + None + + Example Usage: + + await client.runs.join( + thread_id="thread_id_to_join", + run_id="run_id_to_join" + ) + + """ # noqa: E501 + return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET") + async def delete(self, thread_id: str, run_id: str) -> None: """Delete a run. From f93512e3b33346e67880bb55e3a2649147ca64cf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 12:33:25 -0700 Subject: [PATCH 02/16] Remove unused when values for Interrupt --- libs/langgraph/langgraph/constants.py | 4 ++-- libs/langgraph/langgraph/errors.py | 2 +- libs/langgraph/tests/test_pregel.py | 2 +- libs/langgraph/tests/test_pregel_async.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 51b6e1437..21844a565 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -102,5 +102,5 @@ class Send: @dataclass class Interrupt: - when: Literal["before", "during", "after"] - value: Any = None + value: Any + when: Literal["during"] = "during" diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 27a7689fa..25ba93a72 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -40,7 +40,7 @@ class NodeInterrupt(GraphInterrupt): """Raised by a node to interrupt execution.""" def __init__(self, value: Any) -> None: - super().__init__([Interrupt("during", value)]) + super().__init__([Interrupt(value)]) class EmptyInputError(Exception): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 56e7ed5a3..015bba9a6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6383,7 +6383,7 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None: PregelTask( AnyStr(), "tool_two", - interrupts=(Interrupt("during", "Just because..."),), + interrupts=(Interrupt("Just because..."),), ), ), config=tool_two.checkpointer.get_tuple(thread1).config, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0308deb54..26bea855a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -285,7 +285,7 @@ async def test_dynamic_interrupt( PregelTask( AnyStr(), "tool_two", - interrupts=(Interrupt("during", "Just because..."),), + interrupts=(Interrupt("Just because..."),), ), ), config=tup.config, From 7e32de940548574d35816a80e28bbd1f9581e261 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 12:44:26 -0700 Subject: [PATCH 03/16] Remove current_tasks from checkpoint interface (#1440) * Remove current_tasks from checkpoint interface - Not used, now clear that it can be supported with put_writes(SCHEDULE) * Add comment --- libs/checkpoint-postgres/README.md | 2 -- libs/checkpoint-sqlite/README.md | 4 +--- libs/checkpoint/README.md | 1 - libs/checkpoint/langgraph/checkpoint/base/__init__.py | 7 +------ libs/langgraph/langgraph/pregel/loop.py | 4 ---- libs/langgraph/tests/test_pregel.py | 3 --- libs/langgraph/tests/test_pregel_async.py | 3 --- 7 files changed, 2 insertions(+), 22 deletions(-) diff --git a/libs/checkpoint-postgres/README.md b/libs/checkpoint-postgres/README.md index 24a77673a..24652a2b2 100644 --- a/libs/checkpoint-postgres/README.md +++ b/libs/checkpoint-postgres/README.md @@ -44,7 +44,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -87,7 +86,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint diff --git a/libs/checkpoint-sqlite/README.md b/libs/checkpoint-sqlite/README.md index 86ca9cafc..73fe94333 100644 --- a/libs/checkpoint-sqlite/README.md +++ b/libs/checkpoint-sqlite/README.md @@ -35,7 +35,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -78,7 +77,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint @@ -89,4 +87,4 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: # list checkpoints [c async for c in checkpointer.alist(read_config)] -``` \ No newline at end of file +``` diff --git a/libs/checkpoint/README.md b/libs/checkpoint/README.md index 7f6b26d6e..19c7d3807 100644 --- a/libs/checkpoint/README.md +++ b/libs/checkpoint/README.md @@ -74,7 +74,6 @@ checkpoint = { } }, "pending_sends": [], - "current_tasks": {} } # store checkpoint diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 21d021ad5..a4f76f87b 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -97,9 +97,6 @@ class Checkpoint(TypedDict): pending_sends: List[SendProtocol] """List of packets sent to nodes but not yet processed. Cleared by the next checkpoint.""" - current_tasks: Dict[str, TaskInfo] - """Map from task ID to task info.""" - # TODO remove this def empty_checkpoint() -> Checkpoint: @@ -111,7 +108,6 @@ def empty_checkpoint() -> Checkpoint: channel_versions={}, versions_seen={}, pending_sends=[], - current_tasks={}, ) @@ -124,7 +120,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: channel_versions=checkpoint["channel_versions"].copy(), versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()}, pending_sends=checkpoint.get("pending_sends", []).copy(), - current_tasks=checkpoint.get("current_tasks", {}).copy(), ) @@ -156,7 +151,6 @@ def create_checkpoint( channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), - current_tasks={}, ) @@ -451,3 +445,4 @@ saving regular writes. Each Checkpointer implementation should use this mapping in put_writes. """ WRITES_IDX_MAP = {ERROR: -1} +# TODO To store scheduled status of tasks, add a special channel here diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a501232a9..b4e6d4118 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -155,10 +155,6 @@ class PregelLoop: self.stream_keys = stream_keys self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) - def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: - """Mark tasks as scheduled, to be used by queue-based executors.""" - raise NotImplementedError - def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" if not writes: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 015bba9a6..0f06fec2a 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1428,7 +1428,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "one": { @@ -1487,7 +1486,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "__input__": {}, @@ -1535,7 +1533,6 @@ def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 26bea855a..d24a3d4f2 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1654,7 +1654,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "one": { @@ -1713,7 +1712,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": { "__input__": {}, @@ -1761,7 +1759,6 @@ async def test_pending_writes_resume( "v": 1, "id": AnyStr(), "ts": AnyStr(), - "current_tasks": {}, "pending_sends": [], "versions_seen": {"__input__": {}}, "channel_versions": { From 8a00a0026ed09a4ca6d65c517ab4f2e10810c343 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 12:45:20 -0700 Subject: [PATCH 04/16] checkpoint1.0.4 --- libs/checkpoint/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index b17a77238..7ceea436e 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "1.0.3" +version = "1.0.4" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" From 3ec419a2f621a33e242b16682f22341e1d5dd1c6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 12:45:27 -0700 Subject: [PATCH 05/16] lib0.2.12 --- libs/langgraph/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 1c3b20e9b..bc0b67fb4 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.11" +version = "0.2.12" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" From 45e3d1a3f168e997baaf820af545c1f9ed172433 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 13:51:38 -0700 Subject: [PATCH 06/16] 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: From 70746042049e684d751e8f150803589a69306a74 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 14:23:48 -0700 Subject: [PATCH 07/16] Try to improve async stack traces for exceptions in tasks (#1442) * Try to improve async stack traces for exceptions in tasks * Lint --- libs/langgraph/langgraph/pregel/executor.py | 23 +++++++++++---------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 5da78a26c..3e4243c61 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -120,14 +120,14 @@ class AsyncBackgroundExecutor(AsyncContextManager): def done(self, task: asyncio.Task) -> None: try: - task.result() - except GraphInterrupt: - # This exception is an interruption signal, not an error - # so we don't want to re-raise it on exit - self.tasks.pop(task) - except BaseException: - pass - else: + if exc := task.exception(): + # This exception is an interruption signal, not an error + # so we don't want to re-raise it on exit + if isinstance(exc, GraphInterrupt): + self.tasks.pop(task) + else: + self.tasks.pop(task) + except asyncio.CancelledError: self.tasks.pop(task) async def __aenter__(self) -> Submit: @@ -146,12 +146,13 @@ class AsyncBackgroundExecutor(AsyncContextManager): # wait for all tasks to finish if self.tasks: await asyncio.wait(self.tasks) - # re-raise the first exception that occurred in a task + # if there's already an exception being raised, don't raise another one if exc_type is None: - # if there's already an exception being raised, don't raise another one + # re-raise the first exception that occurred in a task for task in self.tasks: try: - task.result() + if exc := task.exception(): + raise exc except asyncio.CancelledError: pass From 8090ca67c5c78c50d1c5a620df57cd866dc1da29 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Thu, 22 Aug 2024 17:39:21 -0400 Subject: [PATCH 08/16] checkpoint-postgres: pass row_factory in cursor (#1433) --- examples/persistence_postgres.ipynb | 7 +- .../langgraph/checkpoint/postgres/__init__.py | 62 +++++++++-------- .../langgraph/checkpoint/postgres/aio.py | 66 ++++++++++--------- 3 files changed, 71 insertions(+), 64 deletions(-) diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index ef8be17cb..f43fa01d6 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -122,7 +122,7 @@ "metadata": {}, "outputs": [], "source": [ - "DB_URI = \"postgresql://postgres:postgres@localhost:5441/postgres?sslmode=disable\"" + "DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\"" ] }, { @@ -137,7 +137,6 @@ "connection_kwargs ={\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", - " \"row_factory\": dict_row,\n", "}" ] }, @@ -551,9 +550,9 @@ ], "metadata": { "kernelspec": { - "display_name": "langgraph-postgres", + "display_name": "langgraph", "language": "python", - "name": "langgraph-postgres" + "name": "langgraph" }, "language_info": { "codemirror_mode": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 433bc231d..cb4a79c35 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -66,7 +66,7 @@ class PostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ with self.lock: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: version = cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -127,31 +127,33 @@ class PostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - for value in self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + with self._cursor() as cur: + cur.execute(query, args, binary=True) + for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": self._load_blobs(value["channel_values"]), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": self._load_blobs(value["channel_values"]), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - self._load_writes(value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + self._load_writes(value["pending_writes"]), + ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -199,7 +201,7 @@ class PostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" with self._cursor() as cur: - cur = self.conn.execute( + cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -336,7 +338,7 @@ class PostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - with self.conn.cursor(binary=True) as cur: + with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -344,8 +346,10 @@ class PostgresSaver(BasePostgresSaver): elif pipeline: # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock - with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.pipeline(), self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur else: - with self.lock, self.conn.cursor(binary=True) as cur: + with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 79ae1ddf7..569159d91 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver): the first time checkpointer is used. """ async with self.lock: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: try: results = await cur.execute( "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" @@ -110,33 +110,35 @@ class AsyncPostgresSaver(BasePostgresSaver): if limit: query += f" LIMIT {limit}" # if we change this to use .stream() we need to make sure to close the cursor - async for value in await self.conn.execute(query, args, binary=True): - yield CheckpointTuple( - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["checkpoint_id"], + async with self._cursor() as cur: + await cur.execute(query, args, binary=True) + async for value in cur: + yield CheckpointTuple( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["checkpoint_id"], + } + }, + { + **self._load_checkpoint(value["checkpoint"]), + "channel_values": await asyncio.to_thread( + self._load_blobs, value["channel_values"] + ), + }, + self._load_metadata(value["metadata"]), + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - }, - { - **self._load_checkpoint(value["checkpoint"]), - "channel_values": await asyncio.to_thread( - self._load_blobs, value["channel_values"] - ), - }, - self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], - } - } - if value["parent_checkpoint_id"] - else None, - await asyncio.to_thread(self._load_writes, value["pending_writes"]), - ) + if value["parent_checkpoint_id"] + else None, + await asyncio.to_thread(self._load_writes, value["pending_writes"]), + ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database asynchronously. @@ -163,7 +165,7 @@ class AsyncPostgresSaver(BasePostgresSaver): where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1" async with self._cursor() as cur: - cur = await self.conn.execute( + await cur.execute( self.SELECT_SQL + where, args, binary=True, @@ -293,7 +295,7 @@ class AsyncPostgresSaver(BasePostgresSaver): # in multiple threads/coroutines, but only one cursor can be # used at a time try: - async with self.conn.cursor(binary=True) as cur: + async with self.conn.cursor(binary=True, row_factory=dict_row) as cur: yield cur finally: if pipeline: @@ -302,9 +304,11 @@ class AsyncPostgresSaver(BasePostgresSaver): # a connection not in pipeline mode can only be used by one # thread/coroutine at a time, so we acquire a lock async with self.lock, self.conn.pipeline(), self.conn.cursor( - binary=True + binary=True, row_factory=dict_row ) as cur: yield cur else: - async with self.lock, self.conn.cursor(binary=True) as cur: + async with self.lock, self.conn.cursor( + binary=True, row_factory=dict_row + ) as cur: yield cur From 75dec9b924d1c91818a4593c292286c19a359b9f Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Fri, 23 Aug 2024 00:18:06 +0200 Subject: [PATCH 09/16] fix(sdk-js): support sending end events --- libs/sdk-js/src/client.mts | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/libs/sdk-js/src/client.mts b/libs/sdk-js/src/client.mts index bf98dccb4..fbf2dd12c 100644 --- a/libs/sdk-js/src/client.mts +++ b/libs/sdk-js/src/client.mts @@ -571,6 +571,7 @@ export class RunsClient extends BaseClient { ); let parser: EventSourceParser; + let onEndEvent: () => void; const textDecoder = new TextDecoder(); const stream: ReadableStream<{ event: string; data: any }> = ( @@ -594,9 +595,17 @@ export class RunsClient extends BaseClient { }); } }); + onEndEvent = () => { + ctrl.enqueue({ event: "end", data: undefined }); + }; }, async transform(chunk) { - parser.feed(textDecoder.decode(chunk)); + const payload = textDecoder.decode(chunk); + parser.feed(payload); + + // eventsource-parser will ignore events + // that are not terminated by a newline + if (payload.trim() === "event: end") onEndEvent(); }, }), ); From 0b7f451b40e0593feaf3e016265fa02d4ac3ccd4 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Fri, 23 Aug 2024 00:18:28 +0200 Subject: [PATCH 10/16] Bump to 0.0.6 --- libs/sdk-js/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/sdk-js/package.json b/libs/sdk-js/package.json index db509e9c2..a10bcdab8 100644 --- a/libs/sdk-js/package.json +++ b/libs/sdk-js/package.json @@ -1,6 +1,6 @@ { "name": "@langchain/langgraph-sdk", - "version": "0.0.5", + "version": "0.0.6", "description": "Client library for interacting with the LangGraph API", "type": "module", "packageManager": "yarn@1.22.19", From 15c3105748e4a2b0c28ea2d03ccd84f568ce13ab Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 15:36:48 -0700 Subject: [PATCH 11/16] cli0.1.51 --- libs/cli/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index 508571f06..02c373e95 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-cli" -version = "0.1.50" +version = "0.1.51" description = "CLI for interacting with LangGraph API" authors = [] license = "MIT" From c72acc91455ac81b9ddab661f6a8cb9ac21b8204 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Fri, 23 Aug 2024 01:38:53 +0200 Subject: [PATCH 12/16] Add missing status --- libs/sdk-js/src/schema.ts | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/libs/sdk-js/src/schema.ts b/libs/sdk-js/src/schema.ts index 9c2df19a9..1d772b7eb 100644 --- a/libs/sdk-js/src/schema.ts +++ b/libs/sdk-js/src/schema.ts @@ -2,6 +2,16 @@ import type { JSONSchema7 } from "json-schema"; type Optional = T | null | undefined; +type RunStatus = + | "pending" + | "running" + | "error" + | "success" + | "timeout" + | "interrupted"; + +type ThreadStatus = "idle" | "busy" | "interrupted"; + export interface Config { /** * Tags for this call and any sub-calls (eg. a Chain calling an LLM). @@ -80,6 +90,7 @@ export interface Thread { created_at: string; updated_at: string; metadata: Metadata; + status: ThreadStatus; } export interface Cron { @@ -112,12 +123,6 @@ export interface Run { assistant_id: string; created_at: string; updated_at: string; - status: - | "pending" - | "running" - | "error" - | "success" - | "timeout" - | "interrupted"; + status: RunStatus; metadata: Metadata; } From 19b382335f931920633045319102ee75b7da7c64 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 17:03:11 -0700 Subject: [PATCH 13/16] sdk: Add on_disconnect arg to create/wait streaming run --- libs/sdk-py/langgraph_sdk/client.py | 21 +++++++++++++++++---- libs/sdk-py/langgraph_sdk/schema.py | 2 ++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index e3a83c543..fb8435b5c 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -25,6 +25,7 @@ from langgraph_sdk.schema import ( Assistant, Config, Cron, + DisconnectMode, GraphSchema, Metadata, MultitaskStrategy, @@ -963,6 +964,8 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, + webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> AsyncIterator[StreamPart]: ... @@ -980,6 +983,8 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, + webhook: Optional[str] = None, ) -> AsyncIterator[StreamPart]: ... @@ -996,6 +1001,7 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, feedback_keys: Optional[list[str]] = None, + on_disconnect: Optional[DisconnectMode] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> AsyncIterator[StreamPart]: @@ -1019,6 +1025,8 @@ class RunsClient: webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. + disconnect_mode: The disconnect mode to use. + Must be one of 'cancel' or 'continue'. Returns: AsyncIterator[StreamPart]: Asynchronous iterator of stream results. @@ -1061,6 +1069,7 @@ class RunsClient: "webhook": webhook, "checkpoint_id": checkpoint_id, "multitask_strategy": multitask_strategy, + "on_disconnect": on_disconnect, } endpoint = ( f"/threads/{thread_id}/runs/stream" @@ -1129,9 +1138,7 @@ class RunsClient: config: The configuration for the assistant. checkpoint_id: The checkpoint to start streaming from. interrupt_before: Nodes to interrupt immediately before they get executed. - interrupt_after: Nodes to Nodes to interrupt immediately after they get executed. - webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. @@ -1242,6 +1249,8 @@ class RunsClient: checkpoint_id: Optional[str] = None, interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, + webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> Union[list[dict], dict[str, Any]]: ... @@ -1257,6 +1266,8 @@ class RunsClient: config: Optional[Config] = None, interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, + webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, ) -> Union[list[dict], dict[str, Any]]: ... @@ -1272,6 +1283,7 @@ class RunsClient: interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, webhook: Optional[str] = None, + on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, ) -> Union[list[dict], dict[str, Any]]: """Create a run, wait until it finishes and return the final state. @@ -1286,12 +1298,12 @@ class RunsClient: config: The configuration for the assistant. checkpoint_id: The checkpoint to start streaming from. interrupt_before: Nodes to interrupt immediately before they get executed. - interrupt_after: Nodes to Nodes to interrupt immediately after they get executed. - webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. + on_disconnect: The disconnect mode to use. + Must be one of 'cancel' or 'continue'. Returns: Union[list[dict], dict[str, Any]]: The output of the run. @@ -1351,6 +1363,7 @@ class RunsClient: "webhook": webhook, "checkpoint_id": checkpoint_id, "multitask_strategy": multitask_strategy, + "on_disconnect": on_disconnect, } endpoint = ( f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait" diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index a0fac51a0..c3232c88e 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -9,6 +9,8 @@ ThreadStatus = Literal["idle", "busy", "interrupted"] StreamMode = Literal["values", "messages", "updates", "events", "debug"] +DisconnectMode = Literal["cancel", "continue"] + MultitaskStrategy = Literal["reject", "interrupt", "rollback", "enqueue"] OnConflictBehavior = Literal["raise", "do_nothing"] From 6ece7124ed251812bc78f307d464f92d6565b3db Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 17:21:01 -0700 Subject: [PATCH 14/16] Fix param name in docstring --- libs/sdk-py/langgraph_sdk/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index fb8435b5c..24b2ef3b0 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -1025,7 +1025,7 @@ class RunsClient: webhook: Webhook to call after LangGraph API call is done. multitask_strategy: Multitask strategy to use. Must be one of 'reject', 'interrupt', 'rollback', or 'enqueue'. - disconnect_mode: The disconnect mode to use. + on_disconnect: The disconnect mode to use. Must be one of 'cancel' or 'continue'. Returns: From dec7eb6f585f2632fcac7845af81029826ec3ed5 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Thu, 22 Aug 2024 18:54:42 -0700 Subject: [PATCH 15/16] [Docs] Use injected RunnableConfig (#1444) * Pass via type * Format --- .../agent-simulation-evaluation.ipynb | 5 +- .../information-gather-prompting.ipynb | 21 ++++-- examples/create-react-agent-hitl.ipynb | 2 +- .../customer-support/customer-support.ipynb | 24 ++++--- .../human_in_the_loop/review-tool-calls.ipynb | 65 ++++++++++--------- examples/input_output_schema.ipynb | 4 ++ examples/llm-compiler/LLMCompiler.ipynb | 10 ++- examples/many-tools.ipynb | 6 +- examples/pass-run-time-values-to-tools.ipynb | 1 + examples/pass_private_state.ipynb | 4 +- examples/persistence.ipynb | 2 +- examples/persistence_mongodb.ipynb | 14 ++-- examples/persistence_postgres.ipynb | 6 +- examples/persistence_redis.ipynb | 12 +++- examples/reflection/reflection.ipynb | 2 +- examples/reflexion/reflexion.ipynb | 1 + examples/streaming-content.ipynb | 4 +- .../tutorials/rag-agent-testing-local.ipynb | 19 +++--- .../tutorials/tool-calling-agent-local.ipynb | 4 ++ 19 files changed, 127 insertions(+), 79 deletions(-) diff --git a/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb b/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb index c41653c50..02f8782eb 100644 --- a/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb +++ b/examples/chatbot-simulation-evaluation/agent-simulation-evaluation.ipynb @@ -235,7 +235,7 @@ " # Call the chat bot\n", " chat_bot_response = my_chat_bot(messages)\n", " # Respond with an AI Message\n", - " return {\"messages\":[AIMessage(content=chat_bot_response[\"content\"])]}" + " return {\"messages\": [AIMessage(content=chat_bot_response[\"content\"])]}" ] }, { @@ -270,7 +270,7 @@ " # Call the simulated user\n", " response = simulated_user.invoke({\"messages\": new_messages})\n", " # This response is an AI message - we need to flip this to be a human message\n", - " return {\"messages\":[HumanMessage(content=response.content)]}" + " return {\"messages\": [HumanMessage(content=response.content)]}" ] }, { @@ -331,6 +331,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "graph_builder = StateGraph(State)\n", "graph_builder.add_node(\"user\", simulated_user_node)\n", "graph_builder.add_node(\"chat_bot\", chat_bot_node)\n", diff --git a/examples/chatbots/information-gather-prompting.ipynb b/examples/chatbots/information-gather-prompting.ipynb index ad4cdfe35..ad375f02f 100644 --- a/examples/chatbots/information-gather-prompting.ipynb +++ b/examples/chatbots/information-gather-prompting.ipynb @@ -79,7 +79,7 @@ "\n", "\n", "def info_chain(state):\n", - " messages = get_messages_info(state['messages'])\n", + " messages = get_messages_info(state[\"messages\"])\n", " response = llm_with_tool.invoke(messages)\n", " return {\"messages\": [response]}" ] @@ -126,7 +126,7 @@ "\n", "\n", "def prompt_gen_chain(state):\n", - " messages = get_prompt_messages(state['messages'])\n", + " messages = get_prompt_messages(state[\"messages\"])\n", " response = llm.invoke(messages)\n", " return {\"messages\": [response]}" ] @@ -158,7 +158,7 @@ "\n", "\n", "def get_state(state) -> Literal[\"add_tool_message\", \"info\", \"__end__\"]:\n", - " messages = state['messages']\n", + " messages = state[\"messages\"]\n", " if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:\n", " return \"add_tool_message\"\n", " elif not isinstance(messages[-1], HumanMessage):\n", @@ -190,9 +190,11 @@ "from typing import Annotated\n", "from typing_extensions import TypedDict\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "memory = MemorySaver()\n", "workflow = StateGraph(State)\n", "workflow.add_node(\"info\", info_chain)\n", @@ -201,9 +203,14 @@ "\n", "@workflow.add_node\n", "def add_tool_message(state: State):\n", - " return {\"messages\": [ToolMessage(\n", - " content=\"Prompt generated!\", tool_call_id=state['messages'][-1].tool_calls[0][\"id\"]\n", - " )]}\n", + " return {\n", + " \"messages\": [\n", + " ToolMessage(\n", + " content=\"Prompt generated!\",\n", + " tool_call_id=state[\"messages\"][-1].tool_calls[0][\"id\"],\n", + " )\n", + " ]\n", + " }\n", "\n", "\n", "workflow.add_conditional_edges(\"info\", get_state)\n", @@ -364,7 +371,7 @@ " for output in graph.stream(\n", " {\"messages\": [HumanMessage(content=user)]}, config=config, stream_mode=\"updates\"\n", " ):\n", - " last_message = next(iter(output.values()))['messages'][-1]\n", + " last_message = next(iter(output.values()))[\"messages\"][-1]\n", " last_message.pretty_print()\n", "\n", " if output and \"prompt\" in output:\n", diff --git a/examples/create-react-agent-hitl.ipynb b/examples/create-react-agent-hitl.ipynb index 39392f205..2bc6d306f 100644 --- a/examples/create-react-agent-hitl.ipynb +++ b/examples/create-react-agent-hitl.ipynb @@ -239,7 +239,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.1" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/customer-support/customer-support.ipynb b/examples/customer-support/customer-support.ipynb index cf451188f..07927338c 100644 --- a/examples/customer-support/customer-support.ipynb +++ b/examples/customer-support/customer-support.ipynb @@ -225,7 +225,14 @@ "\n", "Define the (`fetch_user_flight_information`) tool to let the agent see the current user's flight information. Then define tools to search for flights and manage the passenger's bookings stored in the SQL database.\n", "\n", - "We use `ensure_config` to pass in the `passenger_id` in via configurable parameters. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information." + "We the can [access the RunnableConfig](https://python.langchain.com/v0.2/docs/how_to/tool_configure/#inferring-by-parameter-type) for a given run to check the `passenger_id` of the user accessing this application. The LLM never has to provide these explicitly, they are provided for a given invocation of the graph so that each user cannot access other passengers' booking information.\n", + "\n", + "
\n", + "

Compatibility

\n", + "

\n", + " This tutorial expects `langchain-core>=0.2.16` to use the injected RunnableConfig. Prior to that, you'd use `ensure_config` to collect the config from context.\n", + "

\n", + "
\n" ] }, { @@ -240,18 +247,17 @@ "from typing import Optional\n", "\n", "import pytz\n", - "from langchain_core.runnables import ensure_config\n", + "from langchain_core.runnables import RunnableConfig\n", "\n", "\n", "@tool\n", - "def fetch_user_flight_information() -> list[dict]:\n", + "def fetch_user_flight_information(config: RunnableConfig) -> list[dict]:\n", " \"\"\"Fetch all tickets for the user along with corresponding flight information and seat assignments.\n", "\n", " Returns:\n", " A list of dictionaries where each dictionary contains the ticket details,\n", " associated flight details, and the seat assignments for each ticket belonging to the user.\n", " \"\"\"\n", - " config = ensure_config() # Fetch from the context\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -328,9 +334,10 @@ "\n", "\n", "@tool\n", - "def update_ticket_to_new_flight(ticket_no: str, new_flight_id: int) -> str:\n", + "def update_ticket_to_new_flight(\n", + " ticket_no: str, new_flight_id: int, *, config: RunnableConfig\n", + ") -> str:\n", " \"\"\"Update the user's ticket to a new valid flight.\"\"\"\n", - " config = ensure_config()\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -396,9 +403,8 @@ "\n", "\n", "@tool\n", - "def cancel_ticket(ticket_no: str) -> str:\n", + "def cancel_ticket(ticket_no: str, *, config: RunnableConfig) -> str:\n", " \"\"\"Cancel the user's ticket and remove it from the database.\"\"\"\n", - " config = ensure_config()\n", " configuration = config.get(\"configurable\", {})\n", " passenger_id = configuration.get(\"passenger_id\", None)\n", " if not passenger_id:\n", @@ -4407,7 +4413,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.2" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/human_in_the_loop/review-tool-calls.ipynb b/examples/human_in_the_loop/review-tool-calls.ipynb index 5a4aeaf56..84373b90b 100644 --- a/examples/human_in_the_loop/review-tool-calls.ipynb +++ b/examples/human_in_the_loop/review-tool-calls.ipynb @@ -141,15 +141,18 @@ " print(\"----\")\n", " return \"Sunny!\"\n", "\n", - "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools([weather_search])\n", + "\n", + "model = ChatAnthropic(model_name=\"claude-3-5-sonnet-20240620\").bind_tools(\n", + " [weather_search]\n", + ")\n", + "\n", "\n", "class State(MessagesState):\n", " \"\"\"Simple state.\"\"\"\n", "\n", + "\n", "def call_llm(state):\n", - " return {\n", - " \"messages\": [model.invoke(state['messages'])]\n", - " }\n", + " return {\"messages\": [model.invoke(state[\"messages\"])]}\n", "\n", "\n", "def human_review_node(state):\n", @@ -159,28 +162,30 @@ "def run_tool(state):\n", " new_messages = []\n", " tools = {\"weather_search\": weather_search}\n", - " tool_calls = state['messages'][-1].tool_calls\n", + " tool_calls = state[\"messages\"][-1].tool_calls\n", " for tool_call in tool_calls:\n", - " tool = tools[tool_call['name']]\n", - " result = tool.invoke(tool_call['args'])\n", - " new_messages.append({\n", - " \"role\": \"tool\",\n", - " \"name\": tool_call['name'],\n", - " \"content\": result,\n", - " \"tool_call_id\": tool_call['id']\n", - " })\n", + " tool = tools[tool_call[\"name\"]]\n", + " result = tool.invoke(tool_call[\"args\"])\n", + " new_messages.append(\n", + " {\n", + " \"role\": \"tool\",\n", + " \"name\": tool_call[\"name\"],\n", + " \"content\": result,\n", + " \"tool_call_id\": tool_call[\"id\"],\n", + " }\n", + " )\n", " return {\"messages\": new_messages}\n", "\n", "\n", "def route_after_llm(state) -> Literal[END, \"human_review_node\"]:\n", - " if len(state['messages'][-1].tool_calls) == 0:\n", + " if len(state[\"messages\"][-1].tool_calls) == 0:\n", " return END\n", " else:\n", " return \"human_review_node\"\n", "\n", "\n", "def route_after_human(state) -> Literal[\"run_tool\", \"call_llm\"]:\n", - " if isinstance(state['messages'][-1], AIMessage):\n", + " if isinstance(state[\"messages\"][-1], AIMessage):\n", " return \"run_tool\"\n", " else:\n", " return \"call_llm\"\n", @@ -460,35 +465,35 @@ "print(\"Current State:\")\n", "print(state.values)\n", "print(\"\\nCurrent Tool Call ID:\")\n", - "current_content = state.values['messages'][-1].content\n", - "current_id = state.values['messages'][-1].id\n", - "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n", + "current_content = state.values[\"messages\"][-1].content\n", + "current_id = state.values[\"messages\"][-1].id\n", + "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n", "print(tool_call_id)\n", "\n", "# We now need to construct a replacement tool call.\n", "# We will change the argument to be `San Francisco, USA`\n", "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n", "new_message = {\n", - " \"role\": \"assistant\", \n", + " \"role\": \"assistant\",\n", " \"content\": current_content,\n", " \"tool_calls\": [\n", " {\n", " \"id\": tool_call_id,\n", " \"name\": \"weather_search\",\n", - " \"args\": {\"city\": \"San Francisco, USA\"}\n", + " \"args\": {\"city\": \"San Francisco, USA\"},\n", " }\n", " ],\n", " # This is important - this needs to be the same as the message you replacing!\n", " # Otherwise, it will show up as a separate message\n", - " \"id\": current_id\n", + " \"id\": current_id,\n", "}\n", "graph.update_state(\n", " # This is the config which represents this thread\n", - " thread, \n", + " thread,\n", " # This is the updated value we want to push\n", - " {\"messages\": [new_message]}, \n", + " {\"messages\": [new_message]},\n", " # We push this update acting as our human_review_node\n", - " as_node=\"human_review_node\"\n", + " as_node=\"human_review_node\",\n", ")\n", "\n", "# Let's now continue executing from here\n", @@ -595,26 +600,26 @@ "print(\"Current State:\")\n", "print(state.values)\n", "print(\"\\nCurrent Tool Call ID:\")\n", - "tool_call_id = state.values['messages'][-1].tool_calls[0]['id']\n", + "tool_call_id = state.values[\"messages\"][-1].tool_calls[0][\"id\"]\n", "print(tool_call_id)\n", "\n", "# We now need to construct a replacement tool call.\n", "# We will change the argument to be `San Francisco, USA`\n", "# Note that we could change any number of arguments or tool names - it just has to be a valid one\n", "new_message = {\n", - " \"role\": \"tool\", \n", + " \"role\": \"tool\",\n", " # This is our natural language feedback\n", " \"content\": \"User requested changes: pass in the country as well\",\n", " \"name\": \"weather_search\",\n", - " \"tool_call_id\": tool_call_id\n", + " \"tool_call_id\": tool_call_id,\n", "}\n", "graph.update_state(\n", " # This is the config which represents this thread\n", - " thread, \n", + " thread,\n", " # This is the updated value we want to push\n", - " {\"messages\": [new_message]}, \n", + " {\"messages\": [new_message]},\n", " # We push this update acting as our human_review_node\n", - " as_node=\"human_review_node\"\n", + " as_node=\"human_review_node\",\n", ")\n", "\n", "# Let's now continue executing from here\n", diff --git a/examples/input_output_schema.ipynb b/examples/input_output_schema.ipynb index 16779ba6e..4837b40a8 100644 --- a/examples/input_output_schema.ipynb +++ b/examples/input_output_schema.ipynb @@ -33,15 +33,19 @@ "from langgraph.graph import StateGraph, START, END\n", "from typing import TypedDict\n", "\n", + "\n", "class InputState(TypedDict):\n", " question: str\n", "\n", + "\n", "class OutputState(TypedDict):\n", " answer: str\n", "\n", + "\n", "def answer_node(state: InputState):\n", " return {\"answer\": \"bye\"}\n", "\n", + "\n", "graph = StateGraph(input=InputState, output=OutputState)\n", "graph.add_node(answer_node)\n", "graph.add_edge(START, \"answer_node\")\n", diff --git a/examples/llm-compiler/LLMCompiler.ipynb b/examples/llm-compiler/LLMCompiler.ipynb index bfe5736d8..2f2aa81ed 100644 --- a/examples/llm-compiler/LLMCompiler.ipynb +++ b/examples/llm-compiler/LLMCompiler.ipynb @@ -526,7 +526,7 @@ " \"tasks\": tasks,\n", " }\n", " )\n", - " return {\"messages\":[scheduled_tasks]}" + " return {\"messages\": [scheduled_tasks]}" ] }, { @@ -653,7 +653,7 @@ " )\n", " ]\n", " else:\n", - " return {\"messages\":response + [AIMessage(content=decision.action.response)]}\n", + " return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n", "\n", "\n", "def select_recent_messages(state) -> dict:\n", @@ -726,9 +726,11 @@ "from langgraph.graph.message import add_messages\n", "from typing import Annotated\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "graph_builder = StateGraph(State)\n", "\n", "# 1. Define vertices\n", @@ -794,7 +796,9 @@ } ], "source": [ - "for step in chain.stream({\"messages\":[HumanMessage(content=\"What's the GDP of New York?\")]}):\n", + "for step in chain.stream(\n", + " {\"messages\": [HumanMessage(content=\"What's the GDP of New York?\")]}\n", + "):\n", " print(step)\n", " print(\"---\")" ] diff --git a/examples/many-tools.ipynb b/examples/many-tools.ipynb index b415f5065..107c9921f 100644 --- a/examples/many-tools.ipynb +++ b/examples/many-tools.ipynb @@ -328,9 +328,9 @@ " \"set more_information_needed False and populate a blank string for the query.\"\n", " )\n", " input_messages = [system] + state[\"messages\"]\n", - " response = llm.bind_tools(\n", - " [QueryForTools], tool_choice=True\n", - " ).invoke(input_messages)\n", + " response = llm.bind_tools([QueryForTools], tool_choice=True).invoke(\n", + " input_messages\n", + " )\n", " query = response.tool_calls[0][\"args\"][\"query\"]\n", " tool_documents = vector_store.similarity_search(query)\n", " if hack_remove_tool_condition:\n", diff --git a/examples/pass-run-time-values-to-tools.ipynb b/examples/pass-run-time-values-to-tools.ipynb index 3c74bd06d..ab6680230 100644 --- a/examples/pass-run-time-values-to-tools.ipynb +++ b/examples/pass-run-time-values-to-tools.ipynb @@ -329,6 +329,7 @@ "\n", "tools = [get_context, cite_context_sources]\n", "\n", + "\n", "# Define the function that calls the model\n", "def call_model(state, config):\n", " messages = state[\"messages\"]\n", diff --git a/examples/pass_private_state.ipynb b/examples/pass_private_state.ipynb index 2e60d805e..d34ecc7e1 100644 --- a/examples/pass_private_state.ipynb +++ b/examples/pass_private_state.ipynb @@ -72,12 +72,12 @@ "# Node to retrieve documents\n", "def retrieve_documents(state: QueryOutputState) -> DocumentOutputState:\n", " # Replace this with real logic\n", - " return {\"docs\": [state['query']] * 2}\n", + " return {\"docs\": [state[\"query\"]] * 2}\n", "\n", "\n", "# Node to generate answer\n", "def generate(state: GenerateInputState) -> OverallState:\n", - " return {\"answer\": \"\\n\\n\".join(state['docs'] + [state['question']])}\n", + " return {\"answer\": \"\\n\\n\".join(state[\"docs\"] + [state[\"question\"]])}\n", "\n", "\n", "graph = StateGraph(OverallState)\n", diff --git a/examples/persistence.ipynb b/examples/persistence.ipynb index 869043cf1..715f7c0b5 100644 --- a/examples/persistence.ipynb +++ b/examples/persistence.ipynb @@ -587,7 +587,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.9" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/examples/persistence_mongodb.ipynb b/examples/persistence_mongodb.ipynb index 99ce73057..54dfdc638 100644 --- a/examples/persistence_mongodb.ipynb +++ b/examples/persistence_mongodb.ipynb @@ -630,7 +630,7 @@ " upsert=True,\n", " )\n", " )\n", - " await self.db[\"checkpoint_writes\"].bulk_write(operations)\n" + " await self.db[\"checkpoint_writes\"].bulk_write(operations)" ] }, { @@ -685,7 +685,9 @@ "metadata": {}, "outputs": [], "source": [ - "with MongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + "with MongoDBSaver.from_conn_info(\n", + " host=\"localhost\", port=27017, db_name=\"checkpoints\"\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"1\"}}\n", " res = graph.invoke({\"messages\": [(\"human\", \"what's the weather in sf\")]}, config)\n", @@ -796,10 +798,14 @@ "metadata": {}, "outputs": [], "source": [ - "async with AsyncMongoDBSaver.from_conn_info(host=\"localhost\", port=27017, db_name=\"checkpoints\") as checkpointer:\n", + "async with AsyncMongoDBSaver.from_conn_info(\n", + " host=\"localhost\", port=27017, db_name=\"checkpoints\"\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"2\"}}\n", - " res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n", + " res = await graph.ainvoke(\n", + " {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n", + " )\n", "\n", " latest_checkpoint = await checkpointer.aget(config)\n", " latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n", diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index f43fa01d6..5ea838cee 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -134,7 +134,7 @@ "source": [ "from psycopg.rows import dict_row\n", "\n", - "connection_kwargs ={\n", + "connection_kwargs = {\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", "}" @@ -165,7 +165,7 @@ " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", - " kwargs=connection_kwargs\n", + " kwargs=connection_kwargs,\n", ")\n", "\n", "with pool.connection() as conn:\n", @@ -393,7 +393,7 @@ " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", - " kwargs=connection_kwargs\n", + " kwargs=connection_kwargs,\n", ") as pool, pool.connection() as conn:\n", " checkpointer = AsyncPostgresSaver(conn)\n", "\n", diff --git a/examples/persistence_redis.ipynb b/examples/persistence_redis.ipynb index 3b2ad1170..34d78a53f 100644 --- a/examples/persistence_redis.ipynb +++ b/examples/persistence_redis.ipynb @@ -530,7 +530,9 @@ "\n", " @classmethod\n", " @asynccontextmanager\n", - " async def from_conn_info(cls, *, host: str, port: int, db: int) -> AsyncIterator[\"AsyncRedisSaver\"]:\n", + " async def from_conn_info(\n", + " cls, *, host: str, port: int, db: int\n", + " ) -> AsyncIterator[\"AsyncRedisSaver\"]:\n", " conn = None\n", " try:\n", " conn = AsyncRedis(host=host, port=port, db=db)\n", @@ -887,10 +889,14 @@ "metadata": {}, "outputs": [], "source": [ - "async with AsyncRedisSaver.from_conn_info(host=\"localhost\", port=6379, db=0) as checkpointer:\n", + "async with AsyncRedisSaver.from_conn_info(\n", + " host=\"localhost\", port=6379, db=0\n", + ") as checkpointer:\n", " graph = create_react_agent(model, tools=tools, checkpointer=checkpointer)\n", " config = {\"configurable\": {\"thread_id\": \"2\"}}\n", - " res = await graph.ainvoke({\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config)\n", + " res = await graph.ainvoke(\n", + " {\"messages\": [(\"human\", \"what's the weather in nyc\")]}, config\n", + " )\n", "\n", " latest_checkpoint = await checkpointer.aget(config)\n", " latest_checkpoint_tuple = await checkpointer.aget_tuple(config)\n", diff --git a/examples/reflection/reflection.ipynb b/examples/reflection/reflection.ipynb index 810990e38..ca42160bd 100644 --- a/examples/reflection/reflection.ipynb +++ b/examples/reflection/reflection.ipynb @@ -269,7 +269,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", - " \n", + "\n", "async def generation_node(state: Sequence[BaseMessage]):\n", " return await generate.ainvoke({\"messages\": state})\n", "\n", diff --git a/examples/reflexion/reflexion.ipynb b/examples/reflexion/reflexion.ipynb index 670e6eb5b..8c183c817 100644 --- a/examples/reflexion/reflexion.ipynb +++ b/examples/reflexion/reflexion.ipynb @@ -392,6 +392,7 @@ "class State(TypedDict):\n", " messages: Annotated[list, add_messages]\n", "\n", + "\n", "MAX_ITERATIONS = 5\n", "builder = StateGraph(State)\n", "builder.add_node(\"draft\", first_responder.respond)\n", diff --git a/examples/streaming-content.ipynb b/examples/streaming-content.ipynb index ff182ca21..8c2c7ab36 100644 --- a/examples/streaming-content.ipynb +++ b/examples/streaming-content.ipynb @@ -68,7 +68,9 @@ " # It's completely optional, but useful if you have many functions with similar names\n", " gen = RunnableGenerator(my_generator).with_config(\n", " tags=[\"should_stream\"],\n", - " callbacks=config.get(\"callbacks\", []) # <-- Propagate callbacks (Python <= 3.10)\n", + " callbacks=config.get(\n", + " \"callbacks\", []\n", + " ), # <-- Propagate callbacks (Python <= 3.10)\n", " )\n", " async for message in gen.astream(state):\n", " messages.append(message)\n", diff --git a/examples/tutorials/rag-agent-testing-local.ipynb b/examples/tutorials/rag-agent-testing-local.ipynb index f9e89e56d..3105d3342 100644 --- a/examples/tutorials/rag-agent-testing-local.ipynb +++ b/examples/tutorials/rag-agent-testing-local.ipynb @@ -169,9 +169,7 @@ "from langchain_core.output_parsers import JsonOutputParser\n", "\n", "# JSON\n", - "llm = ChatOllama(model=\"llama3.1\", \n", - " format=\"json\", \n", - " temperature=0)\n", + "llm = ChatOllama(model=\"llama3.1\", format=\"json\", temperature=0)\n", "\n", "\n", "prompt = PromptTemplate(\n", @@ -210,6 +208,7 @@ "from IPython.display import Image, display\n", "from langgraph.graph import START, END, StateGraph\n", "\n", + "\n", "class GraphState(TypedDict):\n", " \"\"\"\n", " Represents the state of our graph.\n", @@ -381,21 +380,22 @@ "metadata": {}, "outputs": [], "source": [ - "import uuid \n", + "import uuid\n", + "\n", "\n", "def predict_custom_agent_answer(example: dict):\n", - " \n", " config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n", - " \n", + "\n", " state_dict = custom_graph.invoke(\n", " {\"question\": example[\"input\"], \"steps\": []}, config\n", " )\n", - " \n", + "\n", " return {\"response\": state_dict[\"generation\"], \"steps\": state_dict[\"steps\"]}\n", "\n", + "\n", "example = {\"input\": \"What are the types of agent memory?\"}\n", - "#response = predict_custom_agent_answer(example)\n", - "#response" + "# response = predict_custom_agent_answer(example)\n", + "# response" ] }, { @@ -544,6 +544,7 @@ " \"generate_answer\",\n", "]\n", "\n", + "\n", "def check_trajectory_custom(root_run: Run, example: Example) -> dict:\n", " \"\"\"\n", " Check if all expected tools are called in exact order and without any additional tool calls.\n", diff --git a/examples/tutorials/tool-calling-agent-local.ipynb b/examples/tutorials/tool-calling-agent-local.ipynb index 5c60336e8..12c43038f 100644 --- a/examples/tutorials/tool-calling-agent-local.ipynb +++ b/examples/tutorials/tool-calling-agent-local.ipynb @@ -134,6 +134,7 @@ " for d in web_results\n", " ]\n", "\n", + "\n", "# Tool list\n", "tools = [retrieve_documents, web_search]" ] @@ -152,9 +153,11 @@ "from langgraph.graph.message import AnyMessage, add_messages\n", "from typing_extensions import TypedDict\n", "\n", + "\n", "class State(TypedDict):\n", " messages: Annotated[list[AnyMessage], add_messages]\n", "\n", + "\n", "class Assistant:\n", " def __init__(self, runnable: Runnable):\n", " \"\"\"\n", @@ -291,6 +294,7 @@ "source": [ "import uuid\n", "\n", + "\n", "def predict_react_agent_answer(example: dict):\n", " \"\"\"Use this for answer evaluation\"\"\"\n", "\n", From 82db3831995ec3bf1897e5f79a97826b8cf03cb7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 22 Aug 2024 21:36:27 -0700 Subject: [PATCH 16/16] Don't try to serialize exceptions - Store their string repr instead --- .../langgraph/checkpoint/serde/jsonplus.py | 5 ++- libs/langgraph/tests/any_str.py | 18 -------- libs/langgraph/tests/memory_assert.py | 2 - libs/langgraph/tests/test_pregel.py | 10 ++--- libs/langgraph/tests/test_pregel_async.py | 42 +++++++++---------- 5 files changed, 30 insertions(+), 47 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 5465b2802..608a077cd 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -7,6 +7,7 @@ import re from collections import deque from datetime import date, datetime, time, timedelta, timezone from enum import Enum +from inspect import isclass from ipaddress import ( IPv4Address, IPv4Interface, @@ -111,7 +112,7 @@ class JsonPlusSerializer(SerializerProtocol): obj.__class__, method="fromhex", args=[obj.hex()] ) elif isinstance(obj, BaseException): - return self._encode_constructor_args(obj.__class__, args=obj.args) + return repr(obj) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" @@ -135,6 +136,8 @@ class JsonPlusSerializer(SerializerProtocol): method = getattr(cls, value["method"]) else: method = cls + if isclass(method) and issubclass(method, BaseException): + return None if value["args"] and value["kwargs"]: return method(*value["args"], **value["kwargs"]) elif value["args"]: diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index a98962cdc..836cf9371 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -23,24 +23,6 @@ class AnyVersion: return hash(str(self)) -class ExceptionLike: - def __init__(self, exc: Exception) -> None: - self.exc = exc - - def __eq__(self, value: object) -> bool: - return ( - isinstance(value, Exception) - and self.exc.__class__ == value.__class__ - and str(self.exc) == str(value) - ) - - def __hash__(self) -> int: - return hash((self.exc.__class__, str(self.exc))) - - def __repr__(self) -> str: - return str(self.exc) - - class UnsortedSequence: def __init__(self, *values: Any) -> None: self.seq = values diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 0b0bcf62f..624a711c7 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -24,8 +24,6 @@ class NoopSerializer(SerializerProtocol): class MemorySaverAssertImmutable(MemorySaver): - serde = NoopSerializer() - storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]] def __init__( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 06a747a17..7325aa41d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -64,7 +64,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence +from tests.any_str import AnyStr, AnyVersion, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -1307,7 +1307,7 @@ def test_invoke_checkpoint_two( assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 assert checkpoint_tup.pending_writes == [ - (AnyStr(), ERROR, ExceptionLike(ValueError("Input is too large"))) + (AnyStr(), ERROR, "ValueError('Input is too large')") ] # on a new thread, total starts out as 0, so output is 0+5=5 assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 @@ -1374,7 +1374,7 @@ def test_pending_writes_resume( assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", ExceptionLike(ConnectionError("I'm not good"))), + PregelTask(AnyStr(), "two", 'ConnectionError("I\'m not good")'), ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" @@ -1384,7 +1384,7 @@ def test_pending_writes_resume( expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), ERROR, ExceptionLike(ConnectionError("I'm not good"))), + (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), ] assert len(checkpoint.pending_writes) == 3 assert all(w in expected_writes for w in checkpoint.pending_writes) @@ -1518,7 +1518,7 @@ def test_pending_writes_resume( pending_writes=UnsortedSequence( (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4eed26017..a35bf7b00 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -61,7 +61,7 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore -from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence +from tests.any_str import AnyStr, AnyVersion, UnsortedSequence from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -1593,7 +1593,7 @@ async def test_pending_writes_resume( assert state.next == ("one", "two") assert state.tasks == ( PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", ExceptionLike(ValueError("I'm not good"))), + PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" @@ -1603,7 +1603,7 @@ async def test_pending_writes_resume( expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), ERROR, ExceptionLike(ValueError("I'm not good"))), + (AnyStr(), ERROR, 'ValueError("I\'m not good")'), ] assert len(checkpoint.pending_writes) == 3 assert all(w in expected_writes for w in checkpoint.pending_writes) @@ -1741,7 +1741,7 @@ async def test_pending_writes_resume( pending_writes=UnsortedSequence( (AnyStr(), "one", "one"), (AnyStr(), "value", 2), - (AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))), + (AnyStr(), "__error__", 'ValueError("I\'m not good")'), (AnyStr(), "two", "two"), (AnyStr(), "value", 3), ), @@ -3108,7 +3108,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: 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: @@ -3140,22 +3140,22 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> 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" @@ -3176,14 +3176,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), "result for query", - ) + ] ], } }, @@ -3199,14 +3199,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="another", log="tool:search_api:another", ), "result for another", - ), + ], ], } }, @@ -3361,14 +3361,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3401,14 +3401,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: 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=(), @@ -3537,14 +3537,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: { "tools": { "intermediate_steps": [ - ( + [ AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), "result for query", - ) + ] ], } }, @@ -3576,14 +3576,14 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: 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=(),