Changes following review

Author: cbornet

libs/core/langchain_core/runnables/base.py

@@ -80,7 +80,6 @@
     indent_lines_after_first,
     is_async_callable,
     is_async_generator,
-    task_with_context,
 )
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
 from langchain_core.tracers.event_stream import (
@@ -3189,7 +3188,7 @@ async def ainvoke(
                         part = functools.partial(step.ainvoke, input_, config, **kwargs)
                     else:
                         part = functools.partial(step.ainvoke, input_, config)
-                    input_ = await task_with_context(part(), context)
+                    input_ = await coro_with_context(part(), context, create_task=True)
             # finish the root run
         except BaseException as e:
             await run_manager.on_chain_error(e)
@@ -3911,8 +3910,8 @@ async def _ainvoke_step(
                 callbacks=run_manager.get_child(f"map:key:{key}"),
             )
             with set_config_context(child_config) as context:
-                return await task_with_context(
-                    step.ainvoke(input_, child_config), context
+                return await coro_with_context(
+                    step.ainvoke(input_, child_config), context, create_task=True
                 )
 
         # gather results from all steps

libs/core/langchain_core/runnables/utils.py

@@ -20,7 +20,7 @@
     TypeVar,
 )
 
-from typing_extensions import override
+from typing_extensions import override, reveal_type
 
 # Re-export create-model for backwards compatibility
 from langchain_core.utils.pydantic import create_model  # noqa: F401
@@ -132,35 +132,23 @@ def asyncio_accepts_context() -> bool:
 _T = TypeVar("_T")
 
 
-def task_with_context(
-    coro: Coroutine[Any, Any, _T], context: Context
-) -> asyncio.Task[_T]:
-    """Create a task with a context.
-
-    Args:
-        coro: The coroutine to create a task for.
-        context: The context to use.
-
-    Returns:
-        The task with the context.
-    """
-    if asyncio_accepts_context():
-        return asyncio.create_task(coro, context=context)  # type: ignore[arg-type,call-arg,unused-ignore]
-    return asyncio.create_task(coro)
-
-
-def coro_with_context(coro: Awaitable[_T], context: Context) -> Awaitable[_T]:
+def coro_with_context(
+    coro: Awaitable[_T], context: Context, *, create_task: bool = False
+) -> Awaitable[_T]:
     """Await a coroutine with a context.
 
     Args:
         coro: The coroutine to await.
         context: The context to use.
+        create_task: Whether to create a task.
 
     Returns:
         The coroutine with the context.
     """
     if asyncio_accepts_context():
         return asyncio.create_task(coro, context=context)  # type: ignore[arg-type,call-arg,unused-ignore]
+    if create_task:
+        return asyncio.create_task(coro)  # type: ignore[arg-type]
     return coro