From 3ff3def62b4c655906b77c048cd50dfb8b73caee Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 3 May 2024 15:55:59 -0700 Subject: [PATCH] Remove option to only checkpoint at end of run - Now that checkpoint at end of each step adds no latency there is not point to keep this - This will make it easier to add future features --- langgraph/checkpoint/__init__.py | 2 - langgraph/checkpoint/aiosqlite.py | 4 +- langgraph/checkpoint/base.py | 14 - langgraph/checkpoint/memory.py | 4 +- langgraph/checkpoint/sqlite.py | 4 +- langgraph/graph/__init__.py | 4 +- langgraph/pregel/__init__.py | 62 +-- tests/__snapshots__/test_pregel.ambr | 518 +-------------------- tests/__snapshots__/test_pregel_async.ambr | 78 +--- tests/memory_assert.py | 6 +- tests/test_pregel.py | 468 ++++++------------- tests/test_pregel_async.py | 466 ++++++------------ 12 files changed, 303 insertions(+), 1327 deletions(-) diff --git a/langgraph/checkpoint/__init__.py b/langgraph/checkpoint/__init__.py index 5a0f4bd0c..50f9db11b 100644 --- a/langgraph/checkpoint/__init__.py +++ b/langgraph/checkpoint/__init__.py @@ -1,7 +1,6 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, SerializerProtocol, ) from langgraph.checkpoint.memory import MemorySaver @@ -9,7 +8,6 @@ from langgraph.checkpoint.memory import MemorySaver __all__ = [ "BaseCheckpointSaver", "Checkpoint", - "CheckpointAt", "MemorySaver", "SerializerProtocol", ] diff --git a/langgraph/checkpoint/aiosqlite.py b/langgraph/checkpoint/aiosqlite.py index e0758d84d..314448b01 100644 --- a/langgraph/checkpoint/aiosqlite.py +++ b/langgraph/checkpoint/aiosqlite.py @@ -10,7 +10,6 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -80,9 +79,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): conn: aiosqlite.Connection, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ): - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.conn = conn self.lock = asyncio.Lock() self.is_setup = False diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 6bfc79885..70734f1e6 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -14,7 +14,6 @@ from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer -from langgraph.utils import StrEnum class Checkpoint(TypedDict): @@ -71,15 +70,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: ) -class CheckpointAt(StrEnum): - """When to take a checkpoint.""" - - END_OF_STEP = "end_of_step" - """Take a checkpoint at the end of each step.""" - END_OF_RUN = "end_of_run" - """Take a checkpoint at the end of the run.""" - - class CheckpointTuple(NamedTuple): config: RunnableConfig checkpoint: Checkpoint @@ -107,18 +97,14 @@ CheckpointThreadTs = ConfigurableFieldSpec( class BaseCheckpointSaver(ABC): - at: CheckpointAt = CheckpointAt.END_OF_STEP - serde: SerializerProtocol = JsonPlusSerializer() def __init__( self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: self.serde = serde or self.serde - self.at = at or self.at @property def config_specs(self) -> list[ConfigurableFieldSpec]: diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index dab21794f..81eba2438 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -7,7 +7,6 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -45,9 +44,8 @@ class MemorySaver(BaseCheckpointSaver): self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.storage = defaultdict(dict) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index f45c28be4..3c2b63ca3 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -10,7 +10,6 @@ from typing_extensions import Self from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, CheckpointTuple, SerializerProtocol, ) @@ -90,9 +89,8 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): conn: sqlite3.Connection, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.conn = conn self.is_setup = False diff --git a/langgraph/graph/__init__.py b/langgraph/graph/__init__.py index 8fac44cfa..1c0ef3262 100644 --- a/langgraph/graph/__init__.py +++ b/langgraph/graph/__init__.py @@ -1,5 +1,5 @@ from langgraph.graph.graph import END, Graph -from langgraph.graph.message import MessageGraph +from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph -__all__ = ["END", "Graph", "StateGraph", "MessageGraph"] +__all__ = ["END", "Graph", "StateGraph", "MessageGraph", "add_messages"] diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index c29f55b86..ffa14304c 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -58,7 +58,6 @@ from langgraph.channels.base import ( from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, - CheckpointAt, copy_checkpoint, empty_checkpoint, ) @@ -769,9 +768,7 @@ class Pregel( yield from map_output_updates(output_keys, next_tasks) # save end of step checkpoint - if self.checkpointer is not None and ( - self.checkpointer.at == CheckpointAt.END_OF_STEP - ): + if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = self.checkpointer.put( checkpoint_config, checkpoint @@ -799,33 +796,6 @@ class Pregel( # set final channel values as run output run_manager.on_chain_end(read_channels(channels, output_keys)) - - # save end of run checkpoint - if ( - self.checkpointer is not None - and self.checkpointer.at == CheckpointAt.END_OF_RUN - ): - checkpoint = create_checkpoint(checkpoint, channels) - executor.submit( - self.checkpointer.put(checkpoint_config, checkpoint) - ) - checkpoint_config = { - "configurable": { - "thread_id": checkpoint_config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], - } - } - if stream_mode == "debug": - yield map_debug_checkpoint( - step, - checkpoint_config, - channels, - self.stream_channels_asis, - ) - elif self.checkpointer is None and stream_mode == "debug": - yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis - ) except BaseException as e: run_manager.on_chain_error(e) raise @@ -1035,9 +1005,7 @@ class Pregel( yield chunk # save end of step checkpoint - if self.checkpointer is not None and ( - self.checkpointer.at == CheckpointAt.END_OF_STEP - ): + if self.checkpointer is not None: checkpoint = create_checkpoint(checkpoint, channels) checkpoint_config = await self.checkpointer.aput( checkpoint_config, checkpoint @@ -1065,32 +1033,6 @@ class Pregel( # set final channel values as run output await run_manager.on_chain_end(read_channels(channels, output_keys)) - - # save end of run checkpoint - if ( - self.checkpointer is not None - and self.checkpointer.at == CheckpointAt.END_OF_RUN - ): - checkpoint = create_checkpoint(checkpoint, channels) - tasks.append( - asyncio.create_task( - self.checkpointer.aput(checkpoint_config, checkpoint) - ) - ) - checkpoint_config = { - "configurable": { - "thread_id": checkpoint_config["configurable"]["thread_id"], - "thread_ts": checkpoint["ts"], - } - } - if stream_mode == "debug": - yield map_debug_checkpoint( - step, checkpoint_config, channels, self.stream_channels_asis - ) - elif self.checkpointer is None and stream_mode == "debug": - yield map_debug_checkpoint( - step, None, channels, self.stream_channels_asis - ) except BaseException as e: await run_manager.on_chain_error(e) raise diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 1b1866624..1cf4bd12e 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1,5 +1,5 @@ # serializer version: 1 -# name: test_branch_then[end_of_run] +# name: test_branch_then ''' graph TD; __start__ --> prepare; @@ -11,41 +11,7 @@ ''' # --- -# name: test_branch_then[end_of_run].1 - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - prepare([prepare]):::otherclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - finish([finish]):::otherclass; - __start__ --> prepare; - finish --> __end__; - prepare -.-> tool_two_slow; - tool_two_slow --> finish; - prepare -.-> tool_two_fast; - tool_two_fast --> finish; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_branch_then[end_of_step] - ''' - 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[end_of_step].1 +# name: test_branch_then.1 ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; @@ -229,7 +195,7 @@ ''' # --- -# name: test_conditional_graph[end_of_run] +# name: test_conditional_graph ''' { "nodes": [ @@ -294,7 +260,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_run].1 +# name: test_conditional_graph.1 ''' graph TD; __start__ --> agent; @@ -304,7 +270,7 @@ ''' # --- -# name: test_conditional_graph[end_of_run].2 +# name: test_conditional_graph.2 ''' { "nodes": [ @@ -442,7 +408,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_run].3 +# name: test_conditional_graph.3 ''' graph TD; PromptTemplate --> FakeStreamingListLLM; @@ -458,242 +424,13 @@ ''' # --- -# name: test_conditional_graph[end_of_step] - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnableAssign" - ], - "name": "RunnableAssign" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "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[end_of_step].1 - ''' - graph TD; - __start__ --> agent; - tools --> agent; - agent -. continue .-> tools; - agent -. exit .-> __end__; - - ''' -# --- -# name: test_conditional_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": 2, - "type": "schema", - "data": "ParallelInput" - }, - { - "id": 3, - "type": "schema", - "data": "ParallelOutput" - }, - { - "id": 4, - "type": "runnable", - "data": { - "id": [ - "langchain", - "prompts", - "prompt", - "PromptTemplate" - ], - "name": "PromptTemplate" - } - }, - { - "id": 5, - "type": "runnable", - "data": { - "id": [ - "langchain_community", - "llms", - "fake", - "FakeStreamingListLLM" - ], - "name": "FakeStreamingListLLM" - } - }, - { - "id": 6, - "type": "runnable", - "data": { - "id": [ - "langchain_core", - "runnables", - "base", - "RunnableLambda" - ], - "name": "agent_parser" - } - }, - { - "id": 7, - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnablePassthrough" - ], - "name": "RunnablePassthrough" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": 4, - "target": 5 - }, - { - "source": 5, - "target": 6 - }, - { - "source": 2, - "target": 4 - }, - { - "source": 6, - "target": 3 - }, - { - "source": 2, - "target": 7 - }, - { - "source": 7, - "target": 3 - }, - { - "source": "__start__", - "target": 2 - }, - { - "source": "tools", - "target": 2 - }, - { - "source": 3, - "target": "tools", - "data": "continue", - "conditional": true - }, - { - "source": 3, - "target": "__end__", - "data": "exit", - "conditional": true - } - ] - } - ''' -# --- -# name: test_conditional_graph[end_of_step].3 - ''' - graph TD; - PromptTemplate --> FakeStreamingListLLM; - FakeStreamingListLLM --> Lambda_agent_parser_; - Parallel_agent_outcome_Input --> PromptTemplate; - Lambda_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_state_graph[end_of_run] +# name: test_conditional_state_graph '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' # --- -# name: test_conditional_state_graph[end_of_run].1 +# name: test_conditional_state_graph.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": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' # --- -# name: test_conditional_state_graph[end_of_run].2 +# name: test_conditional_state_graph.2 ''' { "nodes": [ @@ -758,7 +495,7 @@ } ''' # --- -# name: test_conditional_state_graph[end_of_run].3 +# name: test_conditional_state_graph.3 ''' graph TD; __start__ --> agent; @@ -768,88 +505,7 @@ ''' # --- -# name: test_conditional_state_graph[end_of_step] - '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}' -# --- -# name: test_conditional_state_graph[end_of_step].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": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}' -# --- -# name: test_conditional_state_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "langchain", - "schema", - "runnable", - "RunnableSequence" - ], - "name": "RunnableSequence" - } - }, - { - "id": "tools", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "utils", - "RunnableCallable" - ], - "name": "tools" - } - } - ], - "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[end_of_step].3 - ''' - graph TD; - __start__ --> agent; - tools --> agent; - agent -. continue .-> tools; - agent -. exit .-> __end__; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_run] +# name: test_in_one_fan_out_state_graph_waiting_edge ''' graph TD; __start__ --> rewrite_query; @@ -862,20 +518,7 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge[end_of_step] - ''' - 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[end_of_run] +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class ''' graph TD; __start__ --> rewrite_query; @@ -888,7 +531,7 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_step] +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' graph TD; __start__ --> rewrite_query; @@ -901,39 +544,13 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run] - ''' - 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[end_of_step] - ''' - 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[end_of_run] +# 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"}}, "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"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "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"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "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, usually passed in as the first of a sequence\\nof input 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 function back to a model.", "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.", "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"}}, "required": ["content", "tool_call_id"]}}}' # --- -# name: test_message_graph[end_of_run].1 +# name: test_message_graph.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"}}, "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"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "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"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "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, usually passed in as the first of a sequence\\nof input 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 function back to a model.", "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.", "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"}}, "required": ["content", "tool_call_id"]}}}' # --- -# name: test_message_graph[end_of_run].2 +# name: test_message_graph.2 ''' { "nodes": [ @@ -998,88 +615,7 @@ } ''' # --- -# name: test_message_graph[end_of_run].3 - ''' - graph TD; - __start__ --> agent; - action --> agent; - agent -. continue .-> action; - agent -. end .-> __end__; - - ''' -# --- -# name: test_message_graph[end_of_step] - '{"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"}}, "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"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "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"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "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, usually passed in as the first of a sequence\\nof input 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 function back to a model.", "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.", "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"}}, "required": ["content", "tool_call_id"]}}}' -# --- -# name: test_message_graph[end_of_step].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"}}, "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"}}, "required": ["name", "args", "id", "error"]}, "AIMessage": {"title": "AIMessage", "description": "Message from an AI.", "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"}}}, "required": ["content"]}, "HumanMessage": {"title": "HumanMessage", "description": "Message from a human.", "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, usually passed in as the first of a sequence\\nof input 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 function back to a model.", "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.", "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"}}, "required": ["content", "tool_call_id"]}}}' -# --- -# name: test_message_graph[end_of_step].2 - ''' - { - "nodes": [ - { - "id": "__start__", - "type": "schema", - "data": "__start__" - }, - { - "id": "__end__", - "type": "schema", - "data": "__end__" - }, - { - "id": "agent", - "type": "runnable", - "data": { - "id": [ - "tests", - "test_pregel", - "FakeFuntionChatModel" - ], - "name": "FakeFuntionChatModel" - } - }, - { - "id": "action", - "type": "runnable", - "data": { - "id": [ - "langgraph", - "prebuilt", - "tool_node", - "ToolNode" - ], - "name": "tools" - } - } - ], - "edges": [ - { - "source": "__start__", - "target": "agent" - }, - { - "source": "action", - "target": "agent" - }, - { - "source": "agent", - "target": "action", - "data": "continue", - "conditional": true - }, - { - "source": "agent", - "target": "__end__", - "data": "end", - "conditional": true - } - ] - } - ''' -# --- -# name: test_message_graph[end_of_step].3 +# name: test_message_graph.3 ''' graph TD; __start__ --> agent; @@ -1456,25 +992,7 @@ ''' # --- -# name: test_start_branch_then[end_of_run] - ''' - %%{init: {'flowchart': {'curve': 'linear'}}}%% - graph TD; - __start__[__start__]:::startclass; - __end__[__end__]:::endclass; - tool_two_slow([tool_two_slow]):::otherclass; - tool_two_fast([tool_two_fast]):::otherclass; - __start__ -.-> tool_two_slow; - tool_two_slow --> __end__; - __start__ -.-> tool_two_fast; - tool_two_fast --> __end__; - classDef startclass fill:#ffdfba; - classDef endclass fill:#baffc9; - classDef otherclass fill:#fad7de; - - ''' -# --- -# name: test_start_branch_then[end_of_step] +# name: test_start_branch_then ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% graph TD; diff --git a/tests/__snapshots__/test_pregel_async.ambr b/tests/__snapshots__/test_pregel_async.ambr index 2d68a82a5..809ff8731 100644 --- a/tests/__snapshots__/test_pregel_async.ambr +++ b/tests/__snapshots__/test_pregel_async.ambr @@ -1,5 +1,5 @@ # serializer version: 1 -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_run] +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class ''' +-----------+ | __start__ | @@ -36,81 +36,7 @@ +---------+ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[end_of_step] - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** ... - * . - ** ... - +--------------+ . - | analyzer_one | . - +--------------+ . - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_run] - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** ... - * . - ** ... - +--------------+ . - | analyzer_one | . - +--------------+ . - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[end_of_step] +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' +-----------+ | __start__ | diff --git a/tests/memory_assert.py b/tests/memory_assert.py index 6bceb865e..625b6f12e 100644 --- a/tests/memory_assert.py +++ b/tests/memory_assert.py @@ -3,7 +3,6 @@ from typing import Any, Optional from langgraph.checkpoint.base import ( Checkpoint, - CheckpointAt, SerializerProtocol, copy_checkpoint, ) @@ -21,17 +20,14 @@ class NoopSerializer(SerializerProtocol): class MemorySaverAssertImmutable(MemorySaver): serde = NoopSerializer() - at = CheckpointAt.END_OF_STEP - storage_for_copies: defaultdict[str, dict[str, Checkpoint]] def __init__( self, *, serde: Optional[SerializerProtocol] = None, - at: Optional[CheckpointAt] = None, ) -> None: - super().__init__(serde=serde, at=at) + super().__init__(serde=serde) self.storage_for_copies = defaultdict(dict) def put(self, config: dict, checkpoint: Checkpoint) -> None: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 2ba1d1d77..9180e71c9 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -16,7 +16,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic -from langgraph.checkpoint.base import CheckpointAt from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, Graph from langgraph.graph.message import MessageGraph @@ -292,17 +291,12 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_two_processes_in_out_interrupt( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one, "two": two}, channels={ @@ -475,12 +469,6 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "step": 1, "payload": {"config": None, "values": {"output": 4, "inbox": []}}, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -627,10 +615,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) -> None: +def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -645,7 +630,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) - | raise_if_above_10 ) - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -686,12 +671,7 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) - assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_invoke_checkpoint_sqlite( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +def test_invoke_checkpoint_sqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -707,7 +687,6 @@ def test_invoke_checkpoint_sqlite( ) with SqliteSaver.from_conn_string(":memory:") as memory: - memory.at = checkpoint_at app = Pregel( nodes={"one": one}, channels={ @@ -992,12 +971,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_conditional_graph( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_conditional_graph(snapshot: SnapshotAssertion) -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM @@ -1199,7 +1173,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1352,7 +1326,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1500,7 +1474,7 @@ def test_conditional_graph( # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1668,12 +1642,7 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_conditional_state_graph( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish @@ -1840,7 +1809,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1958,7 +1927,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], debug=True, ) @@ -2077,7 +2046,7 @@ def test_conditional_state_graph( # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before="*", debug=True, ) @@ -2174,7 +2143,7 @@ def test_conditional_state_graph( # test w interrupt after all app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after="*", ) config = {"configurable": {"thread_id": "4"}} @@ -2796,12 +2765,8 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_message_graph( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, deterministic_uuids: MockerFixture, ) -> None: from copy import deepcopy @@ -3020,7 +2985,7 @@ def test_message_graph( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -3192,7 +3157,7 @@ def test_message_graph( ) app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"], ) config = {"configurable": {"thread_id": "2"}} @@ -3472,12 +3437,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_start_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_start_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -3511,7 +3471,6 @@ def test_start_branch_then( } with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] ) @@ -3569,10 +3528,7 @@ def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) -> None: +def test_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -3619,254 +3575,125 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - } with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at - # test stream_mode=debug tool_two = tool_two_graph.compile(checkpointer=saver) thread10 = {"configurable": {"thread_id": "10"}} - if checkpoint_at is CheckpointAt.END_OF_RUN: - assert [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"my_key": "value", "market": "DE"}, + assert [ + *tool_two.stream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", }, }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 4, - "payload": { - "config": { - "configurable": { - "thread_id": "10", - "thread_ts": AnyStr(), - } - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] - else: - assert [ - *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] + }, + ] tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] @@ -3925,7 +3752,6 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - ) with SqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] ) @@ -3983,12 +3809,7 @@ def test_branch_then(snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt) - ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_in_one_fan_out_state_graph_waiting_edge( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -4054,7 +3875,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4075,12 +3896,8 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4150,7 +3967,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4171,12 +3988,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -4254,7 +4067,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -4275,12 +4088,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpoint_at: CheckpointAt, -) -> None: +def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -4349,7 +4157,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 79e4b3e33..32b19b279 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -25,7 +25,6 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver -from langgraph.checkpoint.base import CheckpointAt from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph from langgraph.prebuilt.chat_agent_executor import ( @@ -270,17 +269,12 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_two_processes_in_out_interrupt( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one, "two": two}, channels={ @@ -457,12 +451,6 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "step": 1, "payload": {"config": None, "values": {"output": 4, "inbox": []}}, }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": {"config": None, "values": {"output": 4, "inbox": []}}, - }, ] @@ -613,12 +601,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_checkpoint( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -633,7 +616,7 @@ async def test_invoke_checkpoint( | raise_if_above_10 ) - memory = MemorySaverAssertImmutable(at=checkpoint_at) + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -674,12 +657,7 @@ async def test_invoke_checkpoint( assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_invoke_checkpoint_aiosqlite( - mocker: MockerFixture, checkpoint_at: CheckpointAt -) -> None: +async def test_invoke_checkpoint_aiosqlite(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -695,7 +673,6 @@ async def test_invoke_checkpoint_aiosqlite( ) async with AsyncSqliteSaver.from_conn_string(":memory:") as memory: - memory.at = checkpoint_at app = Pregel( nodes={"one": one}, channels={ @@ -1003,10 +980,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup_async.call_count == 1, "Expected cleanup to be called once" -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: +async def test_conditional_graph() -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM @@ -1274,7 +1248,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1424,7 +1398,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1575,7 +1549,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1702,10 +1676,7 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: +async def test_conditional_graph_state() -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish @@ -1899,7 +1870,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2022,7 +1993,7 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None: # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -2537,10 +2508,7 @@ async def test_prebuilt_chat() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_message_graph(checkpoint_at: CheckpointAt) -> None: +async def test_message_graph() -> None: from langchain.chat_models.fake import FakeMessagesListChatModel from langchain_community.tools import tool from langchain_core.agents import AgentAction @@ -2709,7 +2677,7 @@ async def test_message_graph(checkpoint_at: CheckpointAt) -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2938,12 +2906,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_start_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +async def test_start_branch_then(snapshot: SnapshotAssertion) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -2966,7 +2929,6 @@ async def test_start_branch_then( } async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] ) @@ -3024,12 +2986,7 @@ async def test_start_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_branch_then( - snapshot: SnapshotAssertion, checkpoint_at: CheckpointAt -) -> None: +async def test_branch_then() -> None: pass class State(TypedDict): @@ -3060,256 +3017,126 @@ async def test_branch_then( } async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at - # test stream_mode=debug tool_two = tool_two_graph.compile(checkpointer=saver) thread10 = {"configurable": {"thread_id": "10"}} - if checkpoint_at is CheckpointAt.END_OF_RUN: - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": None, - "values": {"my_key": "value", "market": "DE"}, + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", + "name": "prepare", + "result": [("my_key", " prepared")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", + "name": "tool_two_slow", + "result": [("my_key", " slow")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": {"my_key": "value prepared slow", "market": "DE"}, + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ["branch:prepare:condition:then"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", + "name": "finish", + "result": [("my_key", " finished")], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "configurable": {"thread_id": "10", "thread_ts": AnyStr()} + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", }, }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": None, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": None, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": None, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 4, - "payload": { - "config": { - "configurable": { - "thread_id": "10", - "thread_ts": AnyStr(), - } - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] - else: - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": "e7879e70-6335-5867-9ec6-957fbb3da6fa", - "name": "prepare", - "result": [("my_key", " prepared")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "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": "122f31bd-0e14-5b8f-91e7-4f241047a3fd", - "name": "tool_two_slow", - "result": [("my_key", " slow")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": {"my_key": "value prepared slow", "market": "DE"}, - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition:then"], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": "48a16051-2c14-5ff5-9cfe-e8c7c32d5c83", - "name": "finish", - "result": [("my_key", " finished")], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "configurable": {"thread_id": "10", "thread_ts": AnyStr()} - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - }, - }, - ] + }, + ] tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"] @@ -3368,7 +3195,6 @@ async def test_branch_then( ) async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: - saver.at = checkpoint_at tool_two = tool_two_graph.compile( checkpointer=saver, interrupt_after=["prepare"] ) @@ -3426,12 +3252,7 @@ async def test_branch_then( ) -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_in_one_fan_out_state_graph_waiting_edge( - checkpoint_at: CheckpointAt, -) -> None: +async def test_in_one_fan_out_state_graph_waiting_edge() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3495,7 +3316,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3519,12 +3340,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -3593,7 +3410,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3617,12 +3434,8 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( snapshot: SnapshotAssertion, - checkpoint_at: CheckpointAt, ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -3700,7 +3513,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3724,12 +3537,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ] -@pytest.mark.parametrize( - "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] -) -async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpoint_at: CheckpointAt, -) -> None: +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3798,7 +3606,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}}