From 9766068896a7e5910bd8d92bc4c3ecee971531d4 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 20 Nov 2024 15:20:53 -0800
Subject: [PATCH] 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