Merge pull request #2491 from langchain-ai/nc/20nov/stream-putnowait-loop

lib: For subgraphs / stream modes call stream.put as a callback in the original event loop

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,12 +1811,14 @@ class Pregel(PregelProtocol):
             # set up messages stream mode
             if "messages" in stream_modes:
                 run_manager.inheritable_handlers.append(
-                    StreamMessagesHandler(stream.put_nowait)
+                    StreamMessagesHandler(stream_put)
                 )
             # 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 +1845,9 @@ class Pregel(PregelProtocol):
                 )
                 # enable subgraph streaming
                 if subgraphs:
-                    loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream
+                    loop.config[CONF][CONFIG_KEY_STREAM] = StreamProtocol(
+                        stream_put, stream_modes
+                    )
                 # enable concurrent streaming
                 if subgraphs or "messages" in stream_modes or "custom" in stream_modes:
 

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()

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]]
 

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():

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