From 5f9b0bc0ae35f18658158ddff9b3b58cb11db462 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 6 Aug 2024 15:41:52 -0700 Subject: [PATCH 1/2] Add UntrackedValue to mark a state key as not checkpointable --- libs/langgraph/langgraph/channels/__init__.py | 6 ++ .../langgraph/channels/untracked_value.py | 62 +++++++++++++++++++ libs/langgraph/langgraph/graph/state.py | 2 + libs/langgraph/tests/test_pregel.py | 14 +---- libs/langgraph/tests/test_pregel_async.py | 9 +-- 5 files changed, 74 insertions(+), 19 deletions(-) create mode 100644 libs/langgraph/langgraph/channels/untracked_value.py diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index 975fb93e0..144056a58 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,11 +1,17 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue +from langgraph.channels.any_value import AnyValue __all__ = [ "LastValue", "Topic", "Context", "BinaryOperatorAggregate", + "UntrackedValue", + "EphemeralValue", + "AnyValue", ] diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py new file mode 100644 index 000000000..094d6c3ca --- /dev/null +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -0,0 +1,62 @@ +from contextlib import contextmanager +from typing import Generator, Generic, Optional, Sequence, Type + +from langchain_core.runnables import RunnableConfig +from typing_extensions import Self + +from langgraph.channels.base import BaseChannel, Value +from langgraph.errors import EmptyChannelError, InvalidUpdateError + + +class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): + """Stores the last value received, never checkpointed.""" + + def __init__(self, typ: Type[Value], guard: bool = True) -> None: + self.typ = typ + self.guard = guard + + def __eq__(self, value: object) -> bool: + return isinstance(value, UntrackedValue) and value.guard == self.guard + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + def checkpoint(self) -> Value: + raise EmptyChannelError() + + @contextmanager + def from_checkpoint( + self, checkpoint: Optional[Value], config: RunnableConfig + ) -> Generator[Self, None, None]: + empty = self.__class__(self.typ, self.guard) + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> bool: + if len(values) == 0: + return False + if len(values) != 1 and self.guard: + raise InvalidUpdateError( + "EphemeralValue can only receive one value per step." + ) + + self.value = values[-1] + return True + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index ffe35e8bc..4198f4ac6 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -705,6 +705,8 @@ def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]: meta = typ.__metadata__ if len(meta) >= 1 and isinstance(meta[-1], BaseChannel): return meta[-1] + elif len(meta) >= 1 and isclass(meta[-1]) and issubclass(meta[-1], BaseChannel): + return meta[-1](typ.__origin__ if hasattr(typ, "__origin__") else typ) return None diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 08cc53fbd..5e0e9aaf9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -39,6 +39,7 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -2562,7 +2563,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: from langchain_core.tools import tool class AgentState(TypedDict, total=False): - input: str + input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] session: Annotated[httpx.Client, Context(httpx.Client)] @@ -2753,7 +2754,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -2791,7 +2791,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2856,7 +2855,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", @@ -2914,7 +2912,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -2952,7 +2949,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -3017,7 +3013,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", @@ -3066,7 +3061,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "intermediate_steps": [], }, next=("agent",), @@ -3088,7 +3082,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -3132,7 +3125,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -3205,7 +3197,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -3249,7 +3240,6 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None: assert app_w_interrupt.get_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 23065b3d2..d20d7229f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -39,6 +39,7 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic +from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -2723,7 +2724,7 @@ async def test_conditional_graph_state() -> None: await session.aclose() class AgentState(TypedDict): - input: str + input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] context: Annotated[MyPydanticContextModel, Context(make_context)] @@ -2938,7 +2939,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -2982,7 +2982,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -3051,7 +3050,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", @@ -3115,7 +3113,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), @@ -3157,7 +3154,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -3226,7 +3222,6 @@ async def test_conditional_graph_state() -> None: assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ - "input": "what is weather in sf", "agent_outcome": AgentFinish( return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", From b0b6c0ae7b3b7d49395e89900f143463bc188316 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 6 Aug 2024 19:46:39 -0400 Subject: [PATCH 2/2] lint + comment --- libs/langgraph/langgraph/channels/__init__.py | 2 +- libs/langgraph/langgraph/channels/untracked_value.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index 144056a58..6f9ba2119 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,10 +1,10 @@ +from langgraph.channels.any_value import AnyValue from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue -from langgraph.channels.any_value import AnyValue __all__ = [ "LastValue", diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index 094d6c3ca..989bba35e 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): return False if len(values) != 1 and self.guard: raise InvalidUpdateError( - "EphemeralValue can only receive one value per step." + "UntrackedValue can only receive one value per step." ) self.value = values[-1]