From 91653cf2ab7f71f1f853f6cd974efcfef24b3160 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 7 Jun 2024 17:45:20 -0400 Subject: [PATCH] fix ci --- tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 96603dc0c..3fe2a6973 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1941,12 +1941,12 @@ def test_conditional_entrypoint_to_multiple_state_graph( locations: list[str] results: Annotated[list[str], operator.add] - def get_weather(state: str) -> str: + def get_weather(state: OverallState) -> OverallState: location = state["location"] weather = "sunny" if len(location) > 2 else "cloudy" return {"results": [f"It's {weather} in {location}"]} - def continue_to_weather(state: OverallState) -> str | list[str]: + def continue_to_weather(state: OverallState) -> list[Send]: return [ Send("get_weather", {"location": location}) for location in state["locations"]