ContextThreadPoolExecutor Class — langchain Architecture

Source Code

libs/core/langchain_core/runnables/config.py lines 527–576

class ContextThreadPoolExecutor(ThreadPoolExecutor):
    """ThreadPoolExecutor that copies the context to the child thread."""

    def submit(  # type: ignore[override]
        self,
        func: Callable[P, T],
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> Future[T]:
        """Submit a function to the executor.

        Args:
            func: The function to submit.
            *args: The positional arguments to the function.
            **kwargs: The keyword arguments to the function.

        Returns:
            The future for the function.
        """
        return super().submit(
            cast("Callable[..., T]", partial(copy_context().run, func, *args, **kwargs))
        )

    def map(
        self,
        fn: Callable[..., T],
        *iterables: Iterable[Any],
        **kwargs: Any,
    ) -> Iterator[T]:
        """Map a function to multiple iterables.

        Args:
            fn: The function to map.
            *iterables: The iterables to map over.
            timeout: The timeout for the map.
            chunksize: The chunksize for the map.

        Returns:
            The iterator for the mapped function.
        """
        contexts = [copy_context() for _ in range(len(iterables[0]))]  # type: ignore[arg-type]

        def _wrapped_fn(*args: Any) -> T:
            return contexts.pop().run(fn, *args)

        return super().map(
            _wrapped_fn,
            *iterables,
            **kwargs,
        )