From 3b0f27c777fe20e457b0797325fac05958f71583 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 19 Sep 2024 16:46:23 -0700 Subject: [PATCH 1/3] Add stream_mode=custom - adds the ability for nodes (including in subgraphs) to emit chunks directly to the output stream, emitted chunks can have any type - when stream_mode=custom isnt requested by the caller emitted chunks are ignored --- libs/langgraph/langgraph/constants.py | 8 +++- libs/langgraph/langgraph/pregel/__init__.py | 15 ++++++- libs/langgraph/langgraph/pregel/types.py | 7 +++- libs/langgraph/langgraph/utils/runnable.py | 46 ++++++++++++++++----- libs/langgraph/tests/test_pregel.py | 21 ++++++++-- libs/langgraph/tests/test_pregel_async.py | 25 +++++++++-- 6 files changed, 100 insertions(+), 22 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index dd7efd6f7..e8719e664 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,11 +1,13 @@ from dataclasses import dataclass -from typing import Any, Literal +from types import MappingProxyType +from typing import Any, Literal, Mapping INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" CONFIG_KEY_STREAM = "__pregel_stream" +CONFIG_KEY_STREAM_WRITER = "__pregel_stream_writer" CONFIG_KEY_STORE = "__pregel_store" CONFIG_KEY_RESUMING = "__pregel_resuming" CONFIG_KEY_TASK_ID = "__pregel_task_id" @@ -34,6 +36,8 @@ RESERVED = { CONFIG_KEY_READ, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_STREAM, + CONFIG_KEY_STREAM_WRITER, CONFIG_KEY_STORE, CONFIG_KEY_RESUMING, CONFIG_KEY_TASK_ID, @@ -51,6 +55,8 @@ END = "__end__" NS_SEP = "|" NS_END = ":" +EMPTY_MAP: Mapping[str, Any] = MappingProxyType({}) + class Send: """A message or packet to send to a specific node in the graph. diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 9f16ac698..f1a9aa578 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -60,6 +60,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, CONFIG_KEY_STREAM, + CONFIG_KEY_STREAM_WRITER, CONFIG_KEY_TASK_ID, INTERRUPT, NS_END, @@ -1219,6 +1220,11 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): run_manager.inheritable_handlers.append( StreamMessagesHandler(stream.put) ) + # set up custom stream mode + if "custom" in stream_modes: + config["configurable"][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( + ((), "custom", c) + ) with SyncPregelLoop( input, stream=StreamProtocol(stream.put, stream_modes), @@ -1240,7 +1246,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if subgraphs: loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming - if subgraphs or "messages" in stream_modes: + if subgraphs or "messages" in stream_modes or "custom" in stream_modes: # we are careful to have a single waiter live at any one time # because on exit we increment semaphore count by exactly 1 waiter: Optional[concurrent.futures.Future] = None @@ -1435,6 +1441,11 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): run_manager.inheritable_handlers.append( StreamMessagesHandler(stream.put_nowait) ) + # set up custom stream mode + if "custom" in stream_modes: + config["configurable"][CONFIG_KEY_STREAM_WRITER] = ( + lambda c: stream.put_nowait(((), "custom", c)) + ) async with AsyncPregelLoop( input, stream=StreamProtocol(stream.put_nowait, stream_modes), @@ -1456,7 +1467,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): if subgraphs: loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming - if subgraphs or "messages" in stream_modes: + if subgraphs or "messages" in stream_modes or "custom" in stream_modes: def get_waiter() -> asyncio.Task[None]: return aioloop.create_task(stream.wait()) diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 218456085..452b35328 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -107,7 +107,7 @@ class StateSnapshot(NamedTuple): All = Literal["*"] -StreamMode = Literal["values", "updates", "debug", "messages"] +StreamMode = Literal["values", "updates", "debug", "messages", "custom"] """How the stream method should emit outputs. - 'values': Emit all values of the state for each step. @@ -115,4 +115,9 @@ StreamMode = Literal["values", "updates", "debug", "messages"] that were returned by the node(s) **after** each step. - 'debug': Emit debug events for each step. - 'messages': Emit LLM messages token-by-token. +- 'custom': Emit custom output `write: StreamWriter` kwarg of each node. """ + +StreamWriter = Callable[[Any], None] +"""Callable that accepts a single argument and writes it to the output stream. +Only available when using stream_mode="custom".""" diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 56d5d5df4..7471f0008 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -30,10 +30,12 @@ from langchain_core.runnables.config import ( run_in_executor, var_child_runnable_config, ) -from langchain_core.runnables.utils import Input, accepts_config +from langchain_core.runnables.utils import Input from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard +from langgraph.constants import CONFIG_KEY_STREAM_WRITER +from langgraph.pregel.types import StreamWriter from langgraph.utils.config import ( ensure_config, get_async_callback_manager_for_config, @@ -57,6 +59,19 @@ class StrEnum(str, enum.Enum): ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11) +KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( + ( + sys.intern("writer"), + (StreamWriter, inspect.Parameter.empty), + CONFIG_KEY_STREAM_WRITER, + lambda _: None, + ), +) +"""List of kwargs that can be passed to functions, and their corresponding +config keys, default values and type annotations.""" + +VALID_KINDS = (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) + class RunnableCallable(Runnable): """A much simpler version of RunnableLambda that requires sync and async functions.""" @@ -86,15 +101,20 @@ class RunnableCallable(Runnable): except AttributeError: pass self.func = func - if func is not None: - self.func_accepts_config = accepts_config(func) self.afunc = afunc - if afunc is not None: - self.afunc_accepts_config = accepts_config(afunc) self.tags = tags self.kwargs = kwargs self.trace = trace self.recurse = recurse + # check signature + params = inspect.signature(func or afunc).parameters + self.func_accepts_config = "config" in params + self.func_accepts: dict[str, bool] = {} + for kw, typ, _, _ in KWARGS_CONFIG_KEYS: + p = params.get(kw) + self.func_accepts[kw] = ( + p is not None and p.annotation in typ and p.kind in VALID_KINDS + ) def __repr__(self) -> str: repr_args = { @@ -113,11 +133,14 @@ class RunnableCallable(Runnable): "\nEither initialize with a synchronous function or invoke" " via the async API (ainvoke, astream, etc.)" ) + if config is None: + config = ensure_config() kwargs = {**self.kwargs, **kwargs} if self.func_accepts_config: kwargs["config"] = config - if config is None: - config = ensure_config() + for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + if self.func_accepts[kw]: + kwargs[kw] = config["configurable"].get(ck, defv) context = copy_context() if self.trace: callback_manager = get_callback_manager_for_config(config, self.tags) @@ -149,11 +172,14 @@ class RunnableCallable(Runnable): ) -> Any: if not self.afunc: return self.invoke(input, config) - kwargs = {**self.kwargs, **kwargs} - if self.afunc_accepts_config: - kwargs["config"] = config if config is None: config = ensure_config() + kwargs = {**self.kwargs, **kwargs} + if self.func_accepts_config: + kwargs["config"] = config + for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + if self.func_accepts[kw]: + kwargs[kw] = config["configurable"].get(ck, defv) context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f42370fed..62c38caeb 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -70,7 +70,7 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from langgraph.pregel.types import PregelTask +from langgraph.pregel.types import PregelTask, StreamWriter from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS @@ -10440,11 +10440,13 @@ def test_weather_subgraph( class SubGraphState(MessagesState): city: str - def model_node(state: SubGraphState): + def model_node(state: SubGraphState, writer: StreamWriter): + writer(" very") result = weather_model.invoke(state["messages"]) return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} - def weather_node(state: SubGraphState): + def weather_node(state: SubGraphState, writer: StreamWriter): + writer(" good") result = get_weather.invoke({"city": state["city"]}) return {"messages": [{"role": "assistant", "content": result}]} @@ -10479,7 +10481,8 @@ def test_weather_subgraph( ] ) - def router_node(state: RouterState): + def router_node(state: RouterState, writer: StreamWriter): + writer("I'm") system_message = "Classify the incoming query as either about weather or not." messages = [{"role": "system", "content": system_message}] + state["messages"] route = router_model.invoke(messages) @@ -10510,8 +10513,18 @@ def test_weather_subgraph( assert graph.get_graph(xray=1).draw_mermaid() == snapshot config = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + # run with custom output + assert [c for c in graph.stream(inputs, thread2, stream_mode="custom")] == [ + "I'm", + " very", + ] + assert [c for c in graph.stream(None, thread2, stream_mode="custom")] == [ + " good", + ] + # run until interrupt assert [ c diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index fec2aeafa..640658399 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -68,7 +68,7 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from langgraph.pregel.types import PregelTask +from langgraph.pregel.types import PregelTask, StreamWriter from langgraph.store.memory import MemoryStore from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ( @@ -9051,11 +9051,13 @@ async def test_weather_subgraph( class SubGraphState(MessagesState): city: str - def model_node(state: SubGraphState): + def model_node(state: SubGraphState, writer: StreamWriter): + writer(" very") result = weather_model.invoke(state["messages"]) return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]} - def weather_node(state: SubGraphState): + def weather_node(state: SubGraphState, writer: StreamWriter): + writer(" good") result = get_weather.invoke({"city": state["city"]}) return {"messages": [{"role": "assistant", "content": result}]} @@ -9090,7 +9092,8 @@ async def test_weather_subgraph( ] ) - def router_node(state: RouterState): + def router_node(state: RouterState, writer: StreamWriter): + writer("I'm") system_message = "Classify the incoming query as either about weather or not." messages = [{"role": "system", "content": system_message}] + state["messages"] route = router_model.invoke(messages) @@ -9128,8 +9131,22 @@ async def test_weather_subgraph( assert graph.get_graph(xray=1).draw_mermaid() == snapshot config = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + # run with custom output + assert [ + c async for c in graph.astream(inputs, thread2, stream_mode="custom") + ] == [ + "I'm", + " very", + ] + assert [ + c async for c in graph.astream(None, thread2, stream_mode="custom") + ] == [ + " good", + ] + # run until interrupt assert [ c From 1cf117ed9d4d0d2899abe591dbaf64cf159ee635 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 19 Sep 2024 17:09:04 -0700 Subject: [PATCH 2/3] Lint --- libs/langgraph/langgraph/utils/runnable.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 7471f0008..bfc2a627b 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -107,7 +107,9 @@ class RunnableCallable(Runnable): self.trace = trace self.recurse = recurse # check signature - params = inspect.signature(func or afunc).parameters + if func is None and afunc is None: + raise ValueError("At least one of func or afunc must be provided.") + params = inspect.signature(cast(Callable, func or afunc)).parameters self.func_accepts_config = "config" in params self.func_accepts: dict[str, bool] = {} for kw, typ, _, _ in KWARGS_CONFIG_KEYS: From 531890e35a35f9b381d375167a7b104184378153 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 20 Sep 2024 08:49:00 -0700 Subject: [PATCH 3/3] Update types.py --- libs/langgraph/langgraph/pregel/types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 452b35328..d34845483 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -120,4 +120,5 @@ StreamMode = Literal["values", "updates", "debug", "messages", "custom"] StreamWriter = Callable[[Any], None] """Callable that accepts a single argument and writes it to the output stream. -Only available when using stream_mode="custom".""" +Always injected into nodes if requested, +but it's a no-op when not using stream_mode="custom"."""