Merge pull request #1240 from langchain-ai/nc/6aug/untracked-value

Add UntrackedValue to mark a state key as not checkpointable
This commit is contained in:
Nuno Campos
2024-08-06 17:15:19 -07:00
committed by GitHub
5 changed files with 74 additions and 19 deletions
@@ -1,11 +1,17 @@
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
__all__ = [
"LastValue",
"Topic",
"Context",
"BinaryOperatorAggregate",
"UntrackedValue",
"EphemeralValue",
"AnyValue",
]
@@ -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(
"UntrackedValue 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()
+2
View File
@@ -714,6 +714,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
+2 -12
View File
@@ -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,
@@ -2576,7 +2577,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)]
@@ -2767,7 +2768,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"
),
@@ -2805,7 +2805,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",
@@ -2870,7 +2869,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",
@@ -2928,7 +2926,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"
),
@@ -2966,7 +2963,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",
@@ -3031,7 +3027,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",
@@ -3080,7 +3075,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",),
@@ -3102,7 +3096,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"
),
@@ -3146,7 +3139,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"
),
@@ -3219,7 +3211,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"
),
@@ -3263,7 +3254,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"
),
+2 -7
View File
@@ -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,
@@ -2737,7 +2738,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)]
@@ -2952,7 +2953,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",
@@ -2996,7 +2996,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",
@@ -3065,7 +3064,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",
@@ -3129,7 +3127,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"
),
@@ -3171,7 +3168,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",
@@ -3240,7 +3236,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",