From 65b64c13312fdc69dd2b8df4b50673171892a528 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 09:35:06 -0700 Subject: [PATCH 01/27] Support optional types from typeddict In the input schema --- libs/langgraph/langgraph/graph/state.py | 24 +++++----- libs/langgraph/langgraph/utils.py | 29 ++++++++++-- libs/langgraph/tests/test_state.py | 24 +++++++++- libs/langgraph/tests/test_utils.py | 63 ++++++++++++++++++++++++- 4 files changed, 121 insertions(+), 19 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 1416d14d0..10e8ca5ff 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -20,9 +20,7 @@ from typing import ( from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike -from langchain_core.runnables.utils import ( - create_model, -) +from langchain_core.runnables.utils import create_model from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate @@ -33,14 +31,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN from langgraph.errors import InvalidUpdateError -from langgraph.graph.graph import ( - END, - START, - Branch, - CompiledGraph, - Graph, - Send, -) +from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, @@ -53,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils import RunnableCallable, coerce_to_runnable +from langgraph.utils import RunnableCallable, coerce_to_runnable, is_optional_type logger = logging.getLogger(__name__) @@ -498,7 +489,14 @@ class CompiledStateGraph(CompiledGraph): return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ - k: (self.channels[k].UpdateType, None) + k: ( + self.channels[k].UpdateType, + ( + None + if is_optional_type(self.channels[k].UpdateType) + else ... + ), + ) for k in self.builder.schemas[self.builder.input] if isinstance(self.channels[k], BaseChannel) }, diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 64cb21128..9aa1c257c 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -4,7 +4,15 @@ import inspect import sys from contextvars import copy_context from functools import partial, wraps -from typing import Any, AsyncIterator, Awaitable, Callable, Optional +from typing import ( + Any, + AsyncIterator, + Awaitable, + Callable, + Optional, + Union, + get_origin, +) from langchain_core.runnables.base import ( Runnable, @@ -34,8 +42,6 @@ except ImportError: class StrEnum(str, enum.Enum): """A string enum.""" - pass - class RunnableCallable(Runnable): """A much simpler version of RunnableLambda that requires sync and async functions.""" @@ -183,3 +189,20 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab f"Expected a Runnable, callable or dict." f"Instead got an unsupported type: {type(thing)}" ) + + +def is_optional_type(type_: Any) -> bool: + """Check if a type is Optional.""" + + if hasattr(type_, "__origin__") and hasattr(type_, "__args__"): + origin = get_origin(type_) + if origin is Optional: + return True + if origin is Union: + return any( + arg is type(None) or is_optional_type(arg) for arg in type_.__args__ + ) + return origin is None + if hasattr(type_, "__bound__") and type_.__bound__ is not None: + return is_optional_type(type_.__bound__) + return type_ is None diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 4088e38c9..884895162 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,5 +1,5 @@ from typing import Annotated as Annotated2 -from typing import Any +from typing import Any, Optional import pytest from langchain_core.runnables import RunnableConfig @@ -86,3 +86,25 @@ def test_state_schema_with_type_hint(): for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): node_name = actions[i].__name__ assert c[node_name] == output_state + + +def test_state_schema_optional_values(): + class InputState(TypedDict): + val1: str + val2: Optional[str] + + class State(InputState): + val4: dict + + builder = StateGraph(State, input=InputState) + builder.add_node("n", lambda x: x) + builder.add_edge("__start__", "n") + graph = builder.compile() + model = graph.input_schema + json_schema = model.schema() + expected_required = {"val1"} + expected_optional = {"val2"} + assert set(json_schema["required"]) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 1e3bf8ea1..d5bfb53c7 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -1,7 +1,18 @@ import functools import sys import uuid -from typing import TypedDict +from typing import ( + Any, + Callable, + Dict, + ForwardRef, + List, + Literal, + Optional, + TypedDict, + TypeVar, + Union, +) from unittest.mock import patch import langsmith @@ -9,7 +20,7 @@ import pytest from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph -from langgraph.utils import is_async_callable, is_async_generator +from langgraph.utils import is_async_callable, is_async_generator, is_optional_type def test_is_async() -> None: @@ -119,3 +130,51 @@ async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) - with langsmith.tracing_context(enabled=True): res = await rt_graph.ainvoke({"foo": 1}) assert isinstance(res["node_run_id"], uuid.UUID) + + +def test_is_optional_type(): + assert is_optional_type(None) + assert not is_optional_type(type(None)) + assert is_optional_type(Optional[list]) + assert not is_optional_type(int) + assert is_optional_type(Optional[Literal[1, 2, 3]]) + assert not is_optional_type(Literal[1, 2, 3]) + assert is_optional_type(Optional[List[int]]) + assert is_optional_type(Optional[Dict[str, int]]) + assert not is_optional_type(List[Optional[int]]) + assert is_optional_type(Union[Optional[str], Optional[int]]) + assert is_optional_type( + Union[ + Union[Optional[str], Optional[int]], Union[Optional[float], Optional[dict]] + ] + ) + assert not is_optional_type(Union[Union[str, int], Union[float, dict]]) + + assert is_optional_type(Union[int, None]) + assert is_optional_type(Union[str, None, int]) + assert is_optional_type(Union[None, str, int]) + assert not is_optional_type(Union[int, str]) + + assert not is_optional_type(Any) # Do we actually want this? + assert is_optional_type(Optional[Any]) + + class MyClass: + pass + + assert is_optional_type(Optional[MyClass]) + assert not is_optional_type(MyClass) + assert is_optional_type(Optional[ForwardRef("MyClass")]) + assert not is_optional_type(ForwardRef("MyClass")) + + assert is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]]) + assert not is_optional_type(Union[List[int], Dict[str, Optional[int]]]) + + assert is_optional_type(Optional[Callable[[int], str]]) + assert not is_optional_type(Callable[[int], Optional[str]]) + + T = TypeVar("T") + assert is_optional_type(Optional[T]) + assert not is_optional_type(T) + + U = TypeVar("U", bound=Optional[T]) + assert is_optional_type(U) From a07ec9d8c800be0283fa6332248751fa92ce1db5 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 10:21:57 -0700 Subject: [PATCH 02/27] Handle total=false --- libs/langgraph/langgraph/graph/state.py | 7 ++++++- libs/langgraph/tests/test_state.py | 18 +++++++++++++----- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 10e8ca5ff..31fdfb4e3 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -486,6 +486,10 @@ class CompiledStateGraph(CompiledGraph): __root__=(self.channels[keys[0]].UpdateType, None), ) else: + is_total_false = ( + hasattr(self.builder.input, "__total__") + and self.builder.input.__total__ is False + ) return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ @@ -493,7 +497,8 @@ class CompiledStateGraph(CompiledGraph): self.channels[k].UpdateType, ( None - if is_optional_type(self.channels[k].UpdateType) + if is_total_false + or is_optional_type(self.channels[k].UpdateType) else ... ), ) diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 884895162..a3a0d2806 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -88,8 +88,9 @@ def test_state_schema_with_type_hint(): assert c[node_name] == output_state -def test_state_schema_optional_values(): - class InputState(TypedDict): +@pytest.mark.parametrize("total_", [True, False]) +def test_state_schema_optional_values(total_: bool): + class InputState(TypedDict, total=total_): # type: ignore val1: str val2: Optional[str] @@ -102,9 +103,16 @@ def test_state_schema_optional_values(): graph = builder.compile() model = graph.input_schema json_schema = model.schema() - expected_required = {"val1"} - expected_optional = {"val2"} - assert set(json_schema["required"]) == expected_required + + if total_ is False: + expected_required = set() + expected_optional = {"val2", "val1"} + else: + expected_required = {"val1"} + + expected_optional = {"val2"} + + assert set(json_schema.get("required", set())) == expected_required assert ( set(json_schema["properties"].keys()) == expected_required | expected_optional ) From 51f58279baf8d1eb97d4198b6b5daccdf2fe1c98 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 15:56:22 -0700 Subject: [PATCH 03/27] Update snapshot --- libs/langgraph/langgraph/graph/state.py | 13 +- libs/langgraph/langgraph/utils.py | 89 ++++++++++++-- .../tests/__snapshots__/test_pregel.ambr | 6 +- libs/langgraph/tests/test_state.py | 18 ++- libs/langgraph/tests/test_utils.py | 113 +++++++++++++----- 5 files changed, 182 insertions(+), 57 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 31fdfb4e3..8f4d8fa3d 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -44,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils import RunnableCallable, coerce_to_runnable, is_optional_type +from langgraph.utils import RunnableCallable, coerce_to_runnable, field_is_optional logger = logging.getLogger(__name__) @@ -486,10 +486,6 @@ class CompiledStateGraph(CompiledGraph): __root__=(self.channels[keys[0]].UpdateType, None), ) else: - is_total_false = ( - hasattr(self.builder.input, "__total__") - and self.builder.input.__total__ is False - ) return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ @@ -497,8 +493,11 @@ class CompiledStateGraph(CompiledGraph): self.channels[k].UpdateType, ( None - if is_total_false - or is_optional_type(self.channels[k].UpdateType) + if field_is_optional( + k, + self.channels[k].UpdateType, + self.builder.input, + ) else ... ), ) diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 9aa1c257c..81ab9b543 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -4,15 +4,7 @@ import inspect import sys from contextvars import copy_context from functools import partial, wraps -from typing import ( - Any, - AsyncIterator, - Awaitable, - Callable, - Optional, - Union, - get_origin, -) +from typing import Any, AsyncIterator, Awaitable, Callable, Optional, Type, Union from langchain_core.runnables.base import ( Runnable, @@ -27,7 +19,14 @@ from langchain_core.runnables.config import ( var_child_runnable_config, ) from langchain_core.runnables.utils import accepts_config -from typing_extensions import TypeGuard +from typing_extensions import ( + Annotated, + NotRequired, + ReadOnly, + Required, + TypeGuard, + get_origin, +) try: from langchain_core.runnables.config import _set_config_context @@ -191,7 +190,7 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab ) -def is_optional_type(type_: Any) -> bool: +def _is_optional_type(type_: Any) -> bool: """Check if a type is Optional.""" if hasattr(type_, "__origin__") and hasattr(type_, "__args__"): @@ -200,9 +199,73 @@ def is_optional_type(type_: Any) -> bool: return True if origin is Union: return any( - arg is type(None) or is_optional_type(arg) for arg in type_.__args__ + arg is type(None) or _is_optional_type(arg) for arg in type_.__args__ ) return origin is None if hasattr(type_, "__bound__") and type_.__bound__ is not None: - return is_optional_type(type_.__bound__) + return _is_optional_type(type_.__bound__) return type_ is None + + +def _is_required_type(type_: Any) -> Optional[bool]: + """Check if an annotation is marked as Required/NotRequired. + + Returns: + - True if required + - False if not required + - None if not annotated with either + """ + origin = get_origin(type_) + if origin is Annotated or origin: + # See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated + return _is_required_type(type_.__args__[0]) + if origin is Required: + return True + if origin is NotRequired: + return False + return None + + +def _is_readonly_type(type_: Any) -> bool: + """Check if an annotation is marked as ReadOnly. + + Returns: + - True if is read only + - False if not read only + """ + + # See: https://typing.readthedocs.io/en/latest/spec/typeddict.html#typing-readonly-type-qualifier + origin = get_origin(type_) + if origin is Annotated: + return _is_readonly_type(type_.__args__[0]) + if origin is ReadOnly: + return True + return False + + +_DEFAULT_KEYS = frozenset() + + +def field_is_optional(name: str, type_: Any, schema: Type[Any]) -> bool: + """Determine if the field is optional for a graph's input. + + This is based on: + If TypedDict: + - Required/NotRequired + - total=False -> everything optional + - Type annotation (Optional/Union[None]) + """ + optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS) + if name in optional_keys: + # Either total=False or explicit NotRequired. + # No type annotation trumps this. + return True + if _is_required_type(type_): + # Handle Required[] + # (we already handled NotRequired and total=False) + return False + # Note, we ignore ReadOnly attributes, + # as they don't make much sense. (we don't care if you mutate the state in your node) + # and mutating state in your node has no effect on our graph state. + # Base case is the annotation + return _is_optional_type(type_) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index ba58cda8c..725955315 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -366,7 +366,7 @@ ''' # --- # name: test_conditional_entrypoint_to_multiple_state_graph - '{"title": "LangGraphInput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}, "required": ["locations", "results"]}' # --- # name: test_conditional_entrypoint_to_multiple_state_graph.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}}' @@ -4855,7 +4855,7 @@ ''' # --- # name: test_prebuilt_tool_chat - '{"title": "LangGraphInput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "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", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' + '{"title": "LangGraphInput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"], "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "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", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' # --- # name: test_prebuilt_tool_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "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", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' @@ -5094,7 +5094,7 @@ '{"title": "LangGraphConfig", "type": "object", "properties": {"configurable": {"$ref": "#/definitions/Configurable"}}, "definitions": {"Configurable": {"title": "Configurable", "type": "object", "properties": {"tools": {"title": "Tools", "type": "array", "items": {"type": "string"}}}}}}' # --- # name: test_state_graph_w_config_inherited_state_keys.1 - '{"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"]}}}' + '{"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"}]}}}, "required": ["input"], "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_state_graph_w_config_inherited_state_keys.2 '{"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"]}}}' diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index a3a0d2806..005f0ffd1 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,5 +1,5 @@ from typing import Annotated as Annotated2 -from typing import Any, Optional +from typing import Any, NotRequired, Optional, Required import pytest from langchain_core.runnables import RunnableConfig @@ -90,11 +90,19 @@ def test_state_schema_with_type_hint(): @pytest.mark.parametrize("total_", [True, False]) def test_state_schema_optional_values(total_: bool): - class InputState(TypedDict, total=total_): # type: ignore + class SomeParentState(TypedDict): + val0a: str + val0b: Optional[str] + + class InputState(SomeParentState, total=total_): # type: ignore val1: str val2: Optional[str] + val3: Required[str] + val4: NotRequired[dict] + val5: Annotated[Required[str], "foo"] + val6: Annotated[NotRequired[str], "bar"] - class State(InputState): + class State(InputState): # this would be ignored val4: dict builder = StateGraph(State, input=InputState) @@ -112,6 +120,10 @@ def test_state_schema_optional_values(total_: bool): expected_optional = {"val2"} + # The others should always have precedence based on the required annotation + expected_required |= {"val0a", "val3", "val5"} + expected_optional |= {"val0b", "val4", "val6"} + assert set(json_schema.get("required", set())) == expected_required assert ( set(json_schema["properties"].keys()) == expected_required | expected_optional diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index d5bfb53c7..2e5d2b3bb 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -17,10 +17,16 @@ from unittest.mock import patch import langsmith import pytest +from typing_extensions import Annotated, NotRequired, Required from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph -from langgraph.utils import is_async_callable, is_async_generator, is_optional_type +from langgraph.utils import ( + _is_optional_type, + field_is_optional, + is_async_callable, + is_async_generator, +) def test_is_async() -> None: @@ -133,48 +139,93 @@ async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) - def test_is_optional_type(): - assert is_optional_type(None) - assert not is_optional_type(type(None)) - assert is_optional_type(Optional[list]) - assert not is_optional_type(int) - assert is_optional_type(Optional[Literal[1, 2, 3]]) - assert not is_optional_type(Literal[1, 2, 3]) - assert is_optional_type(Optional[List[int]]) - assert is_optional_type(Optional[Dict[str, int]]) - assert not is_optional_type(List[Optional[int]]) - assert is_optional_type(Union[Optional[str], Optional[int]]) - assert is_optional_type( + assert _is_optional_type(None) + assert not _is_optional_type(type(None)) + assert _is_optional_type(Optional[list]) + assert not _is_optional_type(int) + assert _is_optional_type(Optional[Literal[1, 2, 3]]) + assert not _is_optional_type(Literal[1, 2, 3]) + assert _is_optional_type(Optional[List[int]]) + assert _is_optional_type(Optional[Dict[str, int]]) + assert not _is_optional_type(List[Optional[int]]) + assert _is_optional_type(Union[Optional[str], Optional[int]]) + assert _is_optional_type( Union[ Union[Optional[str], Optional[int]], Union[Optional[float], Optional[dict]] ] ) - assert not is_optional_type(Union[Union[str, int], Union[float, dict]]) + assert not _is_optional_type(Union[Union[str, int], Union[float, dict]]) - assert is_optional_type(Union[int, None]) - assert is_optional_type(Union[str, None, int]) - assert is_optional_type(Union[None, str, int]) - assert not is_optional_type(Union[int, str]) + assert _is_optional_type(Union[int, None]) + assert _is_optional_type(Union[str, None, int]) + assert _is_optional_type(Union[None, str, int]) + assert not _is_optional_type(Union[int, str]) - assert not is_optional_type(Any) # Do we actually want this? - assert is_optional_type(Optional[Any]) + assert not _is_optional_type(Any) # Do we actually want this? + assert _is_optional_type(Optional[Any]) class MyClass: pass - assert is_optional_type(Optional[MyClass]) - assert not is_optional_type(MyClass) - assert is_optional_type(Optional[ForwardRef("MyClass")]) - assert not is_optional_type(ForwardRef("MyClass")) + assert _is_optional_type(Optional[MyClass]) + assert not _is_optional_type(MyClass) + assert _is_optional_type(Optional[ForwardRef("MyClass")]) + assert not _is_optional_type(ForwardRef("MyClass")) - assert is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]]) - assert not is_optional_type(Union[List[int], Dict[str, Optional[int]]]) + assert _is_optional_type(Optional[Union[List[int], Dict[str, Optional[int]]]]) + assert not _is_optional_type(Union[List[int], Dict[str, Optional[int]]]) - assert is_optional_type(Optional[Callable[[int], str]]) - assert not is_optional_type(Callable[[int], Optional[str]]) + assert _is_optional_type(Optional[Callable[[int], str]]) + assert not _is_optional_type(Callable[[int], Optional[str]]) T = TypeVar("T") - assert is_optional_type(Optional[T]) - assert not is_optional_type(T) + assert _is_optional_type(Optional[T]) + assert not _is_optional_type(T) - U = TypeVar("U", bound=Optional[T]) - assert is_optional_type(U) + U = TypeVar("U", bound=Optional[T]) # type: ignore + assert _is_optional_type(U) + + +def test_is_required(): + class MyBaseTypedDict(TypedDict): + val_1: Required[Optional[str]] + val_2: Required[str] + val_3: NotRequired[str] + val_4: NotRequired[Optional[str]] + val_5: Annotated[NotRequired[int], "foo"] + val_6: NotRequired[Annotated[int, "foo"]] + val_7: Annotated[Required[int], "foo"] + val_8: Required[Annotated[int, "foo"]] + val_9: Optional[str] + val_10: str + + annos = MyBaseTypedDict.__annotations__ + assert not field_is_optional("val_1", annos["val_1"], MyBaseTypedDict) + assert not field_is_optional("val_2", annos["val_2"], MyBaseTypedDict) + assert field_is_optional("val_3", annos["val_3"], MyBaseTypedDict) + assert field_is_optional("val_4", annos["val_4"], MyBaseTypedDict) + # See https://peps.python.org/pep-0655/#interaction-with-annotated + assert field_is_optional("val_5", annos["val_5"], MyBaseTypedDict) + assert field_is_optional("val_6", annos["val_6"], MyBaseTypedDict) + assert not field_is_optional("val_7", annos["val_7"], MyBaseTypedDict) + assert not field_is_optional("val_8", annos["val_8"], MyBaseTypedDict) + assert field_is_optional("val_9", annos["val_9"], MyBaseTypedDict) + assert not field_is_optional("val_10", annos["val_10"], MyBaseTypedDict) + + class MyChildDict(MyBaseTypedDict): + val_11: int + val_11b: Optional[int] + val_11c: Union[int, None, str] + + class MyGrandChildDict(MyChildDict, total=False): + val_12: int + val_13: Required[str] + + cannos = MyChildDict.__annotations__ + gcannos = MyGrandChildDict.__annotations__ + assert not field_is_optional("val_11", cannos["val_11"], MyChildDict) + assert field_is_optional("val_11b", cannos["val_11b"], MyChildDict) + assert field_is_optional("val_11c", cannos["val_11c"], MyChildDict) + assert field_is_optional("val_12", gcannos["val_12"], MyGrandChildDict) + assert field_is_optional("val_9", gcannos["val_9"], MyGrandChildDict) + assert not field_is_optional("val_13", gcannos["val_13"], MyGrandChildDict) From 55263c640f5a737782730cbc4cde852a6a65b262 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 16:28:01 -0700 Subject: [PATCH 04/27] Return a default value instead --- libs/langgraph/langgraph/graph/state.py | 6 ++--- libs/langgraph/langgraph/utils.py | 12 +++++---- libs/langgraph/tests/test_utils.py | 34 ++++++++++++------------- 3 files changed, 26 insertions(+), 26 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 8f4d8fa3d..f576a98cd 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -44,7 +44,7 @@ from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils import RunnableCallable, coerce_to_runnable, field_is_optional +from langgraph.utils import RunnableCallable, coerce_to_runnable, get_field_default logger = logging.getLogger(__name__) @@ -492,13 +492,11 @@ class CompiledStateGraph(CompiledGraph): k: ( self.channels[k].UpdateType, ( - None - if field_is_optional( + get_field_default( k, self.channels[k].UpdateType, self.builder.input, ) - else ... ), ) for k in self.builder.schemas[self.builder.input] diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 81ab9b543..1fbcc94e3 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -246,8 +246,8 @@ def _is_readonly_type(type_: Any) -> bool: _DEFAULT_KEYS = frozenset() -def field_is_optional(name: str, type_: Any, schema: Type[Any]) -> bool: - """Determine if the field is optional for a graph's input. +def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any: + """Determine the default value for a field in a state schema. This is based on: If TypedDict: @@ -259,13 +259,15 @@ def field_is_optional(name: str, type_: Any, schema: Type[Any]) -> bool: if name in optional_keys: # Either total=False or explicit NotRequired. # No type annotation trumps this. - return True + return None if _is_required_type(type_): # Handle Required[] # (we already handled NotRequired and total=False) - return False + return ... # Note, we ignore ReadOnly attributes, # as they don't make much sense. (we don't care if you mutate the state in your node) # and mutating state in your node has no effect on our graph state. # Base case is the annotation - return _is_optional_type(type_) + if _is_optional_type(type_): + return None + return ... diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 2e5d2b3bb..936fd2847 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -23,7 +23,7 @@ from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.utils import ( _is_optional_type, - field_is_optional, + get_field_default, is_async_callable, is_async_generator, ) @@ -200,17 +200,17 @@ def test_is_required(): val_10: str annos = MyBaseTypedDict.__annotations__ - assert not field_is_optional("val_1", annos["val_1"], MyBaseTypedDict) - assert not field_is_optional("val_2", annos["val_2"], MyBaseTypedDict) - assert field_is_optional("val_3", annos["val_3"], MyBaseTypedDict) - assert field_is_optional("val_4", annos["val_4"], MyBaseTypedDict) + assert get_field_default("val_1", annos["val_1"], MyBaseTypedDict) == ... + assert get_field_default("val_2", annos["val_2"], MyBaseTypedDict) == ... + assert get_field_default("val_3", annos["val_3"], MyBaseTypedDict) is None + assert get_field_default("val_4", annos["val_4"], MyBaseTypedDict) is None # See https://peps.python.org/pep-0655/#interaction-with-annotated - assert field_is_optional("val_5", annos["val_5"], MyBaseTypedDict) - assert field_is_optional("val_6", annos["val_6"], MyBaseTypedDict) - assert not field_is_optional("val_7", annos["val_7"], MyBaseTypedDict) - assert not field_is_optional("val_8", annos["val_8"], MyBaseTypedDict) - assert field_is_optional("val_9", annos["val_9"], MyBaseTypedDict) - assert not field_is_optional("val_10", annos["val_10"], MyBaseTypedDict) + assert get_field_default("val_5", annos["val_5"], MyBaseTypedDict) is None + assert get_field_default("val_6", annos["val_6"], MyBaseTypedDict) is None + assert get_field_default("val_7", annos["val_7"], MyBaseTypedDict) == ... + assert get_field_default("val_8", annos["val_8"], MyBaseTypedDict) == ... + assert get_field_default("val_9", annos["val_9"], MyBaseTypedDict) is None + assert get_field_default("val_10", annos["val_10"], MyBaseTypedDict) == ... class MyChildDict(MyBaseTypedDict): val_11: int @@ -223,9 +223,9 @@ def test_is_required(): cannos = MyChildDict.__annotations__ gcannos = MyGrandChildDict.__annotations__ - assert not field_is_optional("val_11", cannos["val_11"], MyChildDict) - assert field_is_optional("val_11b", cannos["val_11b"], MyChildDict) - assert field_is_optional("val_11c", cannos["val_11c"], MyChildDict) - assert field_is_optional("val_12", gcannos["val_12"], MyGrandChildDict) - assert field_is_optional("val_9", gcannos["val_9"], MyGrandChildDict) - assert not field_is_optional("val_13", gcannos["val_13"], MyGrandChildDict) + assert get_field_default("val_11", cannos["val_11"], MyChildDict) == ... + assert get_field_default("val_11b", cannos["val_11b"], MyChildDict) is None + assert get_field_default("val_11c", cannos["val_11c"], MyChildDict) is None + assert get_field_default("val_12", gcannos["val_12"], MyGrandChildDict) is None + assert get_field_default("val_9", gcannos["val_9"], MyGrandChildDict) is None + assert get_field_default("val_13", gcannos["val_13"], MyGrandChildDict) == ... From 1191dc1b4880624fc5c5efe93e8bac24b91c5e49 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Fri, 30 Aug 2024 23:52:41 -0700 Subject: [PATCH 05/27] Handle py3.9 --- libs/langgraph/langgraph/utils.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/utils.py b/libs/langgraph/langgraph/utils.py index 5cab1b7f9..cc3fd4760 100644 --- a/libs/langgraph/langgraph/utils.py +++ b/libs/langgraph/langgraph/utils.py @@ -218,13 +218,13 @@ def _is_required_type(type_: Any) -> Optional[bool]: - None if not annotated with either """ origin = get_origin(type_) - if origin is Annotated or origin: - # See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated - return _is_required_type(type_.__args__[0]) if origin is Required: return True if origin is NotRequired: return False + if origin is Annotated or getattr(origin, "__args__", None): + # See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated + return _is_required_type(type_.__args__[0]) return None @@ -258,14 +258,21 @@ def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any: - Type annotation (Optional/Union[None]) """ optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS) + irq = _is_required_type(type_) if name in optional_keys: # Either total=False or explicit NotRequired. # No type annotation trumps this. + if irq: + # Unless it's earlier versions of python & explicit Required + return ... + return None + if irq is not None: + if irq: + # Handle Required[] + # (we already handled NotRequired and total=False) + return ... + # Handle NotRequired[] for earlier versions of python return None - if _is_required_type(type_): - # Handle Required[] - # (we already handled NotRequired and total=False) - return ... # Note, we ignore ReadOnly attributes, # as they don't make much sense. (we don't care if you mutate the state in your node) # and mutating state in your node has no effect on our graph state. From 3a07fe2efb6567acff6d792230f5369d7aa4abe6 Mon Sep 17 00:00:00 2001 From: gbaian10 <34255899+gbaian10@users.noreply.github.com> Date: Sun, 1 Sep 2024 07:35:59 +0800 Subject: [PATCH 06/27] docs: fix agent_supervisor.ipynb (#1562) fix #1561 --- examples/multi_agent/agent_supervisor.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/multi_agent/agent_supervisor.ipynb b/examples/multi_agent/agent_supervisor.ipynb index bd2a7c6ce..b1ecfe31b 100644 --- a/examples/multi_agent/agent_supervisor.ipynb +++ b/examples/multi_agent/agent_supervisor.ipynb @@ -219,7 +219,7 @@ "workflow = StateGraph(AgentState)\n", "workflow.add_node(\"Researcher\", research_node)\n", "workflow.add_node(\"Coder\", code_node)\n", - "workflow.add_node(\"supervisor\", supervisor_chain)" + "workflow.add_node(\"supervisor\", supervisor_agent)" ] }, { From 8c698fc36252bb6a84b09d98c4d0bab2fbf98328 Mon Sep 17 00:00:00 2001 From: Reuben M <120103767+McCReuben@users.noreply.github.com> Date: Sat, 31 Aug 2024 19:41:48 -0400 Subject: [PATCH 07/27] Return replan response in dict. Add Replan example (#1566) --- examples/llm-compiler/LLMCompiler.ipynb | 43 ++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/examples/llm-compiler/LLMCompiler.ipynb b/examples/llm-compiler/LLMCompiler.ipynb index 98f5416d8..3ae67e33b 100644 --- a/examples/llm-compiler/LLMCompiler.ipynb +++ b/examples/llm-compiler/LLMCompiler.ipynb @@ -646,11 +646,12 @@ "def _parse_joiner_output(decision: JoinOutputs) -> List[BaseMessage]:\n", " response = [AIMessage(content=f\"Thought: {decision.thought}\")]\n", " if isinstance(decision.action, Replan):\n", - " return response + [\n", + " return {\"messages\": response + [\n", " SystemMessage(\n", " content=f\"Context from last attempt: {decision.action.feedback}\"\n", " )\n", " ]\n", + " }\n", " else:\n", " return {\"messages\": response + [AIMessage(content=decision.action.response)]}\n", "\n", @@ -933,6 +934,46 @@ "print(step['join']['messages'][-1].content)" ] }, + { + "cell_type": "markdown", + "id": "f9487866", + "metadata": {}, + "source": [ + "#### Complex Replanning Example\n", + "\n", + "This question is likely to prompt the Replan functionality, but it may need to be run multiple times to see this in action." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "391d6931", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'plan_and_schedule': {'messages': [FunctionMessage(content=\"[{'url': 'https://www.timeanddate.com/weather/japan/tokyo', 'content': '88 / 84 °F. 13. 87 / 82 °F. 14. 84 / 80 °F. Detailed forecast for 14 days. Need some help? Current weather in Tokyo and forecast for today, tomorrow, and next 14 days.'}]\", additional_kwargs={'idx': 1, 'args': {'query': 'current temperature in Tokyo'}}, name='tavily_search_results_json', tool_call_id=1), FunctionMessage(content='join', additional_kwargs={'idx': 2, 'args': ()}, name='join', tool_call_id=2)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides the current temperature in Tokyo but does not explicitly state which temperature (88 / 84 °F) corresponds to the current condition. It seems to be a range, possibly the day's high and low. Without a clear indication of the exact current temperature, it's challenging to provide a precise flashcard summary.\", id='8ef2a131-69db-4180-a76e-fd9d6f4037c1'), SystemMessage(content='Context from last attempt: The information provided does not explicitly state the current temperature in Tokyo; it provides a temperature range without specifying which is the current temperature. Need to find a source that gives the exact current temperature in Tokyo for a precise flashcard summary.', id='f5bd752c-b068-459a-8d9e-bd1f1b5fa4fe')]}}\n", + "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 3, 'args': ()}, name='join', tool_call_id=3)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='3cc41891-4f47-4453-8edf-b989926ab25e'), SystemMessage(content='Context from last attempt: The search did not provide an exact current temperature for Tokyo, making it impossible to create a precise flashcard. A source that explicitly states the current temperature is needed for an accurate response.', id='96290b41-a4c4-4ab5-829a-89cc31dfe6c8')]}}\n", + "{'plan_and_schedule': {'messages': [FunctionMessage(content='join', additional_kwargs={'idx': 4, 'args': ()}, name='join', tool_call_id=4)]}}\n", + "{'join': {'messages': [AIMessage(content=\"Thought: The search result provides a temperature range for Tokyo but does not specify the current temperature. This makes it challenging to create a precise flashcard without an exact current temperature. The user's request cannot be fully satisfied without this detail.\", id='4724b242-ddb8-47e6-b235-de25de54fe45'), AIMessage(content='I was unable to find the exact current temperature in Tokyo. However, the temperature range for today in Tokyo is between 88°F and 84°F. For the most accurate and up-to-date temperature, I recommend checking a reliable weather forecasting website or app.', id='40e29a47-a001-4f65-a18f-65c2931d1ae5')]}}\n" + ] + } + ], + "source": [ + "for step in chain.stream({\"messages\":\n", + " [\n", + " HumanMessage(\n", + " content=\"Find the current temperature in Tokyo, then, respond with a flashcard summarizing this information\"\n", + " )\n", + " ]}\n", + "):\n", + " print(step)" + ] + }, { "cell_type": "markdown", "id": "c647d5f3-5e00-4449-9cec-5a9f438c9cff", From a93775413281df9ddf6ba29cc388b2460d94b9af Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 1 Sep 2024 12:15:20 -0700 Subject: [PATCH 08/27] lib 0.2.16 --- 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 36cdb7fdc..b7064ee22 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.15" +version = "0.2.16" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" From d65d0ec665896849bda3f2c370299fc267226fbb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:24:54 -0700 Subject: [PATCH 09/27] test: Ensure that async checkpointers in tests are setup in same event loop as they are used in - This more closely remembers the environment they're used in, so it's what we should be testing - Remove unnecessary pytest-asyncio dependency, use anyio pytest plugin instead - Convert remaining async tests using only memory checkpointer to use all existing ones --- libs/langgraph/Makefile | 2 +- libs/langgraph/poetry.lock | 35 +- libs/langgraph/pyproject.toml | 4 +- .../__snapshots__/test_pregel_async.ambr | 1100 ++ libs/langgraph/tests/checkpoint/__init__.py | 0 libs/langgraph/tests/conftest.py | 105 +- libs/langgraph/tests/test_channels.py | 2 + libs/langgraph/tests/test_interruption.py | 41 +- libs/langgraph/tests/test_prebuilt.py | 74 +- libs/langgraph/tests/test_pregel.py | 28 +- libs/langgraph/tests/test_pregel_async.py | 8989 +++++++++-------- libs/langgraph/tests/test_state.py | 4 +- libs/langgraph/tests/test_store.py | 3 + libs/langgraph/tests/test_utils.py | 2 + 14 files changed, 5975 insertions(+), 4414 deletions(-) delete mode 100644 libs/langgraph/tests/checkpoint/__init__.py diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index ccb8ef59a..d5c6d7ee3 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -15,7 +15,7 @@ coverage: --cov-report term-missing:skip-covered start-postgres: - docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait + docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait --remove-orphans stop-postgres: docker compose -f tests/compose-postgres.yml down -v diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 241d47b6e..50d6c6f89 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -2829,13 +2829,13 @@ diagrams = ["jinja2", "railroad-diagrams"] [[package]] name = "pytest" -version = "7.4.4" +version = "8.3.2" description = "pytest: simple powerful testing with Python" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8"}, - {file = "pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280"}, + {file = "pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5"}, + {file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"}, ] [package.dependencies] @@ -2843,29 +2843,11 @@ colorama = {version = "*", markers = "sys_platform == \"win32\""} exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""} iniconfig = "*" packaging = "*" -pluggy = ">=0.12,<2.0" -tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""} +pluggy = ">=1.5,<2" +tomli = {version = ">=1", markers = "python_version < \"3.11\""} [package.extras] -testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] - -[[package]] -name = "pytest-asyncio" -version = "0.20.3" -description = "Pytest support for asyncio" -optional = false -python-versions = ">=3.7" -files = [ - {file = "pytest-asyncio-0.20.3.tar.gz", hash = "sha256:83cbf01169ce3e8eb71c6c278ccb0574d1a7a3bb8eaaf5e50e0ad342afb33b36"}, - {file = "pytest_asyncio-0.20.3-py3-none-any.whl", hash = "sha256:f129998b209d04fcc65c96fc85c11e5316738358909a8399e93be553d7656442"}, -] - -[package.dependencies] -pytest = ">=6.1.0" - -[package.extras] -docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"] -testing = ["coverage (>=6.2)", "flaky (>=3.5.0)", "hypothesis (>=5.7.1)", "mypy (>=0.931)", "pytest-trio (>=0.7.0)"] +dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] [[package]] name = "pytest-cov" @@ -3069,7 +3051,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -4310,4 +4291,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "7e0d6fd967987fcf46b038c47b5ca81d86cb8fff89704eac2074f2ce55336a2b" +content-hash = "e0787721e6cb80996a39284c74f1ddee4789c6060d6c2272cb1957f4686bd478" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index b7064ee22..d6e04486f 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -14,10 +14,9 @@ langgraph-checkpoint = "^1.0.2" [tool.poetry.group.dev.dependencies] -pytest = "^7.3.0" +pytest = "^8.3.2" pytest-cov = "^4.0.0" pytest-dotenv = "^0.5.2" -pytest-asyncio = "^0.20.3" pytest-mock = "^3.10.0" syrupy = "^4.0.2" httpx = "^0.26.0" @@ -73,7 +72,6 @@ requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" [tool.pytest.ini_options] -asyncio_mode = "auto" # --strict-markers will raise errors on unknown marks. # https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks # diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index c0d2e782c..3aaccffd1 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -36,6 +36,191 @@ +---------+ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[memory] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pipe] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[postgres_aio_pool] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class[sqlite_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2 ''' graph TD; @@ -157,6 +342,611 @@ '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({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio] + ''' + 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_aio].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_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_aio_pipe].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool] + ''' + 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_aio_pool].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio] + ''' + 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_aio].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_via_branch ''' +-----------+ @@ -194,6 +984,191 @@ +---------+ ''' # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pipe] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[postgres_aio_pool] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[sqlite_aio] + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- # name: test_nested_graph ''' +-----------+ @@ -219,3 +1194,128 @@ +---------+ ''' # --- +# name: test_weather_subgraph[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[postgres_aio_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[sqlite_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node
__interrupt = before) + __end__([__end__]):::last + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- diff --git a/libs/langgraph/tests/checkpoint/__init__.py b/libs/langgraph/tests/checkpoint/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index bb6ef9603..ca6351c6b 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,8 +1,6 @@ -import asyncio import sys -from concurrent.futures import ThreadPoolExecutor -from contextlib import asynccontextmanager, contextmanager -from typing import AsyncIterator, Iterator, TypeVar +from contextlib import asynccontextmanager +from typing import AsyncIterator, Optional from uuid import UUID, uuid4 import pytest @@ -10,6 +8,7 @@ from psycopg import AsyncConnection, Connection from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.sqlite import SqliteSaver @@ -19,6 +18,11 @@ from tests.memory_assert import MemorySaverAssertImmutable DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" +@pytest.fixture +def anyio_backend(): + return "asyncio" + + @pytest.fixture() def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: side_effect = ( @@ -27,35 +31,6 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: return mocker.patch("uuid.uuid4", side_effect=side_effect) -""" -pytest-asyncio doesn't support calling async fixtures with getfixturevalue -so we need to use ThreadPoolExecutor to run the async fixture in a thread -https://github.com/pytest-dev/pytest-asyncio/issues/112#issuecomment-462062890 -""" -T = TypeVar("T") - - -def close_loop(loop: asyncio.AbstractEventLoop) -> None: - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.run_until_complete(loop.shutdown_default_executor()) - asyncio.set_event_loop(None) - loop.close() - - -@contextmanager -def agen_to_gen(agen: AsyncIterator[T]) -> Iterator[T]: - with ThreadPoolExecutor(1) as bg: - loop = asyncio.new_event_loop() - bg.submit(asyncio.set_event_loop, loop).result() - try: - yield bg.submit(loop.run_until_complete, agen.__aenter__()).result() - finally: - bg.submit( - loop.run_until_complete, agen.__aexit__(None, None, None) - ).result() - bg.submit(close_loop, loop).result() - - # checkpointer fixtures @@ -70,12 +45,6 @@ def checkpointer_sqlite(): yield checkpointer -@pytest.fixture(scope="function") -def checkpointer_sqlite_aio(): - with agen_to_gen(_checkpointer_sqlite_aio()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_sqlite_aio(): async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: @@ -143,16 +112,10 @@ def checkpointer_postgres_pool(): conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -174,16 +137,10 @@ async def _checkpointer_postgres_aio(): await conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio_pipe(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio_pipe()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -208,16 +165,10 @@ async def _checkpointer_postgres_aio_pipe(): await conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def checkpointer_postgres_aio_pool(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - with agen_to_gen(_checkpointer_postgres_aio_pool()) as checkpointer: - yield checkpointer - - @asynccontextmanager async def _checkpointer_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") database = f"test_{uuid4().hex[:16]}" # create unique db async with await AsyncConnection.connect( @@ -240,6 +191,30 @@ async def _checkpointer_postgres_aio_pool(): await conn.execute(f"DROP DATABASE {database}") +@asynccontextmanager +async def awith_checkpointer( + checkpointer_name: Optional[str], +) -> AsyncIterator[BaseCheckpointSaver]: + if checkpointer_name is None: + yield None + elif checkpointer_name == "memory": + yield MemorySaverAssertImmutable() + elif checkpointer_name == "sqlite_aio": + async with _checkpointer_sqlite_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio": + async with _checkpointer_postgres_aio() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pipe": + async with _checkpointer_postgres_aio_pipe() as checkpointer: + yield checkpointer + elif checkpointer_name == "postgres_aio_pool": + async with _checkpointer_postgres_aio_pool() as checkpointer: + yield checkpointer + else: + raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") + + ALL_CHECKPOINTERS_SYNC = [ "memory", "sqlite", @@ -254,3 +229,7 @@ ALL_CHECKPOINTERS_ASYNC = [ "postgres_aio_pipe", "postgres_aio_pool", ] +ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ + *ALL_CHECKPOINTERS_ASYNC, + None, +] diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index ff393577e..69a624969 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -8,6 +8,8 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.errors import EmptyChannelError, InvalidUpdateError +pytestmark = pytest.mark.anyio + def test_last_value() -> None: with LastValue(int).from_checkpoint(None, {}) as channel: diff --git a/libs/langgraph/tests/test_interruption.py b/libs/langgraph/tests/test_interruption.py index e6682e964..d4e618ddf 100644 --- a/libs/langgraph/tests/test_interruption.py +++ b/libs/langgraph/tests/test_interruption.py @@ -4,12 +4,16 @@ import pytest from pytest_mock import MockerFixture from langgraph.graph import END, START, StateGraph - - -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_SYNC, + awith_checkpointer, ) + +pytestmark = pytest.mark.anyio + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interruption_without_state_updates( request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: @@ -47,12 +51,9 @@ def test_interruption_without_state_updates( assert graph.get_state(thread).next == () -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interruption_without_state_updates_async( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ): """Test interruption without state updates. This test confirms that interrupting doesn't require a state key having been updated in the prev step""" @@ -72,17 +73,17 @@ async def test_interruption_without_state_updates_async( builder.add_edge("step_2", "step_3") builder.add_edge("step_3", END) - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer, interrupt_after="*") - initial_input = {"input": "hello world"} - thread = {"configurable": {"thread_id": "1"}} + initial_input = {"input": "hello world"} + thread = {"configurable": {"thread_id": "1"}} - await graph.ainvoke(initial_input, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_2",) + await graph.ainvoke(initial_input, thread, debug=True) + assert (await graph.aget_state(thread)).next == ("step_2",) - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == ("step_3",) + await graph.ainvoke(None, thread, debug=True) + assert (await graph.aget_state(thread)).next == ("step_3",) - await graph.ainvoke(None, thread, debug=True) - assert (await graph.aget_state(thread)).next == () + await graph.ainvoke(None, thread, debug=True) + assert (await graph.aget_state(thread)).next == () diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index b29074dff..82fd5e613 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -23,8 +23,15 @@ from pydantic import BaseModel as BaseModelV2 from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent from langgraph.prebuilt.tool_node import InjectedState +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_SYNC, + awith_checkpointer, +) from tests.messages import _AnyIdHumanMessage +pytestmark = pytest.mark.anyio + class FakeToolCallingModel(BaseChatModel): def _generate( @@ -53,10 +60,7 @@ class FakeToolCallingModel(BaseChatModel): return self -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite", "postgres", "postgres_pipe"], -) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer: BaseCheckpointSaver = request.getfixturevalue( "checkpointer_" + checkpointer_name @@ -89,43 +93,35 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> assert saved.pending_writes == [] -@pytest.mark.parametrize( - "checkpointer_name", - ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], -) -async def test_no_modifier_async( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_no_modifier_async(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + model = FakeToolCallingModel() - model = FakeToolCallingModel() + agent = create_react_agent(model, [], checkpointer=checkpointer) + inputs = [HumanMessage("hi?")] + thread = {"configurable": {"thread_id": "123"}} + response = await agent.ainvoke({"messages": inputs}, thread, debug=True) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response - agent = create_react_agent(model, [], checkpointer=checkpointer) - inputs = [HumanMessage("hi?")] - thread = {"configurable": {"thread_id": "123"}} - response = await agent.ainvoke({"messages": inputs}, thread, debug=True) - expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} - assert response == expected_response - - if checkpointer: - saved = await checkpointer.aget_tuple(thread) - assert saved is not None - assert saved.checkpoint["channel_values"] == { - "messages": [ - _AnyIdHumanMessage(content="hi?"), - AIMessage(content="hi?", id="0"), - ], - "agent": "agent", - } - assert saved.metadata == { - "parents": {}, - "source": "loop", - "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, - "step": 1, - } - assert saved.pending_writes == [] + if checkpointer: + saved = await checkpointer.aget_tuple(thread) + assert saved is not None + assert saved.checkpoint["channel_values"] == { + "messages": [ + _AnyIdHumanMessage(content="hi?"), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + } + assert saved.metadata == { + "parents": {}, + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] def test_passing_two_modifiers(): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2f85540f1..3ca3906d1 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -74,10 +74,7 @@ from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ALL_CHECKPOINTERS_SYNC from tests.fake_tracer import FakeTracer -from tests.memory_assert import ( - MemorySaverAssertCheckpointMetadata, - MemorySaverNoPending, -) +from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -1626,29 +1623,6 @@ def test_cond_edge_after_send() -> None: assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"] -async def test_checkpointer_null_pending_writes() -> None: - class Node: - def __init__(self, name: str): - self.name = name - setattr(self, "__name__", name) - - def __call__(self, state): - return [self.name] - - builder = StateGraph(Annotated[list, operator.add]) - builder.add_node(Node("1")) - builder.add_edge(START, "1") - graph = builder.compile(checkpointer=MemorySaverNoPending()) - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] - assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 3 - assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ - "1" - ] * 4 - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2f03a546e..4d004075d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -17,6 +17,7 @@ from typing import ( Tuple, TypedDict, Union, + cast, ) from uuid import UUID @@ -43,7 +44,6 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( - BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, @@ -71,14 +71,20 @@ from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence -from tests.conftest import ALL_CHECKPOINTERS_ASYNC +from tests.conftest import ( + ALL_CHECKPOINTERS_ASYNC, + ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, + awith_checkpointer, +) from tests.fake_tracer import FakeTracer from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, - MemorySaverAssertImmutable, + MemorySaverNoPending, ) from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage +pytestmark = pytest.mark.anyio + async def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): @@ -217,11 +223,7 @@ async def test_node_cancellation_on_other_node_exception() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt( - checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -259,59 +261,60 @@ async def test_dynamic_interrupt( "market": "US", } - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile(checkpointer=checkpointer) - # 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 ⛰️", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - interrupts=(Interrupt("Just because..."),), + 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 ⛰️", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + interrupts=(Interrupt("Just because..."),), + ), ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # TODO use aget_state_history + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # TODO use aget_state_history @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( - checkpointer_name: str, request: pytest.FixtureRequest + checkpointer_name: str, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): hello: str @@ -339,23 +342,24 @@ async def test_node_not_cancelled_on_other_node_interrupted( builder.add_node("bad", iambad) builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - graph = builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + thread = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} + assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} + assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} + assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} - assert not inner_task_cancelled - assert awhiles == 2 + assert not inner_task_cancelled + assert awhiles == 2 async def test_step_timeout_on_stream_hang() -> None: @@ -388,12 +392,8 @@ async def test_step_timeout_on_stream_hang() -> None: assert inner_task_cancelled -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_cancel_graph_astream( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) +async def test_cancel_graph_astream(checkpointer_name: str) -> None: class State(TypedDict): value: Annotated[int, operator.add] @@ -426,53 +426,47 @@ async def test_cancel_graph_astream( builder.add_edge(START, "alittlewhile") builder.add_edge(START, "aparallelwhile") builder.add_edge("alittlewhile", "awhile") - graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream - got_event = False - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - async with aclosing(graph.astream({"value": 1}, thread1)) as stream: - async for chunk in stream: - assert chunk == {"alittlewhile": {"value": 2}} - got_event = True - break + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - assert got_event + # test interrupting astream + got_event = False + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + async with aclosing(graph.astream({"value": 1}, thread1)) as stream: + async for chunk in stream: + assert chunk == {"alittlewhile": {"value": 2}} + got_event = True + break - # node aparallelwhile should start, but be cancelled - assert aparallelwhile.started is True - assert aparallelwhile.cancelled is True + assert got_event - # node "awhile" should never start - assert awhile.started is False + # node aparallelwhile should start, but be cancelled + assert aparallelwhile.started is True + assert aparallelwhile.cancelled is True - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ( - "aparallelwhile", - "alittlewhile", - ) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - } + # node "awhile" should never start + assert awhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ( + "aparallelwhile", + "alittlewhile", + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_cancel_graph_astream_events_v2( - request: pytest.FixtureRequest, checkpointer_name: Optional[str] -) -> None: - checkpointer = ( - request.getfixturevalue(f"checkpointer_{checkpointer_name}") - if checkpointer_name - else None - ) - +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC_PLUS_NONE) +async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) -> None: class State(TypedDict): value: int @@ -505,43 +499,45 @@ async def test_cancel_graph_astream_events_v2( builder.add_edge(START, "alittlewhile") builder.add_edge("alittlewhile", "awhile") builder.add_edge("awhile", "anotherwhile") - graph = builder.compile(checkpointer=checkpointer) - # test interrupting astream_events v2 - got_event = False - thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} - async with aclosing( - graph.astream_events({"value": 1}, thread2, version="v2") - ) as stream: - async for chunk in stream: - if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: - got_event = True - assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} - break + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - # did break - assert got_event + # test interrupting astream_events v2 + got_event = False + thread2: RunnableConfig = {"configurable": {"thread_id": "2"}} + async with aclosing( + graph.astream_events({"value": 1}, thread2, version="v2") + ) as stream: + async for chunk in stream: + if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]: + got_event = True + assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}} + break - # node "awhile" maybe starts (impl detail of astream_events) - # if it does start, it must be cancelled - if awhile.started: - assert awhile.cancelled is True + # did break + assert got_event - # node "anotherwhile" should never start - assert anotherwhile.started is False + # node "awhile" maybe starts (impl detail of astream_events) + # if it does start, it must be cancelled + if awhile.started: + assert awhile.cancelled is True - # checkpoint with output of "alittlewhile" should not be saved - if checkpointer is not None: - state = await graph.aget_state(thread2) - assert state is not None - assert state.values == {"value": 2} - assert state.next == ("awhile",) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"alittlewhile": {"value": 2}}, - } + # node "anotherwhile" should never start + assert anotherwhile.started is False + + # checkpoint with output of "alittlewhile" should not be saved + if checkpointer is not None: + state = await graph.aget_state(thread2) + assert state is not None + assert state.values == {"value": 2} + assert state.next == ("awhile",) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"alittlewhile": {"value": 2}}, + } async def test_node_schemas_custom_output() -> None: @@ -819,424 +815,438 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_two_processes_in_out_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} - # start execution, stop at inbox - assert await app.ainvoke(2, thread1) is None + # start execution, stop at inbox + assert await app.ainvoke(2, thread1) is None - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 + # inbox == 3 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None - # inbox == 21 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 + # inbox == 21 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, thread1) is None - assert await app.ainvoke(None, thread1) == 5 + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, thread1) is None + assert await app.ainvoke(None, thread1) == 5 - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) is None - # inbox == 21 - snapshot = await app.aget_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) - # update the state, resume - await app.aupdate_state(thread2, 25, as_node="one") - assert await app.ainvoke(None, thread2) == 26 + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 - # no pending tasks - snapshot = await app.aget_state(thread2) - assert snapshot.next == () + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () - # list history - history = [c async for c in app.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 6, "writes": {"two": 5}}, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 4, - "writes": {"input": 3}, - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 2, - "writes": {"input": 20}, - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 1, "writes": {"two": 4}}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - tasks=(PregelTask(AnyStr(), "two"),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"one": None}, - }, - created_at=AnyStr(), - parent_config=history[7].config, - ), - StateSnapshot( - values={"input": 2}, - tasks=(PregelTask(AnyStr(), "one"),), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"input": 2}, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 6, + "writes": {"two": 5}, + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 4, + "writes": {"input": 3}, + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 2, + "writes": {"input": 20}, + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"two": 4}, + }, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two"),), + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"one": None}, + }, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one"),), + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"input": 2}, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in app.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"two": {"output": 5}}, - ] - assert [ - c async for c in app.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"one": {"inbox": 4}}, - ] + # forking from any previous checkpoint should re-run nodes + assert [ + c async for c in app.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"two": {"output": 5}}, + ] + assert [ + c async for c in app.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"one": {"inbox": 4}}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_fork_always_re_runs_nodes( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + checkpointer_name: str, mocker: MockerFixture ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") add_one = mocker.Mock(side_effect=lambda _: 1) builder = StateGraph(Annotated[int, operator.add]) builder.add_node("add_one", add_one) builder.add_edge(START, "add_one") builder.add_conditional_edges("add_one", lambda cnt: "add_one" if cnt < 6 else END) - graph = builder.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + thread1 = {"configurable": {"thread_id": "1"}} - # start execution, stop at inbox - assert [ - c async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) - ] == [ - ("values", 1), - ("updates", {"add_one": 1}), - ("values", 2), - ("updates", {"add_one": 1}), - ("values", 3), - ("updates", {"add_one": 1}), - ("values", 4), - ("updates", {"add_one": 1}), - ("values", 5), - ("updates", {"add_one": 1}), - ("values", 6), - ] + # start execution, stop at inbox + assert [ + c + async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] - # list history - history = [c async for c in graph.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values=6, - next=(), - tasks=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[1].config, - ), - StateSnapshot( - values=5, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[2].config, - ), - StateSnapshot( - values=4, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[3].config, - ), - StateSnapshot( - values=3, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[4].config, - ), - StateSnapshot( - values=2, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"add_one": 1}, - }, - created_at=AnyStr(), - parent_config=history[5].config, - ), - StateSnapshot( - values=1, - tasks=(PregelTask(AnyStr(), "add_one"),), - next=("add_one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=history[6].config, - ), - StateSnapshot( - values=0, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": 1}, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + tasks=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"add_one": 1}, + }, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + tasks=(PregelTask(AnyStr(), "add_one"),), + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": 1}, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in graph.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in graph.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"add_one": 1}, - ] - assert [ - c async for c in graph.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"add_one": 1}, - {"add_one": 1}, - ] + # forking from any previous checkpoint should re-run nodes + assert [ + c + async for c in graph.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c + async for c in graph.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + ] + assert [ + c + async for c in graph.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: @@ -1508,7 +1518,8 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] -async def test_invoke_checkpoint(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1531,57 +1542,52 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaverAssertImmutable() + async with awith_checkpointer(checkpointer_name) as checkpointer: + 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, + retry_policy=RetryPolicy(), + ) - 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, - retry_policy=RetryPolicy(), - ) - - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 - assert errored_once, "errored and retried" - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = await memory.aget({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 2 + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + checkpoint = await checkpointer.aget({"configurable": {"thread_id": "2"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 5 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class State(TypedDict): value: Annotated[int, operator.add] @@ -1609,214 +1615,219 @@ async def test_pending_writes_resume( builder.add_node("two", two) builder.add_edge(START, "one") builder.add_edge(START, "two") - graph = builder.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke({"value": 1}, thread1) + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke({"value": 1}, thread1) - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 1 + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 1 - # latest checkpoint should be before nodes "one", "two" - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - assert state.tasks == ( - PregelTask(AnyStr(), "one"), - PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), - ) - assert state.metadata == { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - } - # should contain pending write of "one" - checkpoint = await checkpointer.aget_tuple(thread1) - assert checkpoint is not None - # should contain error from "two" - expected_writes = [ - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (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) - # both non-error pending writes come from same task - non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - assert non_error_writes[0][0] == non_error_writes[1][0] - # error write is from the other task - error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) - assert error_write[0] != non_error_writes[0][0] + # latest checkpoint should be before nodes "one", "two" + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + assert state.tasks == ( + PregelTask(AnyStr(), "one"), + PregelTask(AnyStr(), "two", 'ValueError("I\'m not good")'), + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + } + # should contain pending write of "one" + checkpoint = await checkpointer.aget_tuple(thread1) + assert checkpoint is not None + # should contain error from "two" + expected_writes = [ + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (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) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + assert non_error_writes[0][0] == non_error_writes[1][0] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] - # TODO arguably this shouldn't even run the failed task again, - # and should require empty update_state (ie new checkpoint_id) - # in order to try again + # TODO arguably this shouldn't even run the failed task again, + # and should require empty update_state (ie new checkpoint_id) + # in order to try again - # resume execution - with pytest.raises(ValueError, match="I'm not good"): - await graph.ainvoke(None, thread1) + # resume execution + with pytest.raises(ValueError, match="I'm not good"): + await graph.ainvoke(None, thread1) - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 2 + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 2 - # confirm no new checkpoints saved - state_two = await graph.aget_state(thread1) - assert state_two.metadata == state.metadata + # confirm no new checkpoints saved + state_two = await graph.aget_state(thread1) + assert state_two.metadata == state.metadata - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke(None, thread1) == {"value": 6} + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert await graph.ainvoke(None, thread1) == {"value": 6} - # check all final checkpoints - checkpoints = [c async for c in checkpointer.alist(thread1)] - # we should have 3 - assert len(checkpoints) == 3 - # the last one not too interesting for this test - assert checkpoints[0] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "one": { - "start:one": AnyVersion(), + # check all final checkpoints + checkpoints = [c async for c in checkpointer.alist(thread1)] + # we should have 3 + assert len(checkpoints) == 3 + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "one": { + "start:one": AnyVersion(), + }, + "two": { + "start:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, }, - "two": { - "start:two": AnyVersion(), - }, - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - "__interrupt__": { + "channel_versions": { + "one": AnyVersion(), + "two": AnyVersion(), "value": AnyVersion(), "__start__": AnyVersion(), "start:one": AnyVersion(), "start:two": AnyVersion(), }, + "channel_values": {"one": "one", "two": "two", "value": 6}, }, - "channel_versions": { - "one": AnyVersion(), - "two": AnyVersion(), - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + metadata={ + "parents": {}, + "step": 1, + "source": "loop", + "writes": {"one": {"value": 2}, "two": {"value": 3}}, }, - "channel_values": {"one": "one", "two": "two", "value": 6}, - }, - metadata={ - "parents": {}, - "step": 1, - "source": "loop", - "writes": {"one": {"value": 2}, "two": {"value": 3}}, - }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], - } - }, - pending_writes=[], - ) - # the previous one we assert that pending writes contains both - # - original error - # - successful writes from resuming after preventing error - assert checkpoints[1] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "__input__": {}, - "__start__": { + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"][ + "checkpoint_id" + ], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), "__start__": AnyVersion(), + "start:one": AnyVersion(), + "start:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "start:one": "__start__", + "start:two": "__start__", }, }, - "channel_versions": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "start:one": AnyVersion(), - "start:two": AnyVersion(), + metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"][ + "checkpoint_id" + ], + } }, - "channel_values": { - "value": 1, - "start:one": "__start__", - "start:two": "__start__", + pending_writes=UnsortedSequence( + (AnyStr(), "one", "one"), + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ValueError("I\'m not good")'), + (AnyStr(), "two", "two"), + (AnyStr(), "value", 3), + ), + ) + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - }, - metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"], - } - }, - pending_writes=UnsortedSequence( - (AnyStr(), "one", "one"), - (AnyStr(), "value", 2), - (AnyStr(), "__error__", 'ValueError("I\'m not good")'), - (AnyStr(), "two", "two"), - (AnyStr(), "value", 3), - ), - ) - assert checkpoints[2] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 1, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": {"__input__": {}}, - "channel_versions": { - "__start__": AnyVersion(), + checkpoint={ + "v": 1, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), + }, + "channel_values": {"__start__": {"value": 1}}, }, - "channel_values": {"__start__": {"value": 1}}, - }, - metadata={ - "parents": {}, - "step": -1, - "source": "input", - "writes": {"__start__": {"value": 1}}, - }, - parent_config=None, - pending_writes=UnsortedSequence( - (AnyStr(), "value", 1), - (AnyStr(), "start:one", "__start__"), - (AnyStr(), "start:two", "__start__"), - ), - ) + metadata={ + "parents": {}, + "step": -1, + "source": "input", + "writes": {"__start__": {"value": 1}}, + }, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), + ) async def test_cond_edge_after_send() -> None: @@ -1848,9 +1859,8 @@ async def test_cond_edge_after_send() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_checkpoint_three( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str + mocker: MockerFixture, 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: @@ -1865,118 +1875,121 @@ async def test_invoke_checkpoint_three( | raise_if_above_10 ) - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=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 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 == () - - 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 + async with awith_checkpointer(checkpointer_name) as checkpointer: + 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, ) - ] - 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 = {"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 == () + + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config + ) + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await 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 + ) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -2189,9 +2202,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_conditional_graph(checkpointer_name: str) -> None: from copy import deepcopy from langchain_core.agents import AgentAction, AgentFinish @@ -2200,8 +2211,6 @@ async def test_conditional_graph( 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: @@ -2458,82 +2467,75 @@ async def test_conditional_graph( }, ] - # test state get/update methods with interrupt_after + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] + config = {"configurable": {"thread_id": "1"}} - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2541,18 +2543,10 @@ async def test_conditional_graph( ), "input": "what is weather in sf", }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -2560,86 +2554,85 @@ async def test_conditional_graph( log="tool:search_api:a different query", ), "input": "what is weather in sf", - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "input": "what is weather in sf", + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -2656,18 +2649,10 @@ async def test_conditional_graph( log="finish:a really nice answer", ), }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -2684,91 +2669,112 @@ async def test_conditional_graph( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), } - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - ) + # test state get/update methods with interrupt_before - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, + { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2776,18 +2782,10 @@ async def test_conditional_graph( ), "input": "what is weather in sf", }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -2795,86 +2793,85 @@ async def test_conditional_graph( log="tool:search_api:a different query", ), "input": "what is weather in sf", - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "input": "what is weather in sf", + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - await app_w_interrupt.aupdate_state( - config, - { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - }, - ) + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { + await app_w_interrupt.aupdate_state( + config, + { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -2891,18 +2888,10 @@ async def test_conditional_graph( log="finish:a really nice answer", ), }, - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -2919,196 +2908,234 @@ async def test_conditional_graph( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - } + }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test re-invoke to continue with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "3"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], ], - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", + } + }, + { + "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", + ], ], - [ - 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" - ), - } - }, - ] + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] -async def test_conditional_graph_state(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_conditional_graph_state( + mocker: MockerFixture, checkpointer_name: str +) -> None: from langchain_core.agents import AgentAction, AgentFinish from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate @@ -3345,15 +3372,201 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - # test state get/update methods with interrupt_after + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + async with assert_ctx_once(): + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + async with assert_ctx_once(): + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + async with assert_ctx_once(): + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm - async with assert_ctx_once(): assert [ c async for c in app_w_interrupt.astream( @@ -3371,41 +3584,38 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - async with assert_ctx_once(): await app_w_interrupt.aupdate_state( config, { @@ -3417,41 +3627,40 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - async with assert_ctx_once(): assert [c async for c in app_w_interrupt.astream(None, config)] == [ { "tools": { @@ -3478,7 +3687,6 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ] - async with assert_ctx_once(): await app_w_interrupt.aupdate_state( config, { @@ -3489,151 +3697,12 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: }, ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", ), - } - }, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { "intermediate_steps": [ [ AgentAction( @@ -3644,69 +3713,30 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: "result for query", ] ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + tasks=(), + next=(), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) async def test_conditional_entrypoint_graph() -> None: @@ -3983,9 +4013,7 @@ class ToolInput(BaseModel, arbitrary_types_allowed=True): @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_state_graph_packets( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_state_graph_packets(checkpointer_name: str) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -3997,8 +4025,6 @@ async def test_state_graph_packets( ) from langchain_core.tools import tool - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] @@ -4220,63 +4246,20 @@ async def test_state_graph_packets( {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - ] + config = {"configurable": {"thread_id": "1"}} - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ - "ts" - ], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { "agent": { "messages": AIMessage( id="ai1", @@ -4291,47 +4274,68 @@ async def test_state_graph_packets( ) } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=( + await app_w_interrupt.checkpointer.aget_tuple(config) + ).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( id="ai1", content="", tool_calls=[ @@ -4341,97 +4345,50 @@ async def test_state_graph_packets( "args": {"query": "a different query"}, }, ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", ) } }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ) - } - }, - { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ) - }, - }, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { + { "agent": { "messages": AIMessage( id="ai2", @@ -4451,63 +4408,133 @@ async def test_state_graph_packets( ) }, }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - ToolMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": {"messages": AIMessage(content="answer", id="ai2")} + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_message_graph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +async def test_message_graph(checkpointer_name: str) -> None: from langchain_core.agents import AgentAction from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -4515,8 +4542,6 @@ async def test_message_graph( from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage from langchain_core.tools import tool - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -4679,50 +4704,20 @@ async def test_message_graph( {"agent": AIMessage(content="answer", id="ai3")}, ] - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - HumanMessage(content="what is weather in sf"), config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["agent"], ) - ] == [ - { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ) - }, - ] + config = {"configurable": {"thread_id": "1"}} - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { "agent": AIMessage( content="", additional_kwargs={ @@ -4731,43 +4726,59 @@ async def test_message_graph( id="ai1", ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values[-1] - last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"' - await app_w_interrupt.aupdate_state(config, last_message) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"query"', + } + }, + id="ai1", + ) }, - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.additional_kwargs["function_call"]["arguments"] = ( + '"a different query"' + ) + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( content="", additional_kwargs={ "function_call": { @@ -4776,69 +4787,43 @@ async def test_message_graph( } }, id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ) + }, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ) - }, - { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ) - }, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools"),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { + { "agent": AIMessage( content="", additional_kwargs={ @@ -4850,53 +4835,105 @@ async def test_message_graph( id="ai2", ) }, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ] - await app_w_interrupt.aupdate_state( - config, - AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools"),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ) }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) + + await app_w_interrupt.aupdate_state( + config, AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - ) + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + ) async def test_in_one_fan_out_out_one_graph_state() -> None: @@ -5126,11 +5163,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_start_branch_then( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_start_branch_then(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5175,172 +5208,183 @@ async def test_start_branch_then( "market": "US", } - tool_two = tool_two_graph.compile( - store=MemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + async with awith_checkpointer(checkpointer_name) as checkpointer: + 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)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"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={"parents": {}, "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={ - "parents": {}, - "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)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"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={"parents": {}, "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={ + "parents": {}, + "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={"parents": {}, "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={ - "parents": {}, - "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={"parents": {}, "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={ + "parents": {}, + "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={"parents": {}, "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={ - "parents": {}, - "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={ - "parents": {}, - "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={"parents": {}, "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={ + "parents": {}, + "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={ + "parents": {}, + "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, + ) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_branch_then( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_branch_then(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5368,600 +5412,629 @@ async def test_branch_then( "market": "US", } - # 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" + async with awith_checkpointer(checkpointer_name) as checkpointer: + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + }, + "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": { + "parents": {}, + "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": { + "parents": {}, + "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": { + "parents": {}, + "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": { + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + }, + "next": [], + "tasks": [], + }, + }, + ] + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - }, - "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": { - "parents": {}, - "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": { - "parents": {}, - "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": { - "parents": {}, - "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": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - "next": [], - "tasks": [], - }, - }, - ] - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] - ) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - # 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" + thread1 = {"configurable": {"thread_id": "11"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + }, + "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + }, + "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ["start:prepare"], + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": "1a591be4-f85c-558f-8d00-1ccac0d1877f", + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + }, + "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={ + "parents": {}, + "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, ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - }, - "next": ["__start__"], - "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}], + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - }, - "next": ["prepare"], - "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}], + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + + thread2 = {"configurable": {"thread_id": "12"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), + next=("tool_two_fast",), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, }, - }, - { - "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(thread2, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread2)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, }, - }, - { - "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(thread2, limit=2) + ][-1].config, + ) + + tool_two = tool_two_graph.compile( + checkpointer=checkpointer, interrupt_after=["prepare"] + ) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), + next=("tool_two_slow",), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - "next": ["tool_two_slow"], - "tasks": [{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}], + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config=(await tool_two.checkpointer.aget_tuple(thread1)).config, + created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, }, - }, - ] - 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={ - "parents": {}, - "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={ - "parents": {}, - "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, - ) + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) - thread2 = {"configurable": {"thread_id": "12"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - 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={ + "parents": {}, + "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={ + "parents": {}, + "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, + ) - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) - - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow"),), - next=("tool_two_slow",), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread1)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, - ) - - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast"),), - next=("tool_two_fast",), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config=(await tool_two.checkpointer.aget_tuple(thread2)).config, - created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config, - ) - - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state(thread3, {"my_key": "key", "market": "DE"}) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare"),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "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={ - "parents": {}, - "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={ - "parents": {}, - "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={ + "parents": {}, + "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={ + "parents": {}, + "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={ + "parents": {}, + "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: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -6023,31 +6096,33 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, checkpointer_name: str ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -6114,31 +6189,33 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( - snapshot: SnapshotAssertion, mocker: MockerFixture + snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str ) -> None: from langchain_core.pydantic_v1 import BaseModel, ValidationError @@ -6256,75 +6333,77 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + async with assert_ctx_once(): + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + async with assert_ctx_once(): + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "step": 4, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + + async with assert_ctx_once(): + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - "step": 4, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - - async with assert_ctx_once(): - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, + snapshot: SnapshotAssertion, checkpointer_name: str ) -> None: from pydantic import BaseModel, ValidationError @@ -6419,40 +6498,44 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", + } } - } -async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( + checkpointer_name: str, +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -6520,28 +6603,29 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - app_w_interrupt = workflow.compile( - checkpointer=MemorySaverAssertImmutable(), - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config + async with awith_checkpointer(checkpointer_name) as checkpointer: + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_after=["retriever_one"], ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - ] + config = {"configurable": {"thread_id": "1"}} - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: @@ -6839,11 +6923,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -6887,89 +6967,91 @@ async def test_nested_graph_interrupts_parallel( graph.add_edge(["inner", "outer_1"], "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { - "my_key": "", - } + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + "my_key": "", + } - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "got here and there and parallel and back again", - } + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "got here and there and parallel and back again", + } - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": ""}, config, subgraphs=True)] == [ - # we got to parallel node first - ((), {"outer_1": {"my_key": " and parallel"}}), - ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), - ] - assert [c async for c in app.astream(None, config)] == [ - {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, subgraphs=True) + ] == [ + # we got to parallel node first + ((), {"outer_1": {"my_key": " and parallel"}}), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ] + assert [c async for c in app.astream(None, config)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] - # # test interrupts BEFORE the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] - # test interrupts AFTER the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [ - {"my_key": ""}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + ] == [ + {"my_key": ""}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: class State(TypedDict): my_key: str @@ -7017,50 +7099,49 @@ async def test_doubly_nested_graph_interrupts( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value", - } + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + "my_key": "hi my value", + } - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } + assert await app.ainvoke(None, config, debug=True) == { + "my_key": "hi my value here and there and back again", + } - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"parent_1": {"my_key": "hi my value"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, - ] + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + {"parent_1": {"my_key": "hi my value"}}, + ] + assert [c async for c in app.astream(None, config)] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - {"my_key": "my value"}, - {"my_key": "hi my value"}, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - {"my_key": "hi my value"}, - {"my_key": "hi my value here and there"}, - {"my_key": "hi my value here and there and back again"}, - ] + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, stream_mode="values" + ) + ] == [ + {"my_key": "my value"}, + {"my_key": "hi my value"}, + ] + assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + {"my_key": "hi my value"}, + {"my_key": "hi my value here and there"}, + {"my_key": "hi my value here and there and back again"}, + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_nested_graph_state(checkpointer_name: str) -> None: class InnerState(TypedDict): my_key: str my_other_key: str @@ -7106,388 +7187,14 @@ async def test_nested_graph_state( graph.add_edge("inner", "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # now, get_state with subgraphs state - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - name="inner_2", - error=None, - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": { - "": AnyStr(), - }, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - assert history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # get_state_history for a subgraph returns its checkpoints - child_history = [c async for c in app.aget_state_history(history[0].tasks[0].state)] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here", "my_other_key": "hi my value"}, - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="inner_2"),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="inner_1"),), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] - - # resume - await app.ainvoke(None, config, debug=True) - # test state w/ nested subgraph state (after resuming from interrupt) - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # test full history at the end - actual_history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( @@ -7520,11 +7227,62 @@ async def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), + ) + # now, get_state with subgraphs state + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state=StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": { + "": AnyStr(), + }, + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("inner",), config={ "configurable": { "thread_id": "1", @@ -7535,8 +7293,8 @@ async def test_nested_graph_state( metadata={ "parents": {}, "source": "loop", - "writes": None, - "step": 0, + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, }, created_at=AnyStr(), parent_config={ @@ -7546,11 +7304,194 @@ async def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), + ) + # get_state_history returns outer graph checkpoints + history = [c async for c in app.aget_state_history(config)] + assert history == [ + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + # get_state_history for a subgraph returns its checkpoints + child_history = [ + c async for c in app.aget_state_history(history[0].tasks[0].state) + ] + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="inner_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="inner_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ] + + # resume + await app.ainvoke(None, config, debug=True) + # test state w/ nested subgraph state (after resuming from interrupt) + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), config={ "configurable": { "thread_id": "1", @@ -7560,26 +7501,172 @@ async def test_nested_graph_state( }, metadata={ "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, }, created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - # test looking up parent state by checkpoint ID - for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): - assert await app.aget_state(actual_snapshot.config) == expected_snapshot + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # test full history at the end + actual_history = [c async for c in app.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "outer_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert await app.aget_state(actual_snapshot.config) == expected_snapshot @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: class State(TypedDict): my_key: str @@ -7627,277 +7714,37 @@ async def test_doubly_nested_graph_state( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert [ - c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) - ] == [ - ((), {"parent_1": {"my_key": "hi my value"}}), - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ] - # get state without subgraphs - outer_state = await app.aget_state(config) - assert outer_state == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert [ + c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) + ] == [ + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - child_state = await app.aget_state(outer_state.tasks[0].state) - assert ( - child_state.tasks[0] - == StateSnapshot( + ] + # get state without subgraphs + outer_state = await app.aget_state(config) + assert outer_state == StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( AnyStr(), - "child_1", + "child", state={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": AnyStr("child"), } }, ), ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ).tasks[0] - ) - grandchild_state = await app.aget_state(child_state.tasks[0].state) - assert grandchild_state == StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ) - # get state with subgraphs - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - state=StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - state=StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - ), - ), - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile(r"child:.+|child1:") - ): AnyStr(), - } - ), - } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": { - "grandchild_1": {"my_key": "hi my value here"} - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - # resume - assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_2": {"my_key": "hi my value here and there"}}, - ), - ((AnyStr("child:"),), {"child_1": {"my_key": "hi my value here and there"}}), - ((), {"child": {"my_key": "hi my value here and there"}}), - ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), - ] - # get state with and without subgraphs - assert ( - await app.aget_state(config) - == await app.aget_state(config, subgraphs=True) - == StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), + next=("child",), config={ "configurable": { "thread_id": "1", @@ -7908,10 +7755,8 @@ async def test_doubly_nested_graph_state( metadata={ "parents": {}, "source": "loop", - "writes": { - "parent_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, }, created_at=AnyStr(), parent_config={ @@ -7922,13 +7767,227 @@ async def test_doubly_nested_graph_state( } }, ) - ) - # get outer graph history - outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - StateSnapshot( + child_state = await app.aget_state(outer_state.tasks[0].state) + assert ( + child_state.tasks[0] + == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + } + }, + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ).tasks[0] + ) + grandchild_state = await app.aget_state(child_state.tasks[0].state) + assert grandchild_state == StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + ) + # get state with subgraphs + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + state=StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + state=StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "writes": { + "grandchild_1": { + "my_key": "hi my value here" + } + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, + ), + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + # resume + assert [c async for c in app.astream(None, config, subgraphs=True)] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ( + (AnyStr("child:"),), + {"child_1": {"my_key": "hi my value here and there"}}, + ), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ] + # get state with and without subgraphs + assert ( + await app.aget_state(config) + == await app.aget_state(config, subgraphs=True) + == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, tasks=(), next=(), @@ -7957,95 +8016,227 @@ async def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - ), + ) + ) + # get outer graph history + outer_history = [c async for c in app.aget_state_history(config)] + assert ( + outer_history[0] + == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "parent_2": { + "my_key": "hi my value here and there and back again" + } + }, + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("parent_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("parent_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ][0] + ) + # get child graph history + child_history = [ + c async for c in app.aget_state_history(outer_history[2].tasks[0].state) + ] + assert child_history == [ StateSnapshot( values={"my_key": "hi my value here and there"}, - next=("parent_2",), + next=(), config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ - "parents": {}, "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + "parents": {"": AnyStr()}, }, created_at=AnyStr(), parent_config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), } }, - tasks=(PregelTask(id=AnyStr(), name="parent_2"),), + tasks=(), ), StateSnapshot( values={"my_key": "hi my value"}, + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + } + }, tasks=( PregelTask( - AnyStr(), - "child", + id=AnyStr(), + name="child_1", state={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child"), + "checkpoint_ns": AnyStr("child:"), } }, ), ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="parent_1"),), ), StateSnapshot( values={}, @@ -8053,290 +8244,195 @@ async def test_doubly_nested_graph_state( config={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", + "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ - "parents": {}, "source": "input", - "writes": {"my_key": "my value"}, + "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, + "parents": {"": AnyStr()}, }, created_at=AnyStr(), parent_config=None, tasks=(PregelTask(id=AnyStr(), name="__start__"),), ), - ][0] - ) - # get child graph history - child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) - ] - assert child_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="child_1", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - } + ] + # get grandchild graph history + grandchild_history = [ + c async for c in app.aget_state_history(child_history[1].tasks[0].state) + ] + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_2": {"my_key": "hi my value here and there"} }, - ), + "step": 2, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + tasks=(), ), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] - # get grandchild graph history - grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) - ] - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + StateSnapshot( + values={"my_key": "hi my value here"}, + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_2": {"my_key": "hi my value here and there"}}, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(), - ), - StateSnapshot( - values={"my_key": "hi my value here"}, - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + }, + tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("grandchild_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="grandchild_2"),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( + }, + tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": AnyDict( { "": AnyStr(), AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="grandchild_1"),), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ] - # replay grandchild checkpoint - assert [ - c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ) - ] + # replay grandchild checkpoint + assert [ + c + async for c in app.astream( + None, grandchild_history[2].config, subgraphs=True + ) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ) + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_to_nested_graphs( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +async def test_send_to_nested_graphs(checkpointer_name: str) -> None: class OverallState(TypedDict): subjects: list[str] jokes: Annotated[list[str], operator.add] @@ -8370,112 +8466,74 @@ async def test_send_to_nested_graphs( builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} - # invoke and pause at nested interrupt - assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": [], - } - # check state - outer_state = await graph.aget_state(config) + # invoke and pause at nested interrupt + assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + # check state + outer_state = await graph.aget_state(config) - assert outer_state == StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, + assert outer_state == StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), ), - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) - # update state of dogs joke graph - await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} - ) + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + ) - # continue past interrupt - assert await graph.ainvoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } - - actual_snapshot = await graph.aget_state(config) - expected_snapshot = StateSnapshot( - values={ + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert actual_snapshot == expected_snapshot + } - # test full history - actual_history = [c async for c in graph.aget_state_history(config)] - expected_history = [ - StateSnapshot( + actual_snapshot = await graph.aget_state(config) + expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], @@ -8508,53 +8566,247 @@ async def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - ), - StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - next=("generate_joke", "generate_joke"), - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [c async for c in graph.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": [ + "Joke about cats - hohoho", + "Joke about turtles - hohoho", + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about turtles - hohoho"]}, + ] }, - ), - PregelTask( - AnyStr(), - "generate_joke", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, ), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + tasks=( + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + PregelTask( + AnyStr(), + "generate_joke", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + ), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"jokes": []}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + ), + ] + assert actual_history == expected_history + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_weather_subgraph( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import AIMessage, HumanMessage, ToolCall + from langchain_core.tools import tool + + from langgraph.graph import MessagesState + + # setup subgraph + + @tool + def get_weather(city: str): + """Get the weather for a specific city""" + return f"I'ts sunny in {city}!" + + weather_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ) + ] + ) + + class SubGraphState(MessagesState): + city: str + + def model_node(state: SubGraphState): + result = weather_model.invoke(state["messages"]) + return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} + + def weather_node(state: SubGraphState): + result = get_weather.invoke({"city": state["city"]}) + return {"messages": [{"role": "assistant", "content": result}]} + + subgraph = StateGraph(SubGraphState) + subgraph.add_node(model_node) + subgraph.add_node(weather_node) + subgraph.add_edge(START, "model_node") + subgraph.add_edge("model_node", "weather_node") + subgraph.add_edge("weather_node", END) + subgraph = subgraph.compile(interrupt_before=["weather_node"]) + + # setup main graph + + class RouterState(MessagesState): + route: Literal["weather", "other"] + + class Router(TypedDict): + route: Literal["weather", "other"] + + router_model = FakeMessagesListChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ) + ] + ) + + def router_node(state: RouterState): + system_message = "Classify the incoming query as either about weather or not." + messages = [{"role": "system", "content": system_message}] + state["messages"] + route = router_model.invoke(messages) + return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]} + + def normal_llm_node(state: RouterState): + return {"messages": [AIMessage("Hello!")]} + + def route_after_prediction(state: RouterState): + if state["route"] == "weather": + return "weather_graph" + else: + return "normal_llm_node" + + async def weather_graph(state: RouterState): + return await subgraph.ainvoke(state) + + graph = StateGraph(RouterState) + graph.add_node(router_node) + graph.add_node(normal_llm_node) + graph.add_node("weather_graph", weather_graph) + graph.add_edge(START, "router_node") + graph.add_conditional_edges("router_node", route_after_prediction) + graph.add_edge("normal_llm_node", END) + graph.add_edge("weather_graph", END) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = graph.compile(checkpointer=checkpointer) + + assert graph.get_graph(xray=1).draw_mermaid() == snapshot + + config = {"configurable": {"thread_id": "1"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + + # run until interrupt + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ] + + # check current state + state = await graph.aget_state(config) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), + next=("weather_graph",), config={ "configurable": { "thread_id": "1", @@ -8563,16 +8815,264 @@ async def test_send_to_nested_graphs( } }, metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, }, created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("weather_graph:"), + } + }, + ), + ), + ) + + # update + await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) + + # run after update + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (AnyStr("weather_graph:"),), + { + "weather_node": { + "messages": [ + {"role": "assistant", "content": "I'ts sunny in la!"} + ] + } + }, + ), + ( + (), + { + "weather_graph": { + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="I'ts sunny in la!", id=AnyStr()), + ] + } + }, + ), + ] + + # try updating acting as weather node + config = {"configurable": {"thread_id": "14"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ] + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state=StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ) + ], + "city": "San Francisco", + }, + next=("weather_node",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"model_node": {"city": "San Francisco"}}, + "step": 1, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="weather_node"),), + ), + ), + ), + ) + await graph.aupdate_state( + state.tasks[0].state.config, + {"messages": [{"role": "assistant", "content": "rainy"}]}, + as_node="weather_node", + ) + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [ + HumanMessage(content="what's the weather in sf", id=AnyStr()) + ], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + state=StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="rainy", id=AnyStr()), + ], + "city": "San Francisco", + }, + next=(), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "source": "update", + "step": 2, + "writes": { + "weather_node": { + "messages": [ + {"role": "assistant", "content": "rainy"} + ] + } + }, + "parents": {"": AnyStr()}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + ), + ), + ) + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (), + { + "weather_graph": { + "messages": [ + HumanMessage( + content="what's the weather in sf", id=AnyStr() + ), + AIMessage(content="rainy", id=AnyStr()), + ] + } + }, + ), + ] async def test_checkpoint_metadata() -> None: @@ -8731,3 +9231,26 @@ async def test_checkpoint_metadata() -> None: assert chkpnt_tuple.metadata["thread_id"] == "2" assert chkpnt_tuple.metadata["test_config_3"] == "foo" assert chkpnt_tuple.metadata["test_config_4"] == "bar" + + +async def test_checkpointer_null_pending_writes() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return [self.name] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_edge(START, "1") + graph = builder.compile(checkpointer=MemorySaverNoPending()) + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] + assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 3 + assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [ + "1" + ] * 4 diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 4088e38c9..392a3f22a 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,3 +1,4 @@ +import warnings from typing import Annotated as Annotated2 from typing import Any @@ -45,7 +46,8 @@ def test_warns_invalid_schema(schema: Any): ) def test_doesnt_warn_valid_schema(schema: Any): # Assert the function does not raise a warning - with pytest.warns(None): + with warnings.catch_warnings(): + warnings.simplefilter("error") _warn_invalid_state_schema(schema) diff --git a/libs/langgraph/tests/test_store.py b/libs/langgraph/tests/test_store.py index cd53ee407..71494adaf 100644 --- a/libs/langgraph/tests/test_store.py +++ b/libs/langgraph/tests/test_store.py @@ -1,11 +1,14 @@ import asyncio from typing import Any, Optional +import pytest from pytest_mock import MockerFixture from langgraph.store.base import BaseStore from langgraph.store.batch import AsyncBatchedStore +pytestmark = pytest.mark.anyio + async def test_async_batch_store(mocker: MockerFixture) -> None: aget = mocker.stub() diff --git a/libs/langgraph/tests/test_utils.py b/libs/langgraph/tests/test_utils.py index 1e3bf8ea1..3a914c23b 100644 --- a/libs/langgraph/tests/test_utils.py +++ b/libs/langgraph/tests/test_utils.py @@ -11,6 +11,8 @@ from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.utils import is_async_callable, is_async_generator +pytestmark = pytest.mark.anyio + def test_is_async() -> None: async def func() -> None: From ddeb9d73142f14034f948369119273cd468398e7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:42:37 -0700 Subject: [PATCH 10/27] Disable test in old versions of py --- libs/langgraph/tests/test_pregel_async.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4d004075d..d3a6c422f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2,6 +2,7 @@ import asyncio import json import operator import re +import sys from collections import Counter from contextlib import asynccontextmanager, contextmanager from typing import ( @@ -8675,6 +8676,10 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: assert actual_history == expected_history +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_weather_subgraph( checkpointer_name: str, snapshot: SnapshotAssertion From 0cc37992986359d83541fbc1df6c46e6bd413424 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:43:34 -0700 Subject: [PATCH 11/27] checkpoint: Implement sync methods in async checkpointers - These are needed when subgraphs are invoked synchronously in nodes --- .../langgraph/checkpoint/postgres/aio.py | 96 ++++++++++++++++++- .../langgraph/checkpoint/sqlite/aio.py | 88 +++++++++++------ libs/langgraph/tests/test_pregel_async.py | 11 ++- 3 files changed, 165 insertions(+), 30 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 63a2038a3..a931181e1 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,6 +1,6 @@ import asyncio from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Optional, Union +from typing import Any, AsyncIterator, Iterator, List, Optional, Union from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline @@ -51,6 +51,7 @@ class AsyncPostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() @classmethod @asynccontextmanager @@ -329,3 +330,96 @@ class AsyncPostgresSaver(BasePostgresSaver): binary=True, row_factory=dict_row ) as cur: yield cur + + def list( + self, + config: RunnableConfig | None, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + """List checkpoints from the database. + + This method retrieves a list of checkpoint tuples from the Postgres database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. + """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break + + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: + """Get a checkpoint tuple from the database. + + This method retrieves a checkpoint tuple from the Postgres database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() + + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + new_versions: ChannelVersions, + ) -> RunnableConfig: + """Save a checkpoint to the database. + + This method saves a checkpoint to the Postgres database. The checkpoint is associated + with the provided config and its parent config (if any). + + Args: + config (RunnableConfig): The config to associate with the checkpoint. + checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, + config: RunnableConfig, + writes: List[tuple[str, Any]], + task_id: str, + ) -> None: + """Store intermediate writes linked to a checkpoint. + + This method saves intermediate writes associated with a checkpoint to the database. + + Args: + config (RunnableConfig): Configuration of the related checkpoint. + writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. + task_id (str): Identifier for the task creating the writes. + """ + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index dc68a0403..7cdc7c8fb 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -1,11 +1,11 @@ import asyncio -import functools from contextlib import asynccontextmanager from typing import ( Any, AsyncIterator, Dict, Iterator, + List, Optional, Sequence, Tuple, @@ -31,20 +31,6 @@ from langgraph.checkpoint.sqlite.utils import search_where T = TypeVar("T", bound=callable) -def not_implemented_sync_method(func: T) -> T: - @functools.wraps(func) - def wrapper(*args, **kwargs): - raise NotImplementedError( - "The AsyncSqliteSaver does not support synchronous methods. " - "Consider using the SqliteSaver instead.\n" - "from langgraph.checkpoint.sqlite import SqliteSaver\n" - "See https://langchain-ai.github.io/langgraph/reference/checkpoints/langgraph.checkpoint.sqlite.SqliteSaver " - "for more information." - ) - - return wrapper - - class AsyncSqliteSaver(BaseCheckpointSaver): """An asynchronous checkpoint saver that stores checkpoints in a SQLite database. @@ -132,6 +118,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): self.jsonplus_serde = JsonPlusSerializer() self.conn = conn self.lock = asyncio.Lock() + self.loop = asyncio.get_running_loop() self.is_setup = False @classmethod @@ -150,16 +137,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver): async with aiosqlite.connect(conn_string) as conn: yield AsyncSqliteSaver(conn) - @not_implemented_sync_method def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `aget` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. - """ + This method retrieves a checkpoint tuple from the SQLite database based on the + provided config. If the config contains a "checkpoint_id" key, the checkpoint with + the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint + for the given thread ID is retrieved. + + Args: + config (RunnableConfig): The config to use for retrieving the checkpoint. + + Returns: + Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. + """ + return asyncio.run_coroutine_threadsafe( + self.aget_tuple(config), self.loop + ).result() - @not_implemented_sync_method def list( self, config: Optional[RunnableConfig], @@ -168,21 +163,60 @@ class AsyncSqliteSaver(BaseCheckpointSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: - """List checkpoints from the database. + """List checkpoints from the database asynchronously. - Note: - This method is not implemented for the AsyncSqliteSaver. Use `alist` instead. - Or consider using the [SqliteSaver][sqlitesaver] checkpointer. + This method retrieves a list of checkpoint tuples from the SQLite database based + on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first). + + Args: + config (Optional[RunnableConfig]): Base configuration for filtering checkpoints. + filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None. + limit (Optional[int]): Maximum number of checkpoints to return. + + Yields: + Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ + aiter_ = self.alist(config, filter=filter, before=before, limit=limit) + while True: + try: + yield asyncio.run_coroutine_threadsafe( + anext(aiter_), self.loop + ).result() + except StopAsyncIteration: + break - @not_implemented_sync_method def put( self, config: RunnableConfig, checkpoint: Checkpoint, metadata: CheckpointMetadata, + new_versions: ChannelVersions, ) -> RunnableConfig: - """Save a checkpoint to the database. FOO""" + """Save a checkpoint to the database. + + This method saves a checkpoint to the SQLite database. The checkpoint is associated + with the provided config and its parent config (if any). + + Args: + config (RunnableConfig): The config to associate with the checkpoint. + checkpoint (Checkpoint): The checkpoint to save. + metadata (CheckpointMetadata): Additional metadata to save with the checkpoint. + new_versions (ChannelVersions): New channel versions as of this write. + + Returns: + RunnableConfig: Updated configuration after storing the checkpoint. + """ + return asyncio.run_coroutine_threadsafe( + self.aput(config, checkpoint, metadata, new_versions), self.loop + ).result() + + def put_writes( + self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str + ) -> None: + return asyncio.run_coroutine_threadsafe( + self.aput_writes(config, writes, task_id), self.loop + ).result() async def setup(self) -> None: """Set up the checkpoint database asynchronously. diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d3a6c422f..7901ad2bc 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8771,8 +8771,10 @@ async def test_weather_subgraph( else: return "normal_llm_node" - async def weather_graph(state: RouterState): - return await subgraph.ainvoke(state) + def weather_graph(state: RouterState): + # this tests that all async checkpointers tested also implement sync methods + # as the subgraph called with sync invoke will use sync checkpointer methods + return subgraph.invoke(state) graph = StateGraph(RouterState) graph.add_node(router_node) @@ -8783,6 +8785,9 @@ async def test_weather_subgraph( graph.add_edge("normal_llm_node", END) graph.add_edge("weather_graph", END) + def get_first_in_list(): + return [*graph.get_state_history(config, limit=1)][0] + async with awith_checkpointer(checkpointer_name) as checkpointer: graph = graph.compile(checkpointer=checkpointer) @@ -8846,6 +8851,8 @@ async def test_weather_subgraph( ), ), ) + # confirm that list() delegates to alist() correctly + assert await asyncio.to_thread(get_first_in_list) == state # update await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) From 3b2a237deb8c225f573c112d923fdaf9e38b249f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:45:57 -0700 Subject: [PATCH 12/27] Lint --- libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index a931181e1..6140af756 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -333,7 +333,7 @@ class AsyncPostgresSaver(BasePostgresSaver): def list( self, - config: RunnableConfig | None, + config: Optional[RunnableConfig], *, filter: Optional[dict[str, Any]] = None, before: Optional[RunnableConfig] = None, From dd4a15b569e3eec24c04f33b77b6dca0286e1876 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:57:33 -0700 Subject: [PATCH 13/27] Require latest version of langgraph-checkpoint --- libs/checkpoint-postgres/poetry.lock | 7 +++---- libs/checkpoint-postgres/pyproject.toml | 2 +- libs/checkpoint-sqlite/poetry.lock | 7 +++---- libs/checkpoint-sqlite/pyproject.toml | 2 +- 4 files changed, 8 insertions(+), 10 deletions(-) diff --git a/libs/checkpoint-postgres/poetry.lock b/libs/checkpoint-postgres/poetry.lock index f9ca0f97e..520300e91 100644 --- a/libs/checkpoint-postgres/poetry.lock +++ b/libs/checkpoint-postgres/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "annotated-types" @@ -266,7 +266,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.6" +version = "1.0.8" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -766,7 +766,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -972,4 +971,4 @@ watchmedo = ["PyYAML (>=3.10)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0,<4.0" -content-hash = "b139531e8c6f4e24cea4bdfc29d111d1ea000a6a8a604ba81be4bff0977a2466" +content-hash = "e294b6996aa6c8f671e6aaf65be8b4aba94c18e2f237dd3ee0e1b777849ce8a8" diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 077f8dd09..fb68cc181 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }] [tool.poetry.dependencies] python = "^3.9.0,<4.0" -langgraph-checkpoint = "^1.0.1" +langgraph-checkpoint = "^1.0.8" orjson = ">=3.10.1" psycopg = "^3.0.0" psycopg-pool = "^3.0.0" diff --git a/libs/checkpoint-sqlite/poetry.lock b/libs/checkpoint-sqlite/poetry.lock index 0ea747b2c..8d5dbe3af 100644 --- a/libs/checkpoint-sqlite/poetry.lock +++ b/libs/checkpoint-sqlite/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiosqlite" @@ -252,7 +252,7 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0" [[package]] name = "langgraph-checkpoint" -version = "1.0.1" +version = "1.0.8" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -651,7 +651,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -835,4 +834,4 @@ watchmedo = ["PyYAML (>=3.10)"] [metadata] lock-version = "2.0" python-versions = "^3.9.0" -content-hash = "d50ec7c6b55075d19193e080cc95d3b0998b9efa1a0ed407bd6055b5b5e867e0" +content-hash = "752a22dc2b57a0818a3a4d9bf5f62226ba7f8e0c458892551bf8b4c93723dbc1" diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index 536895f77..4b052d387 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }] [tool.poetry.dependencies] python = "^3.9.0" -langgraph-checkpoint = "^1.0.1" +langgraph-checkpoint = "^1.0.8" aiosqlite = "^0.20.0" [tool.poetry.group.dev.dependencies] From 6e9f54e5ae9bd70ac0a114ddbe9b1169b9d9b89f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:57:55 -0700 Subject: [PATCH 14/27] postgres 1.0.6 --- libs/checkpoint-postgres/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index fb68cc181..6e3640e7c 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "1.0.5" +version = "1.0.6" description = "Library with a Postgres implementation of LangGraph checkpoint saver." authors = [] license = "MIT" From af666a87fbf805d1942b5912ebab78822aa100f7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 17:58:01 -0700 Subject: [PATCH 15/27] sqlite 1.0.2 --- libs/checkpoint-sqlite/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index 4b052d387..9d8818619 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-sqlite" -version = "1.0.1" +version = "1.0.2" description = "Library with a SQLite implementation of LangGraph checkpoint saver." authors = [] license = "MIT" From 322d6c8dde4a99a88c76eb0a13d157545a2bfec2 Mon Sep 17 00:00:00 2001 From: Jacob Lee Date: Mon, 2 Sep 2024 18:48:50 -0700 Subject: [PATCH 16/27] patch(docs): Update docs with JS cloud instructions (#1560) * Update docs with JS cloud instructions * Add to index * Add md fence * nits * nits * Update setup_javascript.md --------- Co-authored-by: isaac hershenson --- .../docs/cloud/deployment/setup_javascript.md | 200 ++++++++++++++++++ docs/docs/cloud/faq/studio.md | 18 +- .../cloud/how-tos/interrupt_concurrent.md | 2 +- docs/docs/cloud/quick_start.md | 199 ++++++++++++++--- docs/mkdocs.yml | 1 + 5 files changed, 383 insertions(+), 37 deletions(-) create mode 100644 docs/docs/cloud/deployment/setup_javascript.md diff --git a/docs/docs/cloud/deployment/setup_javascript.md b/docs/docs/cloud/deployment/setup_javascript.md new file mode 100644 index 000000000..fa4c7b823 --- /dev/null +++ b/docs/docs/cloud/deployment/setup_javascript.md @@ -0,0 +1,200 @@ +# How to Set Up a LangGraph.js Application for Deployment + +A [LangGraph.js](https://langchain-ai.github.io/langgraphjs/) application must be configured with a [LangGraph API configuration file](../reference/cli.md#configuration-file) in order to be deployed to LangGraph Cloud (or to be self-hosted). This how-to guide discusses the basic steps to setup a LangGraph.js application for deployment using `package.json` to specify project dependencies. + +This walkthrough is based on [this repository](https://github.com/langchain-ai/langgraphjs-studio-starter), which you can play around with to learn more about how to setup your LangGraph application for deployment. + +The final repo structure will look something like this: + +```bash +my-app/ +├── src # all project code lies within here +│ ├── utils # optional utilities for your graph +│ │ ├── tools.ts # tools for your graph +│ │ ├── nodes.ts # node functions for you graph +│ │ └── state.ts # state definition of your graph +│   └── agent.ts # code for constructing your graph +├── package.json # package dependencies +├── .env # environment variables +└── langgraph.json # configuration file for LangGraph +``` + +After each step, an example file directory is provided to demonstrate how code can be organized. + +## Specify Dependencies + +Dependencies can be specified in a `package.json`. If none of these files is created, then dependencies can be specified later in the [LangGraph API configuration file](#create-langgraph-api-config). + +Example `package.json` file: + +```json +{ + "name": "langgraphjs-studio-starter", + "packageManager": "yarn@1.22.22", + "dependencies": { + "@langchain/community": "^0.2.31", + "@langchain/core": "^0.2.31", + "@langchain/langgraph": "^0.2.0", + "@langchain/openai": "^0.2.8" + } +} +``` + +Example file directory: + +```bash +my-app/ +└── package.json # package dependencies +``` + +## Specify Environment Variables + +Environment variables can optionally be specified in a file (e.g. `.env`). See the [Environment Variables reference](../reference/env_var.md) to configure additional variables for a deployment. + +Example `.env` file: + +``` +MY_ENV_VAR_1=foo +MY_ENV_VAR_2=bar +OPENAI_API_KEY=key +TAVILY_API_KEY=key_2 +``` + +Example file directory: + +```bash +my-app/ +├── package.json +└── .env # environment variables +``` + +## Define Graphs + +Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each compiled graph to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph API configuration file](../reference/cli.md#configuration-file). + +Here is an example `agent.ts`: + +```ts +import type { AIMessage } from "@langchain/core/messages"; +import { TavilySearchResults } from "@langchain/community/tools/tavily_search"; +import { ChatOpenAI } from "@langchain/openai"; + +import { MessagesAnnotation, StateGraph } from "@langchain/langgraph"; +import { ToolNode } from "@langchain/langgraph/prebuilt"; + +const tools = [ + new TavilySearchResults({ maxResults: 3, }), +]; + +// Define the function that calls the model +async function callModel( + state: typeof MessagesAnnotation.State, +) { + /** + * Call the LLM powering our agent. + * Feel free to customize the prompt, model, and other logic! + */ + const model = new ChatOpenAI({ + model: "gpt-4o", + }).bindTools(tools); + + const response = await model.invoke([ + { + role: "system", + content: `You are a helpful assistant. The current date is ${new Date().getTime()}.` + }, + ...state.messages + ]); + + // MessagesAnnotation supports returning a single message or array of messages + return { messages: response }; +} + +// Define the function that determines whether to continue or not +function routeModelOutput(state: typeof MessagesAnnotation.State) { + const messages = state.messages; + const lastMessage: AIMessage = messages[messages.length - 1]; + // If the LLM is invoking tools, route there. + if ((lastMessage?.tool_calls?.length ?? 0) > 0) { + return "tools"; + } + // Otherwise end the graph. + return "__end__"; +} + +// Define a new graph. +// See https://langchain-ai.github.io/langgraphjs/how-tos/define-state/#getting-started for +// more on defining custom graph states. +const workflow = new StateGraph(MessagesAnnotation) + // Define the two nodes we will cycle between + .addNode("callModel", callModel) + .addNode("tools", new ToolNode(tools)) + // Set the entrypoint as `callModel` + // This means that this node is the first one called + .addEdge("__start__", "callModel") + .addConditionalEdges( + // First, we define the edges' source node. We use `callModel`. + // This means these are the edges taken after the `callModel` node is called. + "callModel", + // Next, we pass in the function that will determine the sink node(s), which + // will be called after the source node is called. + routeModelOutput, + // List of the possible destinations the conditional edge can route to. + // Required for conditional edges to properly render the graph in Studio + [ + "tools", + "__end__" + ], + ) + // This means that after `tools` is called, `callModel` node is called next. + .addEdge("tools", "callModel"); + +// Finally, we compile it! +// This compiles it into a graph you can invoke and deploy. +export const graph = workflow.compile(); +``` + +!!! info "Assign `CompiledGraph` to Variable" + The build process for LangGraph Cloud requires that the `CompiledGraph` object be assigned to a variable at the top-level of a JavaScript module (alternatively, you can provide [a function that creates a graph](./graph_rebuild.md)). + +Example file directory: + +```bash +my-app/ +├── src # all project code lies within here +│ ├── utils # optional utilities for your graph +│ │ ├── tools.ts # tools for your graph +│ │ ├── nodes.ts # node functions for you graph +│ │ └── state.ts # state definition of your graph +│   └── agent.ts # code for constructing your graph +├── package.json # package dependencies +├── .env # environment variables +└── langgraph.json # configuration file for LangGraph +``` + +## Create LangGraph API Config + +Create a [LangGraph API configuration file](../reference/cli.md#configuration-file) called `langgraph.json`. See the [LangGraph CLI reference](../reference/cli.md#configuration-file) for detailed explanations of each key in the JSON object of the configuration file. + +Example `langgraph.json` file: + +```json +{ + "node_version": "20", + "dockerfile_lines": [], + "dependencies": ["."], + "graphs": { + "agent": "./src/agent.ts:graph" + }, + "env": ".env" +} +``` + +Note that the variable name of the `CompiledGraph` appears at the end of the value of each subkey in the top-level `graphs` key (i.e. `:`). + +!!! info "Configuration Location" + The LangGraph API configuration file must be placed in a directory that is at the same level or higher than the TypeScript files that contain compiled graphs and associated dependencies. + +## Next + +After you setup your project and place it in a github repo, it's time to [deploy your app](./cloud.md). diff --git a/docs/docs/cloud/faq/studio.md b/docs/docs/cloud/faq/studio.md index cef962f34..070014ce7 100644 --- a/docs/docs/cloud/faq/studio.md +++ b/docs/docs/cloud/faq/studio.md @@ -43,15 +43,23 @@ If you don't define your conditional edges carefully, you might notice extra edg ### Solution 1: Include a path map -The first way to solve this is to add path maps to your conditional edges. A path map is just a dictionary that maps the possible outputs of your router function with the names of the nodes that each output corresponds to. The path map is passed as the third argument to the `add_conditional_edges` function like so: +The first way to solve this is to add path maps to your conditional edges. A path map is just a dictionary or array that maps the possible outputs of your router function with the names of the nodes that each output corresponds to. The path map is passed as the third argument to the `add_conditional_edges` function like so: -```python -graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"}) -``` +=== "Python" + + ```python + graph.add_conditional_edges("node_a", routing_function, {True: "node_b", False: "node_c"}) + ``` + +=== "Javascript" + + ```ts + graph.addConditionalEdges("node_a", routingFunction, ["node_b", "node_c"]); + ``` In this case, the routing function returns either True or False, which map to `node_b` and `node_c` respectively. -### Solution 2: Update the typing of the router +### Solution 2: Update the typing of the router (Python only) Instead of passing a path map, you can also be explicit about the typing of your routing function by specifying the nodes it can map to using the `Literal` python definition. Here is an example of how to define a routing function in that way: diff --git a/docs/docs/cloud/how-tos/interrupt_concurrent.md b/docs/docs/cloud/how-tos/interrupt_concurrent.md index e67fb4ce5..b09f908ef 100644 --- a/docs/docs/cloud/how-tos/interrupt_concurrent.md +++ b/docs/docs/cloud/how-tos/interrupt_concurrent.md @@ -44,7 +44,7 @@ Now, let's import our required packages and instantiate our client, assistant, a const thread = await client.threads.create(); ``` -Now we can start our two runs and join the second on euntil it has completed: +Now we can start our two runs and join the second one until it has completed: === "Python" diff --git a/docs/docs/cloud/quick_start.md b/docs/docs/cloud/quick_start.md index c59a6fb85..755ce0df4 100644 --- a/docs/docs/cloud/quick_start.md +++ b/docs/docs/cloud/quick_start.md @@ -14,13 +14,25 @@ This tutorial will use: 1. Create a new application with the following directory and files: +=== "Python" + / |-- agent.py # code for your LangGraph agent |-- requirements.txt # Python packages required for your graph |-- langgraph.json # configuration file for LangGraph |-- .env # environment files with API keys -2. The `agent.py` file should contain Python code for defining your graph. The following code is a simple example, the important thing is that at some point in your file you compile your graph and assign the compiled graph to a variable (in this case the `graph` variable). This example code uses `create_react_agent`, a prebuilt agent, read more about it [here](..//concepts/agentic_concepts.md#react-agent). +=== "Javascript" + + / + |-- agent.ts # code for your LangGraph agent + |-- package.json # Javascript packages required for your graph + |-- langgraph.json # configuration file for LangGraph + |-- .env # environment files with API keys + +2. The `agent.py`/`agent.ts` file should contain code for defining your graph. The following code is a simple example, the important thing is that at some point in your file you compile your graph and assign the compiled graph to a variable (in this case the `graph` variable). This example code uses `create_react_agent`, a prebuilt agent. You can read more about it [here](../concepts/agentic_concepts.md#react-agent). + +=== "Python" ```python from langchain_anthropic import ChatAnthropic @@ -34,14 +46,53 @@ This tutorial will use: graph = create_react_agent(model, tools) ``` -3. The `requirements.txt` file should contain any dependencies for your graph(s). In this case we only require four packages for our graph to run: +=== "Javascript" - langgraph - langchain_anthropic - tavily-python - langchain_community + ```ts + import { ChatAnthropic } from "@langchain/anthropic"; + import { TavilySearchResults } from "@langchain/community/tools/tavily_search"; + import { createReactAgent } from "@langchain/langgraph/prebuilt"; -4. The [`langgraph.json`][langgraph.json] file is a configuration file that describes what graph(s) you are going to host. In this case we only have one graph to host: the compiled `graph` object from `agent.py`. + const model = new ChatAnthropic({ + model: "claude-3-5-sonnet-20240620", + }); + + const tools = [ + new TavilySearchResults({ maxResults: 3, }), + ]; + + export const graph = createReactAgent({ llm: model, tools }); + ``` + +3. The `requirements.txt`/`package.json` file should contain any dependencies for your graph(s). In this case we only require four packages for our graph to run: + +=== "Python" + + ```python + langgraph + langchain_anthropic + tavily-python + langchain_community + ``` + +=== "Javascript" + + ```js + { + "name": "my-app", + "packageManager": "yarn@1.22.22", + "dependencies": { + "@langchain/community": "^0.2.31", + "@langchain/core": "^0.2.31", + "@langchain/langgraph": "0.2.0", + "@langchain/openai": "^0.2.8" + } + } + ``` + +4. The [`langgraph.json`][langgraph.json] file is a configuration file that describes what graph(s) you are going to host. In this case we only have one graph to host: the compiled `graph` object from `agent.py`/`agent.ts`. + +=== "Python" ```json { @@ -53,7 +104,21 @@ This tutorial will use: } ``` - Learn more about the LangGraph CLI configuration file [here](./reference/cli.md#configuration-file). +=== "Javascript" + + ```json + { + "node_version": "20", + "dockerfile_lines": [], + "dependencies": ["."], + "graphs": { + "agent": "./src/agent.ts:graph" + }, + "env": ".env" + } + ``` + +Learn more about the LangGraph CLI configuration file [here](./reference/cli.md#configuration-file). 5. The `.env` file should have any environment variables needed to run your graph. This will only be used for local testing, so if you are not testing locally you can skip this step. NOTE: if you do add this, you should NOT check this into git. For this graph, we need two environment variables: @@ -206,36 +271,108 @@ export LANGSMITH_API_KEY=... The first thing to do when using the SDK is to setup our client, access our assistant, and create a thread to execute a run on: -```python -from langgraph_sdk import get_client +=== "Python" -# Replace this with the URL of your own deployed graph -URL = "https://chatbot-23a570f3210f52a7b167f09f6158e3b3-ffoprvkqsa-uc.a.run.app" -client = get_client(url=URL) + ```python + from langgraph_sdk import get_client -# Search all hosted graphs -assistants = await client.assistants.search() -# In this example we select the first assistant since we are only hosting a single graph -assistant = assistants[0] + client = get_client(url=) + # get default assistant + assistants = await client.assistants.search() + assistant = [a for a in assistants if not a["config"]][0] + # create thread + thread = await client.threads.create() + print(thread) + ``` -# We create a thread for tracking the state of our run -thread = await client.threads.create() -``` +=== "Javascript" + + ```js + import { Client } from "@langchain/langgraph-sdk"; + + const client = new Client({ apiUrl: }); + // get default assistant + const assistants = await client.assistants.search(); + const assistant = assistants.find(a => !a.config); + // create thread + const thread = await client.threads.create(); + console.log(thread) + ``` + +=== "CURL" + + ```bash + curl --request POST \ + --url /assistants/search \ + --header 'Content-Type: application/json' \ + --data '{ + "limit": 10, + "offset": 0 + }' | jq -c 'map(select(.config == null or .config == {})) | .[0]' && \ + curl --request POST \ + --url /threads \ + --header 'Content-Type: application/json' \ + --data '{}' + ``` We can then execute a run on the thread: -```python -input = {"messages":[{"role": "user", "content": "Hello! My name is Bagatur and I am 26 years old."}]} +=== "Python" -async for chunk in client.runs.stream( - thread['thread_id'], - assistant["assistant_id"], - input=input, - stream_mode="updates", - ): - if chunk.data and chunk.event != "metadata": - print(chunk.data) -``` + ```python + input = {"messages":[{"role": "user", "content": "Hello! My name is Bagatur and I am 26 years old."}]} + + async for chunk in client.runs.stream( + thread['thread_id'], + assistant["assistant_id"], + input=input, + stream_mode="updates", + ): + if chunk.data and chunk.event != "metadata": + print(chunk.data) + ``` + +=== "Javascript" + + ```js + const input = { "messages":[{ "role": "user", "content": "Hello! My name is Bagatur and I am 26 years old." }] }; + + const streamResponse = client.runs.stream( + thread["thread_id"], + assistant["assistant_id"], + { + input, + } + ); + for await (const chunk of streamResponse) { + if (chunk.data && chunk.event !== "metadata" ) { + console.log(chunk.data); + } + } + ``` + +=== "CURL" + + ```bash + curl --request POST \ + --url /threads//runs/stream \ + --header 'Content-Type: application/json' \ + --data "{ + \"assistant_id\": , + \"input\": {\"messages\": [{\"role\": \"human\", \"content\": \"Hello! My name is Bagatur and I am 26 years old.\"}]}, + }" | sed 's/\r$//' | awk ' + /^event:/ { event = $2 } + /^data:/ { + json_data = substr($0, index($0, $2)) + + if (event != "metadata") { + print json_data + } + }' + ``` + + +Output: {'agent': {'messages': [{'content': "Hi Bagatur! It's nice to meet you. How can I assist you today?", 'additional_kwargs': {}, 'response_metadata': {'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-05-13', 'system_fingerprint': 'fp_9cb5d38cf7'}, 'type': 'ai', 'name': None, 'id': 'run-c89118b7-1b1e-42b9-a85d-c43fe99881cd', 'example': False, 'tool_calls': [], 'invalid_tool_calls': [], 'usage_metadata': None}]}} diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index eb6d5d36f..56a9435b2 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -197,6 +197,7 @@ nav: - Setup: - Setup App: "cloud/deployment/setup.md" - Setup App (pyproject.toml): "cloud/deployment/setup_pyproject.md" + - Setup App (JavaScript): "cloud/deployment/setup_javascript.md" - Rebuild Graph at Runtime: "cloud/deployment/graph_rebuild.md" - Test App Locally: "cloud/deployment/test_locally.md" - Deployment: From 8c0b27a85bf58d503c6b79e9f6acc68263e09334 Mon Sep 17 00:00:00 2001 From: gbaian10 <34255899+gbaian10@users.noreply.github.com> Date: Tue, 3 Sep 2024 21:50:00 +0800 Subject: [PATCH 17/27] docs: fix Corrective RAG (CRAG) miss function (#1563) `Corrective RAG (CRAG) using local LLMs` miss `find_tool_calls_react` function --- examples/rag/langgraph_crag_local.ipynb | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/examples/rag/langgraph_crag_local.ipynb b/examples/rag/langgraph_crag_local.ipynb index dbac6f697..d6e791ae3 100644 --- a/examples/rag/langgraph_crag_local.ipynb +++ b/examples/rag/langgraph_crag_local.ipynb @@ -429,6 +429,14 @@ "]\n", "\n", "\n", + "def find_tool_calls_react(messages):\n", + " \"\"\"\n", + " Find all tool calls in the messages returned\n", + " \"\"\"\n", + " tool_calls = [tc['name'] for m in messages['messages'] for tc in getattr(m, 'tool_calls', [])]\n", + " return tool_calls\n", + "\n", + "\n", "def check_trajectory_react(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", From c1d7cbd67c49fbfc7554d4e7c020c1ca964b2297 Mon Sep 17 00:00:00 2001 From: Jacob Lee Date: Tue, 3 Sep 2024 07:47:31 -0700 Subject: [PATCH 18/27] Adds JS deploy guide to index (#1590) --- docs/docs/cloud/how-tos/index.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/docs/cloud/how-tos/index.md b/docs/docs/cloud/how-tos/index.md index 255490d32..90b13c35e 100644 --- a/docs/docs/cloud/how-tos/index.md +++ b/docs/docs/cloud/how-tos/index.md @@ -13,6 +13,7 @@ LangGraph Cloud gives you best in class observability, testing, and hosting serv - [How to set up app for deployment (requirements.txt)](../deployment/setup.md) - [How to set up app for deployment (pyproject.toml)](../deployment/setup_pyproject.md) +- [How to set up app for deployment (JavaScript)](../deployment/setup_javascript.md) - [How to test locally](../deployment/test_locally.md) - [How to deploy to LangGraph cloud](../deployment/cloud.md) - [How to self-host](../deployment/self_hosted.md) From e5b09b8c202de3798044b2067a143be8ba193dec Mon Sep 17 00:00:00 2001 From: gbaian10 <34255899+gbaian10@users.noreply.github.com> Date: Tue, 3 Sep 2024 23:00:14 +0800 Subject: [PATCH 19/27] docs: update `Postgres checkpointer for persistence` (#1556) --- examples/persistence_postgres.ipynb | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/examples/persistence_postgres.ipynb b/examples/persistence_postgres.ipynb index 5ea838cee..b4fa1e2fe 100644 --- a/examples/persistence_postgres.ipynb +++ b/examples/persistence_postgres.ipynb @@ -132,8 +132,6 @@ "metadata": {}, "outputs": [], "source": [ - "from psycopg.rows import dict_row\n", - "\n", "connection_kwargs = {\n", " \"autocommit\": True,\n", " \"prepare_threshold\": 0,\n", @@ -161,15 +159,13 @@ "source": [ "from psycopg_pool import ConnectionPool\n", "\n", - "pool = ConnectionPool(\n", + "with ConnectionPool(\n", " # Example configuration\n", " conninfo=DB_URI,\n", " max_size=20,\n", " kwargs=connection_kwargs,\n", - ")\n", - "\n", - "with pool.connection() as conn:\n", - " checkpointer = PostgresSaver(conn)\n", + ") as pool:\n", + " checkpointer = PostgresSaver(pool)\n", "\n", " # NOTE: you need to call .setup() the first time you're using your checkpointer\n", " checkpointer.setup()\n", @@ -394,8 +390,8 @@ " conninfo=DB_URI,\n", " max_size=20,\n", " kwargs=connection_kwargs,\n", - ") as pool, pool.connection() as conn:\n", - " checkpointer = AsyncPostgresSaver(conn)\n", + ") as pool:\n", + " checkpointer = AsyncPostgresSaver(pool)\n", "\n", " # NOTE: you need to call .setup() the first time you're using your checkpointer\n", " # await checkpointer.setup()\n", From d492475c6133b39ed6aa14af99883e9c7b46906c Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Tue, 3 Sep 2024 11:00:37 -0400 Subject: [PATCH 20/27] checkpoint-sqlite: close connection in from_conn_string (#1589) --- .../langgraph/checkpoint/sqlite/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 04493a381..a5f657824 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -104,10 +104,12 @@ class SqliteSaver(BaseCheckpointSaver): with SqliteSaver.from_conn_string("checkpoints.sqlite") as memory: ... """ - with sqlite3.connect( - conn_string, - # https://ricardoanderegg.com/posts/python-sqlite-thread-safety/ - check_same_thread=False, + with closing( + sqlite3.connect( + conn_string, + # https://ricardoanderegg.com/posts/python-sqlite-thread-safety/ + check_same_thread=False, + ) ) as conn: yield SqliteSaver(conn) From 43d5dbf1741fb81d5fa34b1e34146930115c6702 Mon Sep 17 00:00:00 2001 From: Sreenivasan R S <101383058+sreeniSudharsan@users.noreply.github.com> Date: Tue, 3 Sep 2024 21:42:31 +0530 Subject: [PATCH 21/27] examples: fix missing import introduction (#1584) --- examples/introduction.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/introduction.ipynb b/examples/introduction.ipynb index 3d0287d05..52cd1018b 100644 --- a/examples/introduction.ipynb +++ b/examples/introduction.ipynb @@ -1504,7 +1504,7 @@ "from langgraph.checkpoint.memory import MemorySaver\n", "from langgraph.graph import StateGraph\n", "from langgraph.graph.message import add_messages\n", - "from langgraph.prebuilt import ToolNode\n", + "from langgraph.prebuilt import ToolNode, tools_condition\n", "\n", "\n", "class State(TypedDict):\n", From 3c5913aa29d2c43c20f6378c909b050e92f47f28 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 19:17:38 -0700 Subject: [PATCH 22/27] Remove pydantic base class from Pregel/PregelNode --- libs/langgraph/langgraph/graph/graph.py | 4 + libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/langgraph/pregel/__init__.py | 89 +++++++------ libs/langgraph/langgraph/pregel/read.py | 126 ++++++++++++------ .../tests/__snapshots__/test_pregel.ambr | 60 ++------- libs/langgraph/tests/test_pregel.py | 13 +- 6 files changed, 155 insertions(+), 139 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index d466a3c96..570c06bad 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -424,6 +424,10 @@ class Graph: class CompiledGraph(Pregel): builder: Graph + def __init__(self, *, builder: Graph, **kwargs): + super().__init__(**kwargs) + self.builder = builder + def attach_node(self, key: str, node: NodeSpec) -> None: self.channels[key] = EphemeralValue(Any) self.nodes[key] = ( diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f576a98cd..307b672aa 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -17,10 +17,10 @@ from typing import ( overload, ) -from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike from langchain_core.runnables.utils import create_model +from pydantic import BaseModel from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 01f0ad91d..c5fa760c5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -25,12 +25,10 @@ from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.load.dump import dumpd -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.runnables import ( Runnable, RunnableLambda, RunnableSequence, - RunnableSerializable, ) from langchain_core.runnables.base import Input, Output, coerce_to_runnable from langchain_core.runnables.config import ( @@ -48,6 +46,7 @@ from langchain_core.runnables.utils import ( get_unique_config_specs, ) from langchain_core.tracers._streaming import _StreamingCallbackHandler +from pydantic import BaseModel from typing_extensions import Self from langgraph.channels.base import ( @@ -186,16 +185,10 @@ class Channel: ) -class Pregel( - RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] -): +class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): nodes: Mapping[str, PregelNode] - channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field( - default_factory=dict - ) - - auto_validate: bool = True + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" @@ -205,16 +198,16 @@ class Pregel( stream_channels: Optional[Union[str, Sequence[str]]] = None """Channels to stream, defaults to all channels not in reserved channels""" - interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_after_nodes: Union[All, Sequence[str]] - interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_before_nodes: Union[All, Sequence[str]] input_channels: Union[str, Sequence[str]] step_timeout: Optional[float] = None """Maximum time to wait for a step to complete, in seconds. Defaults to None.""" - debug: bool = Field(default_factory=get_debug) + debug: bool """Whether to print debug information during execution. Defaults to False.""" checkpointer: Optional[BaseCheckpointSaver] = None @@ -232,36 +225,50 @@ class Pregel( name: str = "LangGraph" - class Config: - arbitrary_types_allowed = True + def __init__( + self, + *, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = None, + auto_validate: bool = True, + stream_mode: StreamMode = "values", + output_channels: Union[str, Sequence[str]], + stream_channels: Optional[Union[str, Sequence[str]]] = None, + interrupt_after_nodes: Union[All, Sequence[str]] = (), + interrupt_before_nodes: Union[All, Sequence[str]] = (), + input_channels: Union[str, Sequence[str]], + step_timeout: Optional[float] = None, + debug: bool = False, + checkpointer: Optional[BaseCheckpointSaver] = None, + store: Optional[BaseStore] = None, + retry_policy: Optional[RetryPolicy] = None, + config_type: Optional[Type[Any]] = None, + config: Optional[RunnableConfig] = None, + name: str = "LangGraph", + ) -> None: + self.nodes = nodes + self.channels = channels or {} + self.stream_mode = stream_mode + self.output_channels = output_channels + self.stream_channels = stream_channels + self.interrupt_after_nodes = interrupt_after_nodes + self.interrupt_before_nodes = interrupt_before_nodes + self.input_channels = input_channels + self.step_timeout = step_timeout + self.debug = debug if debug is not None else get_debug() + self.checkpointer = checkpointer + self.store = store + self.retry_policy = retry_policy + self.config_type = config_type + self.config = config + self.name = name + if auto_validate: + self.validate() def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: - return self.copy( - update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})} - ) - - @classmethod - def is_lc_serializable(cls) -> bool: - """Return whether the graph can be serialized by Langchain.""" - return True - - @root_validator(skip_on_failure=True) - def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]: - if not values["auto_validate"]: - return values - validate_graph( - values["nodes"], - values["channels"], - values["input_channels"], - values["output_channels"], - values["stream_channels"], - values["interrupt_after_nodes"], - values["interrupt_before_nodes"], - ) - if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]: - if not values["checkpointer"]: - raise ValueError("Interrupts require a checkpointer") - return values + attrs = {**self.__dict__} + attrs["config"] = merge_configs(self.config, config, kwargs) + return self.__class__(**attrs) def validate(self) -> Self: validate_graph( diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 163665e2e..ca0828fda 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -1,8 +1,16 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import ( + Any, + AsyncIterator, + Callable, + Iterator, + Mapping, + Optional, + Sequence, + Union, +) -from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( Runnable, RunnableConfig, @@ -10,7 +18,7 @@ from langchain_core.runnables import ( RunnableSequence, RunnableSerializable, ) -from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable +from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec @@ -99,20 +107,47 @@ class ChannelRead(RunnableCallable): DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() -class PregelNode(RunnableBindingBase): +class PregelNode(Runnable): channels: Union[list[str], Mapping[str, str]] - triggers: list[str] = Field(default_factory=list) + triggers: list[str] - mapper: Optional[Callable[[Any], Any]] = None + mapper: Optional[Callable[[Any], Any]] - writers: list[Runnable] = Field(default_factory=list) + writers: list[Runnable] - bound: Runnable[Any, Any] = Field(default=DEFAULT_BOUND) + bound: Runnable[Any, Any] - kwargs: Mapping[str, Any] = Field(default_factory=dict) + retry_policy: Optional[RetryPolicy] - retry_policy: Optional[RetryPolicy] = None + config: RunnableConfig + + def __init__( + self, + *, + channels: Union[list[str], Mapping[str, str]], + triggers: Sequence[str], + mapper: Optional[Callable[[Any], Any]] = None, + writers: Optional[list[Runnable]] = None, + tags: Optional[list[str]] = None, + metadata: Optional[Mapping[str, Any]] = None, + bound: Optional[Runnable[Any, Any]] = None, + retry_policy: Optional[RetryPolicy] = None, + config: Optional[RunnableConfig] = None, + ) -> None: + self.channels = channels + self.triggers = list(triggers) + self.mapper = mapper + self.writers = writers or [] + self.bound = bound if bound is not None else DEFAULT_BOUND + self.retry_policy = retry_policy + self.config = merge_configs( + config, {"tags": tags or [], "metadata": metadata or {}} + ) + + def copy(self, update: dict[str, Any]) -> PregelNode: + attrs = {**self.__dict__, **update} + return PregelNode(**attrs) def get_writers(self) -> list[Runnable]: """Get writers with optimizations applied.""" @@ -145,38 +180,6 @@ class PregelNode(RunnableBindingBase): else: return self.bound - def __init__( - self, - *, - channels: Union[list[str], Mapping[str, str]], - triggers: Sequence[str], - mapper: Optional[Callable[[Any], Any]] = None, - writers: Optional[list[Runnable]] = None, - tags: Optional[list[str]] = None, - metadata: Optional[Mapping[str, Any]] = None, - bound: Optional[Runnable[Any, Any]] = None, - kwargs: Optional[Mapping[str, Any]] = None, - config: Optional[RunnableConfig] = None, - retry_policy: Optional[RetryPolicy] = None, - **other_kwargs: Any, - ) -> None: - super().__init__( - channels=channels, - triggers=triggers, - mapper=mapper, - writers=writers or [], - bound=bound or DEFAULT_BOUND, - kwargs=kwargs or {}, - retry_policy=retry_policy, - config=merge_configs( - config, {"tags": tags or [], "metadata": metadata or {}} - ), - **other_kwargs, - ) - - def __repr_args__(self) -> Any: - return [(k, v) for k, v in super().__repr_args__() if k != "bound"] - def join(self, channels: Sequence[str]) -> PregelNode: assert isinstance(channels, list) or isinstance( channels, tuple @@ -226,3 +229,42 @@ class PregelNode(RunnableBindingBase): ], ) -> RunnableSerializable: raise NotImplementedError() + + def invoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return self.bound.invoke(input, merge_configs(self.config, config), **kwargs) + + async def ainvoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return await self.bound.ainvoke( + input, merge_configs(self.config, config), **kwargs + ) + + def stream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Iterator[Output]: + yield from self.bound.stream( + input, merge_configs(self.config, config), **kwargs + ) + + async def astream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> AsyncIterator[Output]: + async for item in self.bound.astream( + input, merge_configs(self.config, config), **kwargs + ): + yield item diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 3c3dd49d0..5ed4cf469 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -3300,11 +3300,7 @@ 'type': 'string', }), }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3323,11 +3319,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3370,11 +3362,7 @@ 'type': 'string', }), }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3393,11 +3381,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3440,11 +3424,7 @@ 'type': 'string', }), }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3463,11 +3443,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3510,11 +3486,7 @@ 'type': 'string', }), }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3533,11 +3505,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3580,11 +3548,7 @@ 'type': 'string', }), }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3603,11 +3567,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3ca3906d1..6109f6ce4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1903,7 +1903,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") two = Channel.subscribe_to("between") | add_one - with pytest.raises(ValueError): + with pytest.raises(TypeError): Pregel(nodes={"one": one, "two": two}) @@ -9784,6 +9784,7 @@ def test_doubly_nested_graph_state( ] +@pytest.mark.repeat(10) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_to_nested_graphs( request: pytest.FixtureRequest, checkpointer_name: str @@ -9943,10 +9944,12 @@ def test_send_to_nested_graphs( graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) # continue past interrupt - assert graph.invoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } + assert sorted( + graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0] + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] actual_snapshot = graph.get_state(config) expected_snapshot = StateSnapshot( From 4793b3f5e1c690a29b0773e36a4fe37bc53b2b78 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Sep 2024 19:40:02 -0700 Subject: [PATCH 23/27] Use a lock for all operations on sqlite checkpointer - Otherwise when used in multiple subgraphs in parallel separate queries can interfere w each other --- .../langgraph/checkpoint/sqlite/__init__.py | 21 ++++++++++--------- .../langgraph/checkpoint/sqlite/aio.py | 10 +++++---- libs/langgraph/tests/test_pregel.py | 1 - 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index a5f657824..31bffd19c 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -165,14 +165,15 @@ class SqliteSaver(BaseCheckpointSaver): Yields: sqlite3.Cursor: A cursor for the SQLite database. """ - self.setup() - cur = self.conn.cursor() - try: - yield cur - finally: - if transaction: - self.conn.commit() - cur.close() + with self.lock: + self.setup() + cur = self.conn.cursor() + try: + yield cur + finally: + if transaction: + self.conn.commit() + cur.close() def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -398,7 +399,7 @@ class SqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( @@ -434,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver): writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. task_id (str): Identifier for the task creating the writes. """ - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 7cdc7c8fb..28ad76109 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -276,7 +276,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): """ await self.setup() checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: # find the latest checkpoint for the thread_id if checkpoint_id := get_checkpoint_id(config): await cur.execute( @@ -371,7 +371,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver): ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur: + async with self.lock, self.conn.execute( + query, params + ) as cur, self.conn.cursor() as wcur: async for ( thread_id, checkpoint_ns, @@ -438,7 +440,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - async with self.conn.execute( + async with self.lock, self.conn.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), @@ -475,7 +477,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): task_id (str): Identifier for the task creating the writes. """ await self.setup() - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: await cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 6109f6ce4..96696fa13 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9784,7 +9784,6 @@ def test_doubly_nested_graph_state( ] -@pytest.mark.repeat(10) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_to_nested_graphs( request: pytest.FixtureRequest, checkpointer_name: str From a99b66ae3f1d08a848aa4e0eb710744ee465f5e6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Sep 2024 09:13:21 -0700 Subject: [PATCH 24/27] Update --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c5fa760c5..d7c265b2d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -238,7 +238,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): interrupt_before_nodes: Union[All, Sequence[str]] = (), input_channels: Union[str, Sequence[str]], step_timeout: Optional[float] = None, - debug: bool = False, + debug: Optional[bool] = None, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, retry_policy: Optional[RetryPolicy] = None, From 4ddb4ae5a037b329cd820a40453aeaaff2b61644 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Sep 2024 09:26:45 -0700 Subject: [PATCH 25/27] Update --- .../tests/__snapshots__/test_pregel.ambr | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 5ed4cf469..3937546f5 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -3300,6 +3300,10 @@ 'type': 'string', }), }), + 'required': list([ + 'query', + 'inner', + ]), 'title': 'LangGraphInput', 'type': 'object', }) @@ -3362,6 +3366,10 @@ 'type': 'string', }), }), + 'required': list([ + 'query', + 'inner', + ]), 'title': 'LangGraphInput', 'type': 'object', }) @@ -3424,6 +3432,10 @@ 'type': 'string', }), }), + 'required': list([ + 'query', + 'inner', + ]), 'title': 'LangGraphInput', 'type': 'object', }) @@ -3486,6 +3498,10 @@ 'type': 'string', }), }), + 'required': list([ + 'query', + 'inner', + ]), 'title': 'LangGraphInput', 'type': 'object', }) @@ -3548,6 +3564,10 @@ 'type': 'string', }), }), + 'required': list([ + 'query', + 'inner', + ]), 'title': 'LangGraphInput', 'type': 'object', }) From 5d1ab66535814d7b7787f75f0fb2794ae5c30436 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Sep 2024 09:41:04 -0700 Subject: [PATCH 26/27] sqlite 1.0.3 --- libs/checkpoint-sqlite/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index 9d8818619..1b9462085 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-sqlite" -version = "1.0.2" +version = "1.0.3" description = "Library with a SQLite implementation of LangGraph checkpoint saver." authors = [] license = "MIT" From 1ae60dccc5ff8183b87cf13a08e1f6a65d4192f5 Mon Sep 17 00:00:00 2001 From: Isaac Francisco <78627776+isahers1@users.noreply.github.com> Date: Tue, 3 Sep 2024 10:12:36 -0700 Subject: [PATCH 27/27] test (#1592) --- libs/sdk-js/src/client.mts | 32 ++++++++++++++++++-------------- libs/sdk-js/src/types.mts | 30 ++++++++++++++++++++++++------ 2 files changed, 42 insertions(+), 20 deletions(-) diff --git a/libs/sdk-js/src/client.mts b/libs/sdk-js/src/client.mts index 03024299d..459e9e6b1 100644 --- a/libs/sdk-js/src/client.mts +++ b/libs/sdk-js/src/client.mts @@ -133,6 +133,7 @@ export class CronsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, }; return this.fetch(`/threads/${threadId}/runs/crons`, { method: "POST", @@ -159,6 +160,7 @@ export class CronsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, }; return this.fetch(`/runs/crons`, { method: "POST", @@ -237,6 +239,8 @@ export class AssistantsClient extends BaseClient { graphId: string; config?: Config; metadata?: Metadata; + assistantId?: string; + ifExists?: OnConflictBehavior; }): Promise { return this.fetch("/assistants", { method: "POST", @@ -244,6 +248,8 @@ export class AssistantsClient extends BaseClient { graph_id: payload.graphId, config: payload.config, metadata: payload.metadata, + assistant_id: payload.assistantId, + if_exists: payload.ifExists, }, }); } @@ -257,7 +263,7 @@ export class AssistantsClient extends BaseClient { async update( assistantId: string, payload: { - graphId: string; + graphId?: string; config?: Config; metadata?: Metadata; }, @@ -518,7 +524,7 @@ export class RunsClient extends BaseClient { stream( threadId: null, assistantId: string, - payload?: Omit, + payload?: Omit, ): AsyncGenerator<{ event: StreamEvent; data: any; @@ -546,8 +552,6 @@ export class RunsClient extends BaseClient { payload?: RunsStreamPayload, ): AsyncGenerator<{ event: StreamEvent; - // TODO: figure out a better way to - // type this without any data: any; }> { const json: Record = { @@ -560,10 +564,11 @@ export class RunsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, checkpoint_id: payload?.checkpointId, + webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, + on_completion: payload?.onCompletion, + on_disconnect: payload?.onDisconnect, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } const endpoint = threadId == null ? `/runs/stream` : `/threads/${threadId}/runs/stream`; @@ -640,10 +645,8 @@ export class RunsClient extends BaseClient { interrupt_after: payload?.interruptAfter, webhook: payload?.webhook, checkpoint_id: payload?.checkpointId, + multitask_strategy: payload?.multitaskStrategy, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } return this.fetch(`/threads/${threadId}/runs`, { method: "POST", json, @@ -654,7 +657,7 @@ export class RunsClient extends BaseClient { async wait( threadId: null, assistantId: string, - payload?: Omit, + payload?: Omit, ): Promise; async wait( @@ -684,10 +687,11 @@ export class RunsClient extends BaseClient { interrupt_before: payload?.interruptBefore, interrupt_after: payload?.interruptAfter, checkpoint_id: payload?.checkpointId, + webhook: payload?.webhook, + multitask_strategy: payload?.multitaskStrategy, + on_completion: payload?.onCompletion, + on_disconnect: payload?.onDisconnect, }; - if (payload?.multitaskStrategy != null) { - json["multitask_strategy"] = payload?.multitaskStrategy; - } const endpoint = threadId == null ? `/runs/wait` : `/threads/${threadId}/runs/wait`; return this.fetch(endpoint, { diff --git a/libs/sdk-js/src/types.mts b/libs/sdk-js/src/types.mts index 872cc5723..2a5d6c4f1 100644 --- a/libs/sdk-js/src/types.mts +++ b/libs/sdk-js/src/types.mts @@ -3,6 +3,8 @@ import { Config, Metadata } from "./schema.js"; export type StreamMode = "values" | "messages" | "updates" | "events" | "debug"; export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue"; export type OnConflictBehavior = "raise" | "do_nothing"; +export type OnCompletionBehavior = "complete" | "continue"; +export type DisconnectMode = "cancel" | "continue"; export type StreamEvent = | "events" | "metadata" @@ -61,6 +63,27 @@ interface RunsInvokePayload { * Abort controller signal to cancel the run. */ signal?: AbortController["signal"]; + + /** + * Behavior to handle run completion. Only relevant if + * there is a pending/inflight run on the same thread. One of: + * - "complete": Complete the run. + * - "continue": Continue the run. + */ + onCompletion?: OnCompletionBehavior; + + /** + * Webhook to call when the run is complete. + */ + webhook?: string; + + /** + * Behavior to handle disconnection. Only relevant if + * there is a pending/inflight run on the same thread. One of: + * - "cancel": Cancel the run. + * - "continue": Continue the run. + */ + onDisconnect?: DisconnectMode; } export interface RunsStreamPayload extends RunsInvokePayload { @@ -82,12 +105,7 @@ export interface RunsStreamPayload extends RunsInvokePayload { feedbackKeys?: string[]; } -export interface RunsCreatePayload extends RunsInvokePayload { - /** - * Webhook to call when the run is complete. - */ - webhook?: string; -} +export interface RunsCreatePayload extends RunsInvokePayload {} export interface CronsCreatePayload extends RunsCreatePayload { /**