From 979b73525666a087f4d3f9e4ec7b539789551c85 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 17 Jan 2024 10:32:33 -0800 Subject: [PATCH 1/4] Allow fan-in in stategraph --- langgraph/channels/last_value.py | 7 ++-- langgraph/graph/state.py | 8 ++++- tests/test_pregel.py | 60 ++++++++++++++++++++++++++++++++ tests/test_pregel_async.py | 60 ++++++++++++++++++++++++++++++++ 4 files changed, 131 insertions(+), 4 deletions(-) diff --git a/langgraph/channels/last_value.py b/langgraph/channels/last_value.py index 858943c91..207423539 100644 --- a/langgraph/channels/last_value.py +++ b/langgraph/channels/last_value.py @@ -14,8 +14,9 @@ from langgraph.channels.base import ( class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): """Stores the last value received, can receive at most one value per step.""" - def __init__(self, typ: Type[Value]) -> None: + def __init__(self, typ: Type[Value], guard: bool = True) -> None: self.typ = typ + self.guard = guard @property def ValueType(self) -> Type[Value]: @@ -29,7 +30,7 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): @contextmanager def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ) + empty = self.__class__(self.typ, self.guard) if checkpoint is not None: empty.value = checkpoint try: @@ -43,7 +44,7 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> None: if len(values) == 0: return - if len(values) != 1: + if len(values) != 1 and self.guard: raise InvalidUpdateError("LastValue can only receive one value per step.") self.value = values[-1] diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index fcfab0425..19be561e9 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -70,6 +70,12 @@ class StateGraph(Graph): ) for key, node in self.nodes.items() } + node_inboxes = { + # we can take any value written to channel because all writers + # write the entire state as of that step, which is equal for all writers + f"{key}:inbox": LastValue(Any, guard=False) + for key in self.nodes + } for key in self.nodes: outgoing = outgoing_edges[key] @@ -97,7 +103,7 @@ class StateGraph(Graph): return Pregel( nodes=nodes, - channels=self.channels, + channels={**self.channels, **node_inboxes}, input=f"{START}:inbox", output=END, hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys, diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f9fef277b..b22d1f689 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1381,3 +1381,63 @@ def test_message_graph() -> None: ] }, ] + + +def test_in_one_fan_out_out_one_graph_state() -> None: + def sorted_add(x: list[str], y: list[str]) -> list[str]: + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "retriever_two": {"docs": ["doc3", "doc4"]}, + "retriever_one": {"docs": ["doc1", "doc2"]}, + }, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index db281e7c6..48e63f226 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1435,3 +1435,63 @@ async def test_message_graph() -> None: ] }, ] + + +async def test_in_one_fan_out_out_one_graph_state() -> None: + def sorted_add(x: list[str], y: list[str]) -> list[str]: + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + async def rewrite_query(data: State) -> State: + return {"query": f'query: {data["query"]}'} + + async def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + async def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + async def qa(data: State) -> State: + return {"answer": ",".join(data["docs"])} + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "retriever_one") + workflow.add_edge("rewrite_query", "retriever_two") + workflow.add_edge("retriever_one", "qa") + workflow.add_edge("retriever_two", "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "retriever_two": {"docs": ["doc3", "doc4"]}, + "retriever_one": {"docs": ["doc1", "doc2"]}, + }, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + { + "__end__": { + "query": "query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + } + }, + ] From 692d1ebe02149e440c0588a5ee3c615536e327f7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 17 Jan 2024 13:50:40 -0800 Subject: [PATCH 2/4] WIP --- langgraph/channels/any_value.py | 55 ++++++++++++++++++++++ langgraph/channels/ephemeral_value.py | 67 +++++++++++++++++++++++++++ langgraph/channels/last_value.py | 7 ++- langgraph/graph/state.py | 11 +++-- 4 files changed, 132 insertions(+), 8 deletions(-) create mode 100644 langgraph/channels/any_value.py create mode 100644 langgraph/channels/ephemeral_value.py diff --git a/langgraph/channels/any_value.py b/langgraph/channels/any_value.py new file mode 100644 index 000000000..2cbb2c25a --- /dev/null +++ b/langgraph/channels/any_value.py @@ -0,0 +1,55 @@ +from contextlib import contextmanager +from typing import Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from langgraph.channels.base import BaseChannel, EmptyChannelError, Value + + +class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): + """Stores the last value received, assumes that if multiple values are + received, they are all equal.""" + + def __init__(self, typ: Type[Value]) -> None: + self.typ = typ + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ) + if checkpoint is not None: + empty.value = checkpoint + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> None: + if len(values) == 0: + return + + self.value = values[-1] + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() diff --git a/langgraph/channels/ephemeral_value.py b/langgraph/channels/ephemeral_value.py new file mode 100644 index 000000000..2baa2f461 --- /dev/null +++ b/langgraph/channels/ephemeral_value.py @@ -0,0 +1,67 @@ +from contextlib import contextmanager +from typing import Generator, Generic, Optional, Sequence, Type + +from typing_extensions import Self + +from langgraph.channels.base import ( + BaseChannel, + EmptyChannelError, + InvalidUpdateError, + Value, +) + + +class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): + """Stores the value received in the step immediately preceding, clears after.""" + + def __init__(self, typ: Type[Value], guard: bool = True) -> None: + self.typ = typ + self.guard = guard + + @property + def ValueType(self) -> Type[Value]: + """The type of the value stored in the channel.""" + return self.typ + + @property + def UpdateType(self) -> Type[Value]: + """The type of the update received by the channel.""" + return self.typ + + @contextmanager + def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + empty = self.__class__(self.typ, self.guard) + if checkpoint is not None: + empty.value = checkpoint + try: + yield empty + finally: + try: + del empty.value + except AttributeError: + pass + + def update(self, values: Sequence[Value]) -> None: + if len(values) == 0: + try: + del self.value + except AttributeError: + pass + finally: + return + if len(values) != 1 and self.guard: + raise InvalidUpdateError("LastValue can only receive one value per step.") + + self.value = values[-1] + + def get(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() diff --git a/langgraph/channels/last_value.py b/langgraph/channels/last_value.py index 207423539..858943c91 100644 --- a/langgraph/channels/last_value.py +++ b/langgraph/channels/last_value.py @@ -14,9 +14,8 @@ from langgraph.channels.base import ( class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): """Stores the last value received, can receive at most one value per step.""" - def __init__(self, typ: Type[Value], guard: bool = True) -> None: + def __init__(self, typ: Type[Value]) -> None: self.typ = typ - self.guard = guard @property def ValueType(self) -> Type[Value]: @@ -30,7 +29,7 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): @contextmanager def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: - empty = self.__class__(self.typ, self.guard) + empty = self.__class__(self.typ) if checkpoint is not None: empty.value = checkpoint try: @@ -44,7 +43,7 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): def update(self, values: Sequence[Value]) -> None: if len(values) == 0: return - if len(values) != 1 and self.guard: + if len(values) != 1: raise InvalidUpdateError("LastValue can only receive one value per step.") self.value = values[-1] diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 19be561e9..785e1d05e 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -6,8 +6,10 @@ from typing import Any, Optional, Sequence, Type from langchain_core.runnables import RunnableLambda, RunnablePassthrough from langchain_core.runnables.base import RunnableLike +from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel, InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, Graph @@ -71,11 +73,12 @@ class StateGraph(Graph): for key, node in self.nodes.items() } node_inboxes = { - # we can take any value written to channel because all writers - # write the entire state as of that step, which is equal for all writers - f"{key}:inbox": LastValue(Any, guard=False) + # we take any value written to channel because all writers + # write the entire state as of that step, which is equal for all + f"{key}:inbox": AnyValue(Any) for key in self.nodes } + node_outboxes = {key: EphemeralValue(Any) for key in self.nodes} for key in self.nodes: outgoing = outgoing_edges[key] @@ -103,7 +106,7 @@ class StateGraph(Graph): return Pregel( nodes=nodes, - channels={**self.channels, **node_inboxes}, + channels={**self.channels, **node_inboxes, **node_outboxes}, input=f"{START}:inbox", output=END, hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys, From 2535795f93f369823ab1e8980a5eaa12e7cfd0ac Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 10 Feb 2024 13:17:26 -0800 Subject: [PATCH 3/4] Comment --- langgraph/graph/state.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 785e1d05e..53613bf4a 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -78,7 +78,11 @@ class StateGraph(Graph): f"{key}:inbox": AnyValue(Any) for key in self.nodes } - node_outboxes = {key: EphemeralValue(Any) for key in self.nodes} + node_outboxes = { + # we clear outbox channels after each step + key: EphemeralValue(Any) + for key in self.nodes + } for key in self.nodes: outgoing = outgoing_edges[key] From 0cc6ce3b504e2ebaee8f0fed9f3832d59315874b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 12 Feb 2024 10:01:24 -0800 Subject: [PATCH 4/4] Lint --- langgraph/pregel/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 016f3a35b..d7fadd9e5 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -789,7 +789,7 @@ def _prepare_next_tasks( checkpoint["channel_versions"][chan] > seen[chan] for chan in proc.triggers ): - # If all channels subscribed by this process have been initialized + # If all channels subscribed by this process are not empty try: val: Any = { k: _read_channel( @@ -819,9 +819,11 @@ def _prepare_next_tasks( elif isinstance(proc, ChannelBatch): # If the channel read by this process was updated if checkpoint["channel_versions"][proc.channel] > seen[proc.channel]: - # Here we don't catch EmptyChannelError because the channel - # must be intialized if the previous `if` condition is true - val = channels[proc.channel].get() + # If the channel subscribed by this process is not empty + try: + val = channels[proc.channel].get() + except EmptyChannelError: + continue if proc.key is not None: val = [{proc.key: v} for v in val]