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 dataclasses
import types import types
import weakref import weakref
from collections.abc import Generator, Sequence from collections.abc import Callable, Generator, Sequence
from typing import Annotated, Any, Optional, Union, get_origin, get_type_hints from typing import Annotated, Any, Optional, Union, cast, get_origin, get_type_hints
from pydantic import BaseModel from pydantic import BaseModel
from pydantic_core import PydanticUndefined
from typing_extensions import NotRequired, ReadOnly, Required from typing_extensions import NotRequired, ReadOnly, Required
from langgraph._internal._typing import MISSING from langgraph._internal._typing import MISSING
@@ -101,6 +102,14 @@ def get_field_default(name: str, type_: Any, schema: type[Any]) -> Any:
return ... return ...
# Handle NotRequired[<type>] for earlier versions of python # Handle NotRequired[<type>] for earlier versions of python
return None 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): if dataclasses.is_dataclass(schema):
field_info = next( field_info = next(
(f for f in dataclasses.fields(schema) if f.name == name), None (f for f in dataclasses.fields(schema) if f.name == name), None
+25 -12
View File
@@ -1,4 +1,5 @@
import collections.abc import collections.abc
import copy
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from typing import Any, Generic 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) super().__init__(typ)
self.operator = operator self.operator = operator
self.default = default
# special forms from typing or collections.abc are not instantiable # special forms from typing or collections.abc are not instantiable
# so we need to replace them with their concrete counterparts # so we need to replace them with their concrete counterparts
typ = _strip_extras(typ) typ = _strip_extras(typ)
@@ -62,17 +69,23 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
typ = set typ = set
if typ in (collections.abc.Mapping, collections.abc.MutableMapping): if typ in (collections.abc.Mapping, collections.abc.MutableMapping):
typ = dict typ = dict
try: if default is not MISSING:
self.value = typ() self.value = copy.deepcopy(default)
except Exception: else:
self.value = MISSING try:
self.value = typ()
except Exception:
self.value = MISSING
def __eq__(self, value: object) -> bool: def __eq__(self, value: object) -> bool:
return isinstance(value, BinaryOperatorAggregate) and ( return isinstance(value, BinaryOperatorAggregate) and (
value.operator is self.operator (
if value.operator.__name__ != "<lambda>" value.operator is self.operator
and self.operator.__name__ != "<lambda>" if value.operator.__name__ != "<lambda>"
else True and self.operator.__name__ != "<lambda>"
else True
)
and value.default == self.default
) )
@property @property
@@ -87,13 +100,13 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
def copy(self) -> Self: def copy(self) -> Self:
"""Return a copy of the channel.""" """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.key = self.key
empty.value = self.value empty.value = self.value
return empty return empty
def from_checkpoint(self, checkpoint: Value) -> Self: 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 empty.key = self.key
if checkpoint is not MISSING: if checkpoint is not MISSING:
empty.value = checkpoint empty.value = checkpoint
+29 -7
View File
@@ -1567,7 +1567,7 @@ def _get_channels(
type_hints = get_type_hints(schema, include_extras=True) type_hints = get_type_hints(schema, include_extras=True)
all_keys = { all_keys = {
name: _get_channel(name, typ) name: _get_channel(name, typ, schema=schema)
for name, typ in type_hints.items() for name, typ in type_hints.items()
if name != "__slots__" if name != "__slots__"
} }
@@ -1580,18 +1580,30 @@ def _get_channels(
@overload @overload
def _get_channel( 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: ... ) -> BaseChannel: ...
@overload @overload
def _get_channel( 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: ... ) -> BaseChannel | ManagedValueSpec: ...
def _get_channel( 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: ) -> BaseChannel | ManagedValueSpec:
# Strip out Required and NotRequired wrappers # Strip out Required and NotRequired wrappers
if hasattr(annotation, "__origin__") and annotation.__origin__ in ( if hasattr(annotation, "__origin__") and annotation.__origin__ in (
@@ -1607,7 +1619,7 @@ def _get_channel(
elif channel := _is_field_channel(annotation): elif channel := _is_field_channel(annotation):
channel.key = name channel.key = name
return channel return channel
elif channel := _is_field_binop(annotation): elif channel := _is_field_binop(annotation, name=name, schema=schema):
channel.key = name channel.key = name
return channel return channel
@@ -1630,7 +1642,12 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None:
return 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__"): if hasattr(typ, "__metadata__"):
meta = typ.__metadata__ meta = typ.__metadata__
if len(meta) >= 1 and callable(meta[-1]): if len(meta) >= 1 and callable(meta[-1]):
@@ -1643,7 +1660,12 @@ def _is_field_binop(typ: type[Any]) -> BinaryOperatorAggregate | None:
) )
== 2 == 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: else:
raise ValueError( raise ValueError(
f"Invalid reducer signature. Expected (a, b) -> c. Got {sig}" 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 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: def test_untracked_value() -> None:
channel = UntrackedValue(dict).from_checkpoint(MISSING) channel = UntrackedValue(dict).from_checkpoint(MISSING)
assert channel.ValueType is dict 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 # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 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 # 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 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"]}