chore: use defaults for channel values

This commit is contained in:
William Fu-Hinthorn
2026-02-18 20:12:54 -08:00
parent c4f5861166
commit a2912cef67
6 changed files with 313 additions and 15 deletions
@@ -101,6 +101,21 @@ 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]
# Check default_factory first (it takes precedence in Pydantic)
if field.default_factory is not None:
return field.default_factory() # type: ignore[call-arg]
# Check if default is set (not PydanticUndefined)
if (
hasattr(field.default, "__class__")
and getattr(field.default.__class__, "__name__", "")
== "PydanticUndefinedType"
):
pass # No default, fall through
else:
return field.default
if dataclasses.is_dataclass(schema):
field_info = next(
(f for f in dataclasses.fields(schema) if f.name == name), None
+17 -8
View File
@@ -48,11 +48,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,10 +68,13 @@ 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 = default
else:
try:
self.value = typ()
except Exception:
self.value = MISSING
def __eq__(self, value: object) -> bool:
return isinstance(value, BinaryOperatorAggregate) and (
@@ -87,13 +96,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}"
+32
View File
@@ -90,6 +90,38 @@ 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_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"]}