Compare commits

...
6 changed files with 356 additions and 21 deletions
+11 -2
View File
@@ -3,10 +3,11 @@ from __future__ import annotations
import dataclasses
import types
import weakref
from collections.abc import Generator, Sequence
from typing import Annotated, Any, Optional, Union, get_origin, get_type_hints
from collections.abc import Callable, Generator, Sequence
from typing import Annotated, Any, Optional, Union, cast, get_origin, get_type_hints
from pydantic import BaseModel
from pydantic_core import PydanticUndefined
from typing_extensions import NotRequired, ReadOnly, Required
from langgraph._internal._typing import MISSING
@@ -101,6 +102,14 @@ def get_field_default(name: str, type_: Any, schema: type[Any]) -> Any:
return ...
# Handle NotRequired[<type>] for earlier versions of python
return None
if isinstance(schema, type) and issubclass(schema, BaseModel):
if name in schema.model_fields:
field = schema.model_fields[name]
if field.default_factory is not None:
factory = cast(Callable[[], Any], field.default_factory)
return factory()
if field.default is not PydanticUndefined:
return field.default
if dataclasses.is_dataclass(schema):
field_info = next(
(f for f in dataclasses.fields(schema) if f.name == name), None
+25 -12
View File
@@ -1,4 +1,5 @@
import collections.abc
import copy
from collections.abc import Callable, Sequence
from typing import Any, Generic
@@ -48,11 +49,17 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
```
"""
__slots__ = ("value", "operator")
__slots__ = ("value", "operator", "default")
def __init__(self, typ: type[Value], operator: Callable[[Value, Value], Value]):
def __init__(
self,
typ: type[Value],
operator: Callable[[Value, Value], Value],
default: Any = MISSING,
):
super().__init__(typ)
self.operator = operator
self.default = default
# special forms from typing or collections.abc are not instantiable
# so we need to replace them with their concrete counterparts
typ = _strip_extras(typ)
@@ -62,17 +69,23 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
typ = set
if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
typ = dict
try:
self.value = typ()
except Exception:
self.value = MISSING
if default is not MISSING:
self.value = copy.deepcopy(default)
else:
try:
self.value = typ()
except Exception:
self.value = MISSING
def __eq__(self, value: object) -> bool:
return isinstance(value, BinaryOperatorAggregate) and (
value.operator is self.operator
if value.operator.__name__ != "<lambda>"
and self.operator.__name__ != "<lambda>"
else True
(
value.operator is self.operator
if value.operator.__name__ != "<lambda>"
and self.operator.__name__ != "<lambda>"
else True
)
and value.default == self.default
)
@property
@@ -87,13 +100,13 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
def copy(self) -> Self:
"""Return a copy of the channel."""
empty = self.__class__(self.typ, self.operator)
empty = self.__class__(self.typ, self.operator, self.default)
empty.key = self.key
empty.value = self.value
return empty
def from_checkpoint(self, checkpoint: Value) -> Self:
empty = self.__class__(self.typ, self.operator)
empty = self.__class__(self.typ, self.operator, self.default)
empty.key = self.key
if checkpoint is not MISSING:
empty.value = checkpoint
+29 -7
View File
@@ -1567,7 +1567,7 @@ def _get_channels(
type_hints = get_type_hints(schema, include_extras=True)
all_keys = {
name: _get_channel(name, typ)
name: _get_channel(name, typ, schema=schema)
for name, typ in type_hints.items()
if name != "__slots__"
}
@@ -1580,18 +1580,30 @@ def _get_channels(
@overload
def _get_channel(
name: str, annotation: Any, *, allow_managed: Literal[False]
name: str,
annotation: Any,
*,
allow_managed: Literal[False],
schema: type[Any] | None = None,
) -> BaseChannel: ...
@overload
def _get_channel(
name: str, annotation: Any, *, allow_managed: Literal[True] = True
name: str,
annotation: Any,
*,
allow_managed: Literal[True] = True,
schema: type[Any] | None = None,
) -> BaseChannel | ManagedValueSpec: ...
def _get_channel(
name: str, annotation: Any, *, allow_managed: bool = True
name: str,
annotation: Any,
*,
allow_managed: bool = True,
schema: type[Any] | None = None,
) -> BaseChannel | ManagedValueSpec:
# Strip out Required and NotRequired wrappers
if hasattr(annotation, "__origin__") and annotation.__origin__ in (
@@ -1607,7 +1619,7 @@ def _get_channel(
elif channel := _is_field_channel(annotation):
channel.key = name
return channel
elif channel := _is_field_binop(annotation):
elif channel := _is_field_binop(annotation, name=name, schema=schema):
channel.key = name
return channel
@@ -1630,7 +1642,12 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
return None
def _is_field_binop(typ: type[Any]) -> BinaryOperatorAggregate | None:
def _is_field_binop(
typ: type[Any],
*,
name: str | None = None,
schema: type[Any] | None = None,
) -> BinaryOperatorAggregate | None:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) >= 1 and callable(meta[-1]):
@@ -1643,7 +1660,12 @@ def _is_field_binop(typ: type[Any]) -> BinaryOperatorAggregate | None:
)
== 2
):
return BinaryOperatorAggregate(typ, meta[-1])
default: Any = MISSING
if name is not None and schema is not None:
field_default = get_field_default(name, typ, schema)
if field_default is not ...:
default = field_default
return BinaryOperatorAggregate(typ, meta[-1], default=default)
else:
raise ValueError(
f"Invalid reducer signature. Expected (a, b) -> c. Got {sig}"
+71
View File
@@ -90,6 +90,77 @@ def test_binop() -> None:
assert channel.get() == 10
def test_binop_with_default() -> None:
# Test that a default value is used instead of typ()
channel = BinaryOperatorAggregate(int, operator.add, default=10).from_checkpoint(
MISSING
)
assert channel.get() == 10
channel.update([5])
assert channel.get() == 15
# Test checkpoint round-trip preserves default for new channels
checkpoint = channel.checkpoint()
restored = BinaryOperatorAggregate(int, operator.add, default=10).from_checkpoint(
checkpoint
)
assert restored.get() == 15
# Test from_checkpoint with MISSING uses default
fresh = BinaryOperatorAggregate(int, operator.add, default=10).from_checkpoint(
MISSING
)
assert fresh.get() == 10
# Test dict default with or_ reducer
channel = BinaryOperatorAggregate(
dict, operator.or_, default={"a": 1}
).from_checkpoint(MISSING)
assert channel.get() == {"a": 1}
channel.update([{"b": 2}])
assert channel.get() == {"a": 1, "b": 2}
def test_binop_with_default_mutable_safety() -> None:
"""Mutable defaults should not be shared across channel instances."""
default = {"a": 1}
ch1 = BinaryOperatorAggregate(dict, operator.or_, default=default).from_checkpoint(
MISSING
)
ch2 = BinaryOperatorAggregate(dict, operator.or_, default=default).from_checkpoint(
MISSING
)
# Mutate ch1's value via a reducer that mutates in-place
def mutating_reducer(a: dict, b: dict) -> dict:
a.update(b)
return a
ch1.operator = mutating_reducer
ch1.update([{"b": 2}])
assert ch1.get() == {"a": 1, "b": 2}
# ch2 should be unaffected
assert ch2.get() == {"a": 1}
# Original default should be unaffected
assert default == {"a": 1}
def test_binop_with_default_multi_invoke() -> None:
"""Defaults should be fresh across multiple from_checkpoint calls."""
template = BinaryOperatorAggregate(dict, operator.or_, default={"a": 1})
# Simulate two separate runs
run1 = template.from_checkpoint(MISSING)
run1.update([{"b": 2}])
assert run1.get() == {"a": 1, "b": 2}
run2 = template.from_checkpoint(MISSING)
assert run2.get() == {"a": 1} # Should NOT see {"b": 2}
def test_untracked_value() -> None:
channel = UntrackedValue(dict).from_checkpoint(MISSING)
assert channel.ValueType is dict
+141
View File
@@ -9059,3 +9059,144 @@ def test_fork_does_not_apply_pending_writes(
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121}
def test_reducer_field_with_pydantic_default() -> None:
"""Test that Annotated reducer fields respect Pydantic Field defaults."""
class State(BaseModel):
query: str
files: Annotated[dict[str, str], operator.or_] = Field(
default_factory=lambda: {"default.txt": "content"}
)
observed_files: list[dict] = []
def node(state: State) -> dict:
observed_files.append(state.files)
return {"files": {"new.txt": "new"}}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
# Invoke without providing files — should use default
result = app.invoke({"query": "test"})
assert observed_files[0] == {"default.txt": "content"}
assert result == {
"query": "test",
"files": {"default.txt": "content", "new.txt": "new"},
}
def test_reducer_field_with_pydantic_default_explicit_value() -> None:
"""Test that an explicit value overrides the default for reducer fields."""
class State(BaseModel):
files: Annotated[dict[str, str], operator.or_] = Field(
default_factory=lambda: {"default.txt": "content"}
)
observed_files: list[dict] = []
def node(state: State) -> dict:
observed_files.append(state.files)
return {}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
# Invoke WITH explicit files — should use provided value, not default
app.invoke({"files": {"custom.txt": "custom"}})
assert observed_files[0] == {"default.txt": "content", "custom.txt": "custom"}
def test_reducer_field_with_dataclass_default() -> None:
"""Test that Annotated reducer fields respect dataclass defaults."""
@dataclass
class State:
query: str
files: Annotated[dict[str, str], operator.or_] = field(
default_factory=lambda: {"default.txt": "content"}
)
observed_files: list[dict] = []
def node(state: State) -> dict:
observed_files.append(state.files)
return {"files": {"new.txt": "new"}}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
result = app.invoke({"query": "test"})
assert observed_files[0] == {"default.txt": "content"}
assert result == {
"query": "test",
"files": {"default.txt": "content", "new.txt": "new"},
}
def test_reducer_field_with_pydantic_none_default() -> None:
"""Test that None default is respected (not confused with 'no default')."""
def _reducer(a: str | None, b: str | None) -> str | None:
if b is not None:
return b
return a
class State(BaseModel):
query: str
data: Annotated[str | None, _reducer] = None
observed: list[Any] = []
def node(state: State) -> dict:
observed.append(state.data)
return {"data": "updated"}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
result = app.invoke({"query": "test"})
assert observed[0] is None # default None was respected
assert result == {"query": "test", "data": "updated"}
def test_reducer_field_with_default_multi_step() -> None:
"""Test that defaults work correctly across multiple graph steps."""
class State(BaseModel):
items: Annotated[list[str], operator.add] = Field(
default_factory=lambda: ["initial"]
)
def step1(state: State) -> dict:
return {"items": ["step1"]}
def step2(state: State) -> dict:
return {"items": ["step2"]}
graph = StateGraph(State)
graph.add_node("step1", step1)
graph.add_node("step2", step2)
graph.set_entry_point("step1")
graph.add_edge("step1", "step2")
graph.set_finish_point("step2")
app = graph.compile()
result = app.invoke({})
# Default ["initial"] + ["step1"] + ["step2"]
assert result == {"items": ["initial", "step1", "step2"]}
+79
View File
@@ -9347,3 +9347,82 @@ async def test_fork_does_not_apply_pending_writes(
# 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121}
async def test_reducer_field_with_pydantic_default() -> None:
"""Test that Annotated reducer fields respect Pydantic Field defaults."""
class State(BaseModel):
query: str
files: Annotated[dict[str, str], operator.or_] = Field(
default_factory=lambda: {"default.txt": "content"}
)
observed_files: list[dict] = []
def node(state: State) -> dict:
observed_files.append(state.files)
return {"files": {"new.txt": "new"}}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
result = await app.ainvoke({"query": "test"})
assert observed_files[0] == {"default.txt": "content"}
assert result == {
"query": "test",
"files": {"default.txt": "content", "new.txt": "new"},
}
async def test_reducer_field_with_pydantic_default_explicit_value() -> None:
"""Test that an explicit value overrides the default for reducer fields."""
class State(BaseModel):
files: Annotated[dict[str, str], operator.or_] = Field(
default_factory=lambda: {"default.txt": "content"}
)
observed_files: list[dict] = []
def node(state: State) -> dict:
observed_files.append(state.files)
return {}
graph = StateGraph(State)
graph.add_node("node", node)
graph.set_entry_point("node")
graph.set_finish_point("node")
app = graph.compile()
await app.ainvoke({"files": {"custom.txt": "custom"}})
assert observed_files[0] == {"default.txt": "content", "custom.txt": "custom"}
async def test_reducer_field_with_default_multi_step() -> None:
"""Test that defaults work correctly across multiple graph steps."""
class State(BaseModel):
items: Annotated[list[str], operator.add] = Field(
default_factory=lambda: ["initial"]
)
def step1(state: State) -> dict:
return {"items": ["step1"]}
def step2(state: State) -> dict:
return {"items": ["step2"]}
graph = StateGraph(State)
graph.add_node("step1", step1)
graph.add_node("step2", step2)
graph.set_entry_point("step1")
graph.add_edge("step1", "step2")
graph.set_finish_point("step2")
app = graph.compile()
result = await app.ainvoke({})
assert result == {"items": ["initial", "step1", "step2"]}