mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 15:42:25 +02:00
125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
import asyncio
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
from typing_extensions import TypedDict
|
|
|
|
from saf_python_sdk.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
|
|
from saf_python_sdk.types import Command, Send
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
@dataclass
|
|
class DataClassPayload:
|
|
value: int
|
|
|
|
|
|
class PydanticPayload(BaseModel):
|
|
value: int
|
|
|
|
|
|
class InnerTypedDict(TypedDict):
|
|
flag: bool
|
|
n: int
|
|
|
|
|
|
@dataclass
|
|
class UpdateElisionState:
|
|
x: int
|
|
dc: DataClassPayload
|
|
model: PydanticPayload
|
|
td: InnerTypedDict
|
|
obj: dict[str, int]
|
|
items: list[int]
|
|
|
|
|
|
def _initial_state() -> UpdateElisionState:
|
|
return UpdateElisionState(
|
|
x=0,
|
|
dc=DataClassPayload(0),
|
|
model=PydanticPayload(value=0),
|
|
td={"flag": False, "n": 0},
|
|
obj={"n": 0},
|
|
items=[0],
|
|
)
|
|
|
|
|
|
async def test_noop_slow_update_does_not_override_fast_update() -> None:
|
|
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
|
|
|
|
async def start_node(state: UpdateElisionState) -> Command:
|
|
return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
|
|
|
|
async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
|
|
state.x = 1
|
|
state.dc.value = 1
|
|
state.model.value = 1
|
|
state.td["flag"] = True
|
|
state.td["n"] = 1
|
|
state.obj["n"] = 1
|
|
state.items.append(1)
|
|
return state
|
|
|
|
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
|
await asyncio.sleep(0.1)
|
|
return state
|
|
|
|
graph.add_entry_node(start_node)
|
|
graph.add_node(fast_node)
|
|
graph.add_finish_node(slow_node)
|
|
|
|
compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
|
|
initial_state: UpdateElisionState = _initial_state()
|
|
result: UpdateElisionState = await compiled.ainvoke(initial_state)
|
|
assert result.x == 1
|
|
assert result.dc == DataClassPayload(1)
|
|
assert result.model.value == 1
|
|
assert result.td == {"flag": True, "n": 1}
|
|
assert result.obj == {"n": 1}
|
|
assert result.items == [0, 1]
|
|
|
|
|
|
async def test_changed_slow_update_overrides_fast_update() -> None:
|
|
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
|
|
|
|
async def start_node(state: UpdateElisionState) -> Command:
|
|
return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
|
|
|
|
async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
|
|
state.x = 1
|
|
state.dc.value = 1
|
|
state.model.value = 1
|
|
state.td["flag"] = True
|
|
state.td["n"] = 1
|
|
state.obj["n"] = 1
|
|
state.items.append(1)
|
|
return state
|
|
|
|
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
|
|
await asyncio.sleep(0.1)
|
|
state.x = 2
|
|
state.dc.value = 2
|
|
state.model.value = 2
|
|
state.td["flag"] = False
|
|
state.td["n"] = 2
|
|
state.obj["n"] = 2
|
|
state.items.append(2)
|
|
return state
|
|
|
|
graph.add_entry_node(start_node)
|
|
graph.add_node(fast_node)
|
|
graph.add_finish_node(slow_node)
|
|
|
|
compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
|
|
initial_state: UpdateElisionState = _initial_state()
|
|
result: UpdateElisionState = await compiled.ainvoke(initial_state)
|
|
assert result.x == 2
|
|
assert result.dc == DataClassPayload(2)
|
|
assert result.model.value == 2
|
|
assert result.td == {"flag": False, "n": 2}
|
|
assert result.obj == {"n": 2}
|
|
assert result.items == [0, 1, 2]
|
|
|