WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-03 00:48:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

Update handling of updates/inputs passed in as pydantic models

Browse Source 
- remove usage of require_at_least_one_of, we shouldn't be enforcing presence of keys in inputs/updates, an empty dict is a valid input/update
- ensure that values that were explcitly set/assigned in pydantic model are saved even if equal to default value
This commit is contained in:
Nuno Campos
2025-03-13 17:31:38 -07:00
parent 6342cd1665
commit 857f3e4a38
6 changed files with 131 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/state.py
+12 -8
View File
@@ -762,23 +762,28 @@ class CompiledStateGraph(CompiledGraph):
else:
updates.extend(_get_updates(i) or ())
return updates
elif get_type_hints(type(input)):
elif (t := type(input)) and get_type_hints(t):
# Pydantic v2
if hasattr(input, "model_fields"):
defaults = {k: v.default for k, v in input.model_fields.items()}
if isinstance(input, BaseModel):
keep: Optional[set[str]] = input.model_fields_set
defaults = {k: v.default for k, v in t.model_fields.items()}
# Pydantic v1
elif hasattr(input, "__fields__"):
defaults = {k: v.default for k, v in input.__fields__.items()}
elif isinstance(input, BaseModelV1):
keep = input.__fields_set__
defaults = {k: v.default for k, v in t.__fields__.items()}
else:
keep = None
defaults = {}
# NOTE: This behavior for Pydantic is somewhat inelegant,
# but we keep around for backwards compatibility
# if input is a Pydantic model, only update values
# that are different from the default values
# that are different from the default values or in the keep set
return [
(k, value)
for k in output_keys
if (value := getattr(input, k, MISSING)) is not MISSING
and value != defaults.get(k)
and (value != defaults.get(k) or (keep is not None and k in keep))
]
else:
msg = create_error_message(
@@ -804,7 +809,6 @@ class CompiledStateGraph(CompiledGraph):
ChannelWrite(
write_entries,
tags=[TAG_HIDDEN],
require_at_least_one_of=output_keys,
),
],
)
libs/langgraph/langgraph/pregel/debug.py
+57 -39
View File
@@ -25,8 +25,10 @@ from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_NS,
ERROR,
INTERRUPT,
MISSING,
NS_END,
NS_SEP,
RETURN,
TAG_HIDDEN,
)
from langgraph.pregel.io import read_channels
@@ -132,7 +134,9 @@ def map_debug_task_results(
"id": task.id,
"name": task.name,
"error": next((w[1] for w in writes if w[0] == ERROR), None),
"result": [w for w in writes if w[0] in stream_channels_list],
"result": [
w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
],
"interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
},
}
@@ -264,49 +268,63 @@ def tasks_w_writes(
) -> tuple[PregelTask, ...]:
"""Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""
pending_writes = pending_writes or []
return tuple(
PregelTask(
task.id,
task.name,
task.path,
next(
(
exc
for tid, n, exc in pending_writes
if tid == task.id and n == ERROR
),
None,
),
tuple(
v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT
),
states.get(task.id) if states else None,
out: list[PregelTask] = []
for task in tasks:
rtn = next(
(
val
for tid, chan, val in pending_writes
if tid == task.id and chan == RETURN
),
MISSING,
)
out.append(
PregelTask(
task.id,
task.name,
task.path,
next(
(
val
for tid, chan, val in pending_writes
if tid == task.id and chan == output_keys
exc
for tid, n, exc in pending_writes
if tid == task.id and n == ERROR
),
None,
)
if isinstance(output_keys, str)
else {
chan: val
for tid, chan, val in pending_writes
if tid == task.id
and (
chan == output_keys
if isinstance(output_keys, str)
else chan in output_keys
),
tuple(
v
for tid, n, v in pending_writes
if tid == task.id and n == INTERRUPT
),
states.get(task.id) if states else None,
(
rtn
if rtn is not MISSING
else next(
(
val
for tid, chan, val in pending_writes
if tid == task.id and chan == output_keys
),
None,
)
}
if isinstance(output_keys, str)
else {
chan: val
for tid, chan, val in pending_writes
if tid == task.id
and (
chan == output_keys
if isinstance(output_keys, str)
else chan in output_keys
)
}
)
if any(
w[0] == task.id and w[1] not in (ERROR, INTERRUPT)
for w in pending_writes
)
else None,
)
if any(
w[0] == task.id and w[1] not in (ERROR, INTERRUPT)
for w in pending_writes
)
else None,
)
for task in tasks
)
return tuple(out)
libs/langgraph/langgraph/pregel/read.py
-1
View File
@@ -201,7 +201,6 @@ class PregelNode(Runnable):
writers[-2] = ChannelWrite(
writes=writers[-2].writes + writers[-1].writes,
tags=writers[-2].tags,
require_at_least_one_of=writers[-2].require_at_least_one_of,
)
writers.pop()
return writers
libs/langgraph/langgraph/pregel/write.py
+2 -13
View File
@@ -49,21 +49,18 @@ class ChannelWrite(RunnableCallable):
writes: list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]]
"""Sequence of write entries or Send objects to write."""
require_at_least_one_of: Optional[Sequence[str]]
"""If defined, at least one of these channels must be written to."""
def __init__(
self,
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
*,
tags: Optional[Sequence[str]] = None,
require_at_least_one_of: Optional[Sequence[str]] = None,
require_at_least_one_of: Optional[Sequence[str]] = None, # ignored
):
super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
self.writes = cast(
list[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], writes
)
self.require_at_least_one_of = require_at_least_one_of
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
@@ -96,7 +93,6 @@ class ChannelWrite(RunnableCallable):
self.do_write(
config,
writes,
self.require_at_least_one_of if input is not None else None,
)
return input
@@ -112,7 +108,6 @@ class ChannelWrite(RunnableCallable):
self.do_write(
config,
writes,
self.require_at_least_one_of if input is not None else None,
)
return input
@@ -120,7 +115,7 @@ class ChannelWrite(RunnableCallable):
def do_write(
config: RunnableConfig,
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
require_at_least_one_of: Optional[Sequence[str]] = None,
require_at_least_one_of: Optional[Sequence[str]] = None, # ignored
) -> None:
# validate
for w in writes:
@@ -151,12 +146,6 @@ class ChannelWrite(RunnableCallable):
tuples.append((w.channel, value))
else:
raise ValueError(f"Invalid write entry: {w}")
# assert required channels
if require_at_least_one_of is not None:
if not {chan for chan, _ in tuples} & set(require_at_least_one_of):
raise InvalidUpdateError(
f"Must write to at least one of {require_at_least_one_of}"
)
write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
write(tuples)
libs/langgraph/tests/test_pregel.py
+60 -35
View File
@@ -2563,7 +2563,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
request: pytest.FixtureRequest,
checkpointer_name: str,
) -> None:
from pydantic.v1 import BaseModel, ValidationError
from pydantic.v1 import BaseModel, Field, ValidationError
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
setup = mocker.Mock()
@@ -2625,6 +2625,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
answer: Optional[str] = None
docs: Optional[list[str]] = None
class UpdateDocs34(BaseModel):
docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"])
def rewrite_query(data: State) -> State:
assert isinstance(data.inner, InnerObject)
return {"query": f"query: {data.query}"}
@@ -2638,7 +2641,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
def retriever_two(data: State) -> State:
time.sleep(0.1)
return {"docs": ["doc3", "doc4"]}
return UpdateDocs34()
def qa(data: State) -> State:
return {"answer": ",".join(data.docs)}
@@ -2734,7 +2737,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
request: pytest.FixtureRequest,
checkpointer_name: str,
) -> None:
from pydantic import BaseModel, ConfigDict, ValidationError
from pydantic import BaseModel, ConfigDict, Field, ValidationError
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
setup = mocker.Mock()
@@ -2787,6 +2790,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
answer: Optional[str] = None
docs: Optional[list[str]] = None
class UpdateDocs34(BaseModel):
docs: list[str] = Field(default_factory=lambda: ["doc3", "doc4"])
class Input(BaseModel):
query: str
inner: InnerObject
@@ -2808,7 +2814,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
def retriever_two(data: State) -> State:
time.sleep(0.1)
return {"docs": ["doc3", "doc4"]}
return UpdateDocs34()
def qa(data: State) -> State:
return {"answer": ",".join(data.docs)}
@@ -5671,37 +5677,6 @@ def test_command_goto_with_static_breakpoints(
assert result == {"foo": "abc|node-1|node-2|node-2"}
def test_nested_graph_state_error_handling():
"""Test error handling when updating state in nested graphs."""
class State(TypedDict):
count: int
def child_node(state: State):
return {"count": state["count"] + 1}
child = StateGraph(State)
child.add_node("child", child_node)
child.add_edge(START, "child")
parent = StateGraph(State)
parent.add_node("child_graph", child.compile())
parent.add_edge(START, "child_graph")
app = parent.compile(checkpointer=MemorySaver())
# Test invalid state update on parent
with pytest.raises(InvalidUpdateError):
app.update_state({"configurable": {"thread_id": "1"}}, {"invalid_key": "value"})
# Test invalid state update on child
with pytest.raises(InvalidUpdateError):
app.update_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}},
{"invalid_key": "value"},
)
def test_parallel_node_execution():
"""Test that parallel nodes execute concurrently."""
@@ -7033,3 +7008,53 @@ def test_interrupt_subgraph_reenter_checkpointer_true(
}
# confirm that we preserve the state values from the previous invocation
assert bar_values == [None, "barbaz", "quxbaz"]
def test_empty_invoke() -> None:
from pydantic import BaseModel
def reducer_merge_dicts(
dict1: dict[Any, Any], dict2: dict[Any, Any]
) -> dict[Any, Any]:
merged = {**dict1, **dict2}
return merged
class SimpleGraphState(BaseModel):
x1: Annotated[list[str], operator.add] = []
x2: Annotated[dict[str, Any], reducer_merge_dicts] = {}
def update_x1_1(state: SimpleGraphState):
print(state)
return {"x1": ["111"]}
def update_x1_2(state: SimpleGraphState):
print(state)
state.x1.append("222")
return {"x1": ["222"]}
def update_x2_1(state: SimpleGraphState):
print(state)
return {"x2": {"111": 111}}
def update_x2_2(state: SimpleGraphState):
print(state)
return {"x2": {"222": 222}}
graph = StateGraph(SimpleGraphState)
graph.add_node("x1_1_node", update_x1_1)
graph.add_node("x1_2_node", update_x1_2)
graph.add_node("x2_1_node", update_x2_1)
graph.add_node("x2_2_node", update_x2_2)
graph.add_edge("x1_1_node", "x1_2_node")
graph.add_edge("x1_2_node", "x2_1_node")
graph.add_edge("x2_1_node", "x2_2_node")
graph.add_edge(START, "x1_1_node")
graph.add_edge("x2_2_node", END)
compiled = graph.compile()
assert compiled.invoke(SimpleGraphState()).get("x2") == {
"111": 111,
"222": 222,
}
libs/langgraph/tests/test_pregel_async.py
-33
View File
@@ -6654,39 +6654,6 @@ async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> N
assert result == {"foo": "abc|node-1|node-2|node-2"}
async def test_nested_graph_state_error_handling():
"""Test error handling when updating state in nested graphs."""
class State(TypedDict):
count: int
def child_node(state: State):
return {"count": state["count"] + 1}
child = StateGraph(State)
child.add_node("child", child_node)
child.add_edge(START, "child")
parent = StateGraph(State)
parent.add_node("child_graph", child.compile())
parent.add_edge(START, "child_graph")
app = parent.compile(checkpointer=MemorySaver())
# Test invalid state update on parent
with pytest.raises(InvalidUpdateError):
await app.aupdate_state(
{"configurable": {"thread_id": "1"}}, {"invalid_key": "value"}
)
# Test invalid state update on child
with pytest.raises(InvalidUpdateError):
await app.aupdate_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child_graph"}},
{"invalid_key": "value"},
)
async def test_parallel_node_execution():
"""Test that parallel nodes execute concurrently."""