core: Use Blockbuster to detect blocking calls in asyncio during tests

libs/core/langchain_core/beta/runnables/context.py

@@ -65,20 +65,11 @@ def _key_from_id(id_: str) -> str:
 
 def _config_with_context(
     config: RunnableConfig,
-    steps: list[Runnable],
+    context_specs: list[tuple[ConfigurableFieldSpec, int]],
     setter: Callable,
     getter: Callable,
     event_cls: Union[type[threading.Event], type[asyncio.Event]],
 ) -> RunnableConfig:
-    if any(k.startswith(CONTEXT_CONFIG_PREFIX) for k in config.get("configurable", {})):
-        return config
-
-    context_specs = [
-        (spec, i)
-        for i, step in enumerate(steps)
-        for spec in step.config_specs
-        if spec.id.startswith(CONTEXT_CONFIG_PREFIX)
-    ]
     grouped_by_key = {
         key: list(group)
         for key, group in groupby(
@@ -121,7 +112,7 @@ def _config_with_context(
     return patch_config(config, configurable=context_funcs)
 
 
-def aconfig_with_context(
+async def aconfig_with_context(
     config: RunnableConfig,
     steps: list[Runnable],
 ) -> RunnableConfig:
@@ -134,7 +125,18 @@ def aconfig_with_context(
     Returns:
         The patched runnable config.
     """
-    return _config_with_context(config, steps, _asetter, _agetter, asyncio.Event)
+    if any(k.startswith(CONTEXT_CONFIG_PREFIX) for k in config.get("configurable", {})):
+        return config
+
+    context_specs = [
+        (spec, i)
+        for i, step in enumerate(steps)
+        for spec in await asyncio.to_thread(getattr, step, "config_specs")
+        if spec.id.startswith(CONTEXT_CONFIG_PREFIX)
+    ]
+    return _config_with_context(
+        config, context_specs, _asetter, _agetter, asyncio.Event
+    )
 
 
 def config_with_context(
@@ -150,7 +152,18 @@ def config_with_context(
     Returns:
         The patched runnable config.
     """
-    return _config_with_context(config, steps, _setter, _getter, threading.Event)
+    if any(k.startswith(CONTEXT_CONFIG_PREFIX) for k in config.get("configurable", {})):
+        return config
+
+    context_specs = [
+        (spec, i)
+        for i, step in enumerate(steps)
+        for spec in step.config_specs
+        if spec.id.startswith(CONTEXT_CONFIG_PREFIX)
+    ]
+    return _config_with_context(
+        config, context_specs, _setter, _getter, threading.Event
+    )
 
 
 @beta()

libs/core/langchain_core/runnables/base.py

@@ -3037,7 +3037,7 @@ async def ainvoke(
         from langchain_core.beta.runnables.context import aconfig_with_context
 
         # setup callbacks and context
-        config = aconfig_with_context(ensure_config(config), self.steps)
+        config = await aconfig_with_context(ensure_config(config), self.steps)
         callback_manager = get_async_callback_manager_for_config(config)
         # start the root run
         run_manager = await callback_manager.on_chain_start(
@@ -3214,7 +3214,7 @@ async def abatch(
 
         # setup callbacks and context
         configs = [
-            aconfig_with_context(c, self.steps)
+            await aconfig_with_context(c, self.steps)
             for c in get_config_list(config, len(inputs))
         ]
         callback_managers = [
@@ -3364,7 +3364,7 @@ async def _atransform(
         from langchain_core.beta.runnables.context import aconfig_with_context
 
         steps = [self.first] + self.middle + [self.last]
-        config = aconfig_with_context(config, self.steps)
+        config = await aconfig_with_context(config, self.steps)
 
         # stream the last steps
         # transform the input stream of each step with the next

libs/core/pyproject.toml

@@ -109,6 +109,7 @@ grandalf = "^0.8"
 responses = "^0.25.0"
 pytest-socket = "^0.7.0"
 pytest-xdist = "^3.6.1"
+blockbuster = "~1.5.9"
 [[tool.poetry.group.test.dependencies.numpy]]
 version = "^1.24.0"
 python = "<3.12"

libs/core/tests/unit_tests/conftest.py

@@ -1,14 +1,38 @@
 """Configuration for unit tests."""
 
-from collections.abc import Sequence
+from collections.abc import Iterator, Sequence
 from importlib import util
 from uuid import UUID
 
 import pytest
+from blockbuster import BlockBuster, blockbuster_ctx
 from pytest import Config, Function, Parser
 from pytest_mock import MockerFixture
 
 
+@pytest.fixture(autouse=True)
+def blockbuster() -> Iterator[BlockBuster]:
+    with blockbuster_ctx() as bb:
+        for func in ["os.stat", "os.path.abspath"]:
+            (
+                bb.functions[func]
+                .can_block_in("langchain_core/_api/internal.py", "is_caller_internal")
+                .can_block_in("langchain_core/runnables/base.py", "__repr__")
+            )
+
+        for func in ["os.stat", "io.TextIOWrapper.read"]:
+            bb.functions[func].can_block_in(
+                "langsmith/client.py", "_default_retry_config"
+            )
+
+        for bb_function in bb.functions.values():
+            bb_function.can_block_in(
+                "freezegun/api.py", "_get_cached_module_attributes"
+            )
+
+        yield bb
+
+
 def pytest_addoption(parser: Parser) -> None:
     """Add custom command line options to pytest."""
     parser.addoption(

libs/core/tests/unit_tests/fake/test_fake_chat_model.py

@@ -191,7 +191,12 @@ async def on_llm_new_token(
     model = GenericFakeChatModel(messages=infinite_cycle)
     tokens: list[str] = []
     # New model
-    results = list(model.stream("meow", {"callbacks": [MyCustomAsyncHandler(tokens)]}))
+    results = [
+        chunk
+        async for chunk in model.astream(
+            "meow", {"callbacks": [MyCustomAsyncHandler(tokens)]}
+        )
+    ]
     assert results == [
         _any_id_ai_message_chunk(content="hello"),
         _any_id_ai_message_chunk(content=" "),

libs/core/tests/unit_tests/language_models/chat_models/test_base.py

@@ -18,6 +18,7 @@
 )
 from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from langchain_core.outputs.llm_result import LLMResult
+from langchain_core.tracers import LogStreamCallbackHandler
 from langchain_core.tracers.base import BaseTracer
 from langchain_core.tracers.context import collect_runs
 from langchain_core.tracers.event_stream import _AstreamEventsCallbackHandler
@@ -304,39 +305,48 @@ def _stream(
 
 
 @pytest.mark.parametrize("disable_streaming", [True, False, "tool_calling"])
-async def test_disable_streaming(
+def test_disable_streaming(
     disable_streaming: Union[bool, Literal["tool_calling"]],
 ) -> None:
     model = StreamingModel(disable_streaming=disable_streaming)
     assert model.invoke([]).content == "invoke"
-    assert (await model.ainvoke([])).content == "invoke"
 
     expected = "invoke" if disable_streaming is True else "stream"
     assert next(model.stream([])).content == expected
-    async for c in model.astream([]):
-        assert c.content == expected
-        break
+    assert (
+        model.invoke([], config={"callbacks": [LogStreamCallbackHandler()]}).content
+        == expected
+    )
+
+    expected = "invoke" if disable_streaming in ("tool_calling", True) else "stream"
+    assert next(model.stream([], tools=[{"type": "function"}])).content == expected
     assert (
         model.invoke(
-            [], config={"callbacks": [_AstreamEventsCallbackHandler()]}
+            [], config={"callbacks": [LogStreamCallbackHandler()]}, tools=[{}]
         ).content
         == expected
     )
+
+
+@pytest.mark.parametrize("disable_streaming", [True, False, "tool_calling"])
+async def test_disable_streaming_async(
+    disable_streaming: Union[bool, Literal["tool_calling"]],
+) -> None:
+    model = StreamingModel(disable_streaming=disable_streaming)
+    assert (await model.ainvoke([])).content == "invoke"
+
+    expected = "invoke" if disable_streaming is True else "stream"
+    async for c in model.astream([]):
+        assert c.content == expected
+        break
     assert (
         await model.ainvoke([], config={"callbacks": [_AstreamEventsCallbackHandler()]})
     ).content == expected
 
     expected = "invoke" if disable_streaming in ("tool_calling", True) else "stream"
-    assert next(model.stream([], tools=[{"type": "function"}])).content == expected
     async for c in model.astream([], tools=[{}]):
         assert c.content == expected
         break
-    assert (
-        model.invoke(
-            [], config={"callbacks": [_AstreamEventsCallbackHandler()]}, tools=[{}]
-        ).content
-        == expected
-    )
     assert (
         await model.ainvoke(
             [], config={"callbacks": [_AstreamEventsCallbackHandler()]}, tools=[{}]
@@ -345,26 +355,31 @@ async def test_disable_streaming(
 
 
 @pytest.mark.parametrize("disable_streaming", [True, False, "tool_calling"])
-async def test_disable_streaming_no_streaming_model(
+def test_disable_streaming_no_streaming_model(
     disable_streaming: Union[bool, Literal["tool_calling"]],
 ) -> None:
     model = NoStreamingModel(disable_streaming=disable_streaming)
     assert model.invoke([]).content == "invoke"
-    assert (await model.ainvoke([])).content == "invoke"
     assert next(model.stream([])).content == "invoke"
-    async for c in model.astream([]):
-        assert c.content == "invoke"
-        break
     assert (
-        model.invoke(
-            [], config={"callbacks": [_AstreamEventsCallbackHandler()]}
-        ).content
+        model.invoke([], config={"callbacks": [LogStreamCallbackHandler()]}).content
         == "invoke"
     )
+    assert next(model.stream([], tools=[{}])).content == "invoke"
+
+
+@pytest.mark.parametrize("disable_streaming", [True, False, "tool_calling"])
+async def test_disable_streaming_no_streaming_model_async(
+    disable_streaming: Union[bool, Literal["tool_calling"]],
+) -> None:
+    model = NoStreamingModel(disable_streaming=disable_streaming)
+    assert (await model.ainvoke([])).content == "invoke"
+    async for c in model.astream([]):
+        assert c.content == "invoke"
+        break
     assert (
         await model.ainvoke([], config={"callbacks": [_AstreamEventsCallbackHandler()]})
     ).content == "invoke"
-    assert next(model.stream([], tools=[{}])).content == "invoke"
     async for c in model.astream([], tools=[{}]):
         assert c.content == "invoke"
         break

libs/core/tests/unit_tests/language_models/chat_models/test_rate_limiting.py

@@ -1,11 +1,20 @@
 import time
 from typing import Optional as Optional
 
+import pytest
+from blockbuster import BlockBuster
+
 from langchain_core.caches import InMemoryCache
 from langchain_core.language_models import GenericFakeChatModel
 from langchain_core.rate_limiters import InMemoryRateLimiter
 
 
+@pytest.fixture(autouse=True)
+def deactivate_blockbuster(blockbuster: BlockBuster) -> None:
+    # Deactivate BlockBuster to not disturb the rate limiter timings
+    blockbuster.deactivate()
+
+
 def test_rate_limit_invoke() -> None:
     """Add rate limiter."""