From 9766068896a7e5910bd8d92bc4c3ecee971531d4 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 15:20:53 -0800 Subject: [PATCH 1/2] lib: For subgraphs / stream modes call stream.put as a callback in the original event loop - This is asynchronous, so we shouldn't use for regular writes to the output stream (ie those from PregelLoop) - For writes from subgraphs / nodes this is fine to use, as we make no guarantees about when those show up anyway --- libs/langgraph/langgraph/pregel/__init__.py | 15 ++++-- libs/langgraph/tests/fake_chat.py | 55 ++++++++++++++++++++- 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 59d2736a0..36150f829 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1806,12 +1806,16 @@ class Pregel(PregelProtocol): # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream.put_nowait) + StreamMessagesHandler( + partial(aioloop.call_soon_threadsafe, stream.put_nowait) + ) ) # set up custom stream mode if "custom" in stream_modes: - config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put_nowait( - ((), "custom", c) + config[CONF][CONFIG_KEY_STREAM_WRITER] = ( + lambda c: aioloop.call_soon_threadsafe( + stream.put_nowait, ((), "custom", c) + ) ) async with AsyncPregelLoop( input, @@ -1838,7 +1842,10 @@ class Pregel(PregelProtocol): ) # enable subgraph streaming if subgraphs: - loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream + loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol( + partial(aioloop.call_soon_threadsafe, stream.put_nowait), + stream_modes, + ) # enable concurrent streaming if subgraphs or "messages" in stream_modes or "custom" in stream_modes: diff --git a/libs/langgraph/tests/fake_chat.py b/libs/langgraph/tests/fake_chat.py index c2a6b9b9e..d4a76ef7c 100644 --- a/libs/langgraph/tests/fake_chat.py +++ b/libs/langgraph/tests/fake_chat.py @@ -1,7 +1,10 @@ import re -from typing import Any, Iterator, List, Optional, cast +from typing import Any, AsyncIterator, Iterator, List, Optional, cast -from langchain_core.callbacks import CallbackManagerForLLMRun +from langchain_core.callbacks import ( + AsyncCallbackManagerForLLMRun, + CallbackManagerForLLMRun, +) from langchain_core.language_models.fake_chat_models import GenericFakeChatModel from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult @@ -84,3 +87,51 @@ class FakeChatModel(GenericFakeChatModel): if run_manager: run_manager.on_llm_new_token("", chunk=chunk) yield chunk + + async def _astream( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> AsyncIterator[ChatGenerationChunk]: + """Stream the output of the model.""" + chat_result = self._generate( + messages, stop=stop, run_manager=run_manager, **kwargs + ) + if not isinstance(chat_result, ChatResult): + raise ValueError( + f"Expected generate to return a ChatResult, " + f"but got {type(chat_result)} instead." + ) + + message = chat_result.generations[0].message + + if not isinstance(message, AIMessage): + raise ValueError( + f"Expected invoke to return an AIMessage, " + f"but got {type(message)} instead." + ) + + content = message.content + + if content: + # Use a regular expression to split on whitespace with a capture group + # so that we can preserve the whitespace in the output. + assert isinstance(content, str) + content_chunks = cast(list[str], re.split(r"(\s)", content)) + + for token in content_chunks: + chunk = ChatGenerationChunk( + message=AIMessageChunk(content=token, id=message.id) + ) + if run_manager: + run_manager.on_llm_new_token(token, chunk=chunk) + yield chunk + else: + args = message.__dict__ + args.pop("type") + chunk = ChatGenerationChunk(message=AIMessageChunk(**args)) + if run_manager: + await run_manager.on_llm_new_token("", chunk=chunk) + yield chunk From a5706627732de5b35398c4765121107870e76caf Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 15:43:46 -0800 Subject: [PATCH 2/2] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 12 +++++++----- libs/langgraph/langgraph/pregel/loop.py | 2 +- libs/langgraph/langgraph/pregel/messages.py | 2 +- libs/langgraph/langgraph/utils/config.py | 4 ++-- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 36150f829..62506142d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -107,6 +107,7 @@ from langgraph.types import ( Checkpointer, LoopProtocol, StateSnapshot, + StreamChunk, StreamMode, ) from langgraph.utils.config import ( @@ -1752,6 +1753,10 @@ class Pregel(PregelProtocol): stream = AsyncQueue() aioloop = asyncio.get_running_loop() + stream_put = cast( + Callable[[StreamChunk], None], + partial(aioloop.call_soon_threadsafe, stream.put_nowait), + ) def output() -> Iterator: while True: @@ -1806,9 +1811,7 @@ class Pregel(PregelProtocol): # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( - StreamMessagesHandler( - partial(aioloop.call_soon_threadsafe, stream.put_nowait) - ) + StreamMessagesHandler(stream_put) ) # set up custom stream mode if "custom" in stream_modes: @@ -1843,8 +1846,7 @@ class Pregel(PregelProtocol): # enable subgraph streaming if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol( - partial(aioloop.call_soon_threadsafe, stream.put_nowait), - stream_modes, + stream_put, stream_modes ) # enable concurrent streaming if subgraphs or "messages" in stream_modes or "custom" in stream_modes: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 6a9b6a95e..2a68b00f2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -110,13 +110,13 @@ from langgraph.types import ( Command, LoopProtocol, PregelExecutableTask, + StreamChunk, StreamProtocol, ) from langgraph.utils.config import patch_configurable V = TypeVar("V") P = ParamSpec("P") -StreamChunk = tuple[tuple[str, ...], str, Any] INPUT_DONE = object() INPUT_RESUMING = object() diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index 2c31de279..989f116dc 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -18,7 +18,7 @@ from langchain_core.outputs import ChatGenerationChunk, LLMResult from langchain_core.tracers._streaming import T, _StreamingCallbackHandler from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM -from langgraph.pregel.loop import StreamChunk +from langgraph.types import StreamChunk Meta = tuple[tuple[str, ...], dict[str, Any]] diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index adb26dd89..5bff9e848 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -1,7 +1,7 @@ import asyncio import sys from collections import ChainMap -from typing import Any, Optional, Sequence +from typing import Any, Optional, Sequence, cast from langchain_core.callbacks import ( AsyncCallbackManager, @@ -281,7 +281,7 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: for k, v in config.items(): if v is not None and k in CONFIG_KEYS: if k == CONF: - empty[k] = v.copy() # type: ignore[literal-required] + empty[k] = cast(dict, v).copy() else: empty[k] = v # type: ignore[literal-required] for k, v in config.items():