1"""`Runnable` that retries a `Runnable` if it fails."""23from typing import (4    TYPE_CHECKING,5    Any,6    TypeVar,7    cast,8)910from tenacity import (11    AsyncRetrying,12    RetryCallState,13    RetryError,14    Retrying,15    retry_if_exception_type,16    stop_after_attempt,17    wait_exponential_jitter,18)19from typing_extensions import TypedDict, override2021from langchain_core.runnables.base import RunnableBindingBase22from langchain_core.runnables.config import RunnableConfig, patch_config23from langchain_core.runnables.utils import Input, Output2425if TYPE_CHECKING:26    from langchain_core.callbacks.manager import (27        AsyncCallbackManagerForChainRun,28        CallbackManagerForChainRun,29    )3031    T = TypeVar("T", CallbackManagerForChainRun, AsyncCallbackManagerForChainRun)32U = TypeVar("U")333435class ExponentialJitterParams(TypedDict, total=False):36    """Parameters for `tenacity.wait_exponential_jitter`."""3738    initial: float39    """Initial wait."""40    max: float41    """Maximum wait."""42    exp_base: float43    """Base for exponential backoff."""44    jitter: float45    """Random additional wait sampled from random.uniform(0, jitter)."""464748class RunnableRetry(RunnableBindingBase[Input, Output]):  # type: ignore[no-redef]49    """Retry a Runnable if it fails.5051    RunnableRetry can be used to add retry logic to any object52    that subclasses the base Runnable.5354    Such retries are especially useful for network calls that may fail55    due to transient errors.5657    The RunnableRetry is implemented as a RunnableBinding. The easiest58    way to use it is through the `.with_retry()` method on all Runnables.5960    Example:61    Here's an example that uses a RunnableLambda to raise an exception6263        ```python64        import time656667        def foo(input) -> None:68            '''Fake function that raises an exception.'''69            raise ValueError(f"Invoking foo failed. At time {time.time()}")707172        runnable = RunnableLambda(foo)7374        runnable_with_retries = runnable.with_retry(75            retry_if_exception_type=(ValueError,),  # Retry only on ValueError76            wait_exponential_jitter=True,  # Add jitter to the exponential backoff77            stop_after_attempt=2,  # Try twice78            exponential_jitter_params={"initial": 2},  # if desired, customize backoff79        )8081        # The method invocation above is equivalent to the longer form below:8283        runnable_with_retries = RunnableRetry(84            bound=runnable,85            retry_exception_types=(ValueError,),86            max_attempt_number=2,87            wait_exponential_jitter=True,88            exponential_jitter_params={"initial": 2},89        )90        ```9192    This logic can be used to retry any Runnable, including a chain of Runnables,93    but in general it's best practice to keep the scope of the retry as small as94    possible. For example, if you have a chain of Runnables, you should only retry95    the Runnable that is likely to fail, not the entire chain.9697    Example:98        ```python99        from langchain_core.chat_models import ChatOpenAI100        from langchain_core.prompts import PromptTemplate101102        template = PromptTemplate.from_template("tell me a joke about {topic}.")103        model = ChatOpenAI(temperature=0.5)104105        # Good106        chain = template | model.with_retry()107108        # Bad109        chain = template | model110        retryable_chain = chain.with_retry()111        ```112    """113114    retry_exception_types: tuple[type[BaseException], ...] = (Exception,)115    """The exception types to retry on. By default all exceptions are retried.116117    In general you should only retry on exceptions that are likely to be118    transient, such as network errors.119120    Good exceptions to retry are all server errors (5xx) and selected client121    errors (4xx) such as 429 Too Many Requests.122    """123124    wait_exponential_jitter: bool = True125    """Whether to add jitter to the exponential backoff."""126127    exponential_jitter_params: ExponentialJitterParams | None = None128    """Parameters for `tenacity.wait_exponential_jitter`. Namely: `initial`,129    `max`, `exp_base`, and `jitter` (all `float` values).130    """131132    max_attempt_number: int = 3133    """The maximum number of attempts to retry the Runnable."""134135    @property136    def _kwargs_retrying(self) -> dict[str, Any]:137        kwargs: dict[str, Any] = {}138139        if self.max_attempt_number:140            kwargs["stop"] = stop_after_attempt(self.max_attempt_number)141142        if self.wait_exponential_jitter:143            kwargs["wait"] = wait_exponential_jitter(144                **(self.exponential_jitter_params or {})145            )146147        if self.retry_exception_types:148            kwargs["retry"] = retry_if_exception_type(self.retry_exception_types)149150        return kwargs151152    def _sync_retrying(self, **kwargs: Any) -> Retrying:153        return Retrying(**self._kwargs_retrying, **kwargs)154155    def _async_retrying(self, **kwargs: Any) -> AsyncRetrying:156        return AsyncRetrying(**self._kwargs_retrying, **kwargs)157158    @staticmethod159    def _patch_config(160        config: RunnableConfig,161        run_manager: "T",162        retry_state: RetryCallState,163    ) -> RunnableConfig:164        attempt = retry_state.attempt_number165        tag = f"retry:attempt:{attempt}" if attempt > 1 else None166        return patch_config(config, callbacks=run_manager.get_child(tag))167168    def _patch_config_list(169        self,170        config: list[RunnableConfig],171        run_manager: list["T"],172        retry_state: RetryCallState,173    ) -> list[RunnableConfig]:174        return [175            self._patch_config(c, rm, retry_state)176            for c, rm in zip(config, run_manager, strict=False)177        ]178179    def _invoke(180        self,181        input_: Input,182        run_manager: "CallbackManagerForChainRun",183        config: RunnableConfig,184        **kwargs: Any,185    ) -> Output:186        for attempt in self._sync_retrying(reraise=True):187            with attempt:188                result = super().invoke(189                    input_,190                    self._patch_config(config, run_manager, attempt.retry_state),191                    **kwargs,192                )193            if attempt.retry_state.outcome and not attempt.retry_state.outcome.failed:194                attempt.retry_state.set_result(result)195        return result196197    @override198    def invoke(199        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any200    ) -> Output:201        return self._call_with_config(self._invoke, input, config, **kwargs)202203    async def _ainvoke(204        self,205        input_: Input,206        run_manager: "AsyncCallbackManagerForChainRun",207        config: RunnableConfig,208        **kwargs: Any,209    ) -> Output:210        async for attempt in self._async_retrying(reraise=True):211            with attempt:212                result = await super().ainvoke(213                    input_,214                    self._patch_config(config, run_manager, attempt.retry_state),215                    **kwargs,216                )217            if attempt.retry_state.outcome and not attempt.retry_state.outcome.failed:218                attempt.retry_state.set_result(result)219        return result220221    @override222    async def ainvoke(223        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any224    ) -> Output:225        return await self._acall_with_config(self._ainvoke, input, config, **kwargs)226227    def _batch(228        self,229        inputs: list[Input],230        run_manager: list["CallbackManagerForChainRun"],231        config: list[RunnableConfig],232        **kwargs: Any,233    ) -> list[Output | Exception]:234        results_map: dict[int, Output] = {}235236        not_set: list[Output] = []237        result = not_set238        try:239            for attempt in self._sync_retrying():240                with attempt:241                    # Retry for inputs that have not yet succeeded242                    # Determine which original indices remain.243                    remaining_indices = [244                        i for i in range(len(inputs)) if i not in results_map245                    ]246                    if not remaining_indices:247                        break248                    pending_inputs = [inputs[i] for i in remaining_indices]249                    pending_configs = [config[i] for i in remaining_indices]250                    pending_run_managers = [run_manager[i] for i in remaining_indices]251                    # Invoke underlying batch only on remaining elements.252                    result = super().batch(253                        pending_inputs,254                        self._patch_config_list(255                            pending_configs, pending_run_managers, attempt.retry_state256                        ),257                        return_exceptions=True,258                        **kwargs,259                    )260                    # Register the results of the inputs that have succeeded, mapping261                    # back to their original indices.262                    first_exception = None263                    for offset, r in enumerate(result):264                        if isinstance(r, Exception):265                            if not first_exception:266                                first_exception = r267                            continue268                        orig_idx = remaining_indices[offset]269                        results_map[orig_idx] = r270                    # If any exception occurred, raise it, to retry the failed ones271                    if first_exception:272                        raise first_exception273                if (274                    attempt.retry_state.outcome275                    and not attempt.retry_state.outcome.failed276                ):277                    attempt.retry_state.set_result(result)278        except RetryError as e:279            if result is not_set:280                result = cast("list[Output]", [e] * len(inputs))281282        outputs: list[Output | Exception] = []283        for idx in range(len(inputs)):284            if idx in results_map:285                outputs.append(results_map[idx])286            else:287                outputs.append(result.pop(0))288        return outputs289290    @override291    def batch(292        self,293        inputs: list[Input],294        config: RunnableConfig | list[RunnableConfig] | None = None,295        *,296        return_exceptions: bool = False,297        **kwargs: Any,298    ) -> list[Output]:299        return self._batch_with_config(300            self._batch, inputs, config, return_exceptions=return_exceptions, **kwargs301        )302303    async def _abatch(304        self,305        inputs: list[Input],306        run_manager: list["AsyncCallbackManagerForChainRun"],307        config: list[RunnableConfig],308        **kwargs: Any,309    ) -> list[Output | Exception]:310        results_map: dict[int, Output] = {}311312        not_set: list[Output] = []313        result = not_set314        try:315            async for attempt in self._async_retrying():316                with attempt:317                    # Retry for inputs that have not yet succeeded318                    # Determine which original indices remain.319                    remaining_indices = [320                        i for i in range(len(inputs)) if i not in results_map321                    ]322                    if not remaining_indices:323                        break324                    pending_inputs = [inputs[i] for i in remaining_indices]325                    pending_configs = [config[i] for i in remaining_indices]326                    pending_run_managers = [run_manager[i] for i in remaining_indices]327                    result = await super().abatch(328                        pending_inputs,329                        self._patch_config_list(330                            pending_configs, pending_run_managers, attempt.retry_state331                        ),332                        return_exceptions=True,333                        **kwargs,334                    )335                    # Register the results of the inputs that have succeeded, mapping336                    # back to their original indices.337                    first_exception = None338                    for offset, r in enumerate(result):339                        if isinstance(r, Exception):340                            if not first_exception:341                                first_exception = r342                            continue343                        orig_idx = remaining_indices[offset]344                        results_map[orig_idx] = r345                    # If any exception occurred, raise it, to retry the failed ones346                    if first_exception:347                        raise first_exception348                if (349                    attempt.retry_state.outcome350                    and not attempt.retry_state.outcome.failed351                ):352                    attempt.retry_state.set_result(result)353        except RetryError as e:354            if result is not_set:355                result = cast("list[Output]", [e] * len(inputs))356357        outputs: list[Output | Exception] = []358        for idx in range(len(inputs)):359            if idx in results_map:360                outputs.append(results_map[idx])361            else:362                outputs.append(result.pop(0))363        return outputs364365    @override366    async def abatch(367        self,368        inputs: list[Input],369        config: RunnableConfig | list[RunnableConfig] | None = None,370        *,371        return_exceptions: bool = False,372        **kwargs: Any,373    ) -> list[Output]:374        return await self._abatch_with_config(375            self._abatch, inputs, config, return_exceptions=return_exceptions, **kwargs376        )377378    # stream() and transform() are not retried because retrying a stream379    # is not very intuitive.