langchain-core 0.4.0.dev0__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain_core/runnables/retry.py

@@ -3,9 +3,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Optional,
     TypeVar,
-    Union,
     cast,
 )
 
@@ -35,7 +33,7 @@ U = TypeVar("U")
 
 
 class ExponentialJitterParams(TypedDict, total=False):
-    """Parameters for ``tenacity.wait_exponential_jitter``."""
+    """Parameters for `tenacity.wait_exponential_jitter`."""
 
     initial: float
     """Initial wait."""
@@ -47,7 +45,7 @@ class ExponentialJitterParams(TypedDict, total=False):
     """Random additional wait sampled from random.uniform(0, jitter)."""
 
 
-class RunnableRetry(RunnableBindingBase[Input, Output]):
+class RunnableRetry(RunnableBindingBase[Input, Output]):  # type: ignore[no-redef]
     """Retry a Runnable if it fails.
 
     RunnableRetry can be used to add retry logic to any object
@@ -62,32 +60,34 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
     Example:
     Here's an example that uses a RunnableLambda to raise an exception
 
-        .. code-block:: python
+        ```python
+        import time
 
-            import time
 
-            def foo(input) -> None:
-                '''Fake function that raises an exception.'''
-                raise ValueError(f"Invoking foo failed. At time {time.time()}")
+        def foo(input) -> None:
+            '''Fake function that raises an exception.'''
+            raise ValueError(f"Invoking foo failed. At time {time.time()}")
 
-            runnable = RunnableLambda(foo)
 
-            runnable_with_retries = runnable.with_retry(
-                retry_if_exception_type=(ValueError,), # Retry only on ValueError
-                wait_exponential_jitter=True, # Add jitter to the exponential backoff
-                stop_after_attempt=2, # Try twice
-                exponential_jitter_params={"initial": 2},  # if desired, customize backoff
-            )
+        runnable = RunnableLambda(foo)
 
-            # The method invocation above is equivalent to the longer form below:
+        runnable_with_retries = runnable.with_retry(
+            retry_if_exception_type=(ValueError,),  # Retry only on ValueError
+            wait_exponential_jitter=True,  # Add jitter to the exponential backoff
+            stop_after_attempt=2,  # Try twice
+            exponential_jitter_params={"initial": 2},  # if desired, customize backoff
+        )
 
-            runnable_with_retries = RunnableRetry(
-                bound=runnable,
-                retry_exception_types=(ValueError,),
-                max_attempt_number=2,
-                wait_exponential_jitter=True,
-                exponential_jitter_params={"initial": 2},
-            )
+        # The method invocation above is equivalent to the longer form below:
+
+        runnable_with_retries = RunnableRetry(
+            bound=runnable,
+            retry_exception_types=(ValueError,),
+            max_attempt_number=2,
+            wait_exponential_jitter=True,
+            exponential_jitter_params={"initial": 2},
+        )
+        ```
 
     This logic can be used to retry any Runnable, including a chain of Runnables,
     but in general it's best practice to keep the scope of the retry as small as
@@ -95,23 +95,21 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
     the Runnable that is likely to fail, not the entire chain.
 
     Example:
+        ```python
+        from langchain_core.chat_models import ChatOpenAI
+        from langchain_core.prompts import PromptTemplate
 
-        .. code-block:: python
-
-            from langchain_core.chat_models import ChatOpenAI
-            from langchain_core.prompts import PromptTemplate
+        template = PromptTemplate.from_template("tell me a joke about {topic}.")
+        model = ChatOpenAI(temperature=0.5)
 
-            template = PromptTemplate.from_template("tell me a joke about {topic}.")
-            model = ChatOpenAI(temperature=0.5)
+        # Good
+        chain = template | model.with_retry()
 
-            # Good
-            chain = template | model.with_retry()
-
-            # Bad
-            chain = template | model
-            retryable_chain = chain.with_retry()
-
-    """  # noqa: E501
+        # Bad
+        chain = template | model
+        retryable_chain = chain.with_retry()
+        ```
+    """
 
     retry_exception_types: tuple[type[BaseException], ...] = (Exception,)
     """The exception types to retry on. By default all exceptions are retried.
@@ -126,9 +124,9 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
     wait_exponential_jitter: bool = True
     """Whether to add jitter to the exponential backoff."""
 
-    exponential_jitter_params: Optional[ExponentialJitterParams] = None
-    """Parameters for ``tenacity.wait_exponential_jitter``. Namely: ``initial``,
-    ``max``, ``exp_base``, and ``jitter`` (all float values).
+    exponential_jitter_params: ExponentialJitterParams | None = None
+    """Parameters for `tenacity.wait_exponential_jitter`. Namely: `initial`,
+    `max`, `exp_base`, and `jitter` (all `float` values).
     """
 
     max_attempt_number: int = 3
@@ -174,7 +172,8 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
         retry_state: RetryCallState,
     ) -> list[RunnableConfig]:
         return [
-            self._patch_config(c, rm, retry_state) for c, rm in zip(config, run_manager)
+            self._patch_config(c, rm, retry_state)
+            for c, rm in zip(config, run_manager, strict=False)
         ]
 
     def _invoke(
@@ -197,7 +196,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
 
     @override
     def invoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
+        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
     ) -> Output:
         return self._call_with_config(self._invoke, input, config, **kwargs)
 
@@ -221,7 +220,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
 
     @override
     async def ainvoke(
-        self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
+        self, input: Input, config: RunnableConfig | None = None, **kwargs: Any
     ) -> Output:
         return await self._acall_with_config(self._ainvoke, input, config, **kwargs)
 
@@ -231,34 +230,43 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
         run_manager: list["CallbackManagerForChainRun"],
         config: list[RunnableConfig],
         **kwargs: Any,
-    ) -> list[Union[Output, Exception]]:
+    ) -> list[Output | Exception]:
         results_map: dict[int, Output] = {}
 
-        def pending(iterable: list[U]) -> list[U]:
-            return [item for idx, item in enumerate(iterable) if idx not in results_map]
-
         not_set: list[Output] = []
         result = not_set
         try:
             for attempt in self._sync_retrying():
                 with attempt:
-                    # Get the results of the inputs that have not succeeded yet.
+                    # Retry for inputs that have not yet succeeded
+                    # Determine which original indices remain.
+                    remaining_indices = [
+                        i for i in range(len(inputs)) if i not in results_map
+                    ]
+                    if not remaining_indices:
+                        break
+                    pending_inputs = [inputs[i] for i in remaining_indices]
+                    pending_configs = [config[i] for i in remaining_indices]
+                    pending_run_managers = [run_manager[i] for i in remaining_indices]
+                    # Invoke underlying batch only on remaining elements.
                     result = super().batch(
-                        pending(inputs),
+                        pending_inputs,
                         self._patch_config_list(
-                            pending(config), pending(run_manager), attempt.retry_state
+                            pending_configs, pending_run_managers, attempt.retry_state
                         ),
                         return_exceptions=True,
                         **kwargs,
                     )
-                    # Register the results of the inputs that have succeeded.
+                    # Register the results of the inputs that have succeeded, mapping
+                    # back to their original indices.
                     first_exception = None
-                    for i, r in enumerate(result):
+                    for offset, r in enumerate(result):
                         if isinstance(r, Exception):
                             if not first_exception:
                                 first_exception = r
                             continue
-                        results_map[i] = r
+                        orig_idx = remaining_indices[offset]
+                        results_map[orig_idx] = r
                     # If any exception occurred, raise it, to retry the failed ones
                     if first_exception:
                         raise first_exception
@@ -271,7 +279,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
             if result is not_set:
                 result = cast("list[Output]", [e] * len(inputs))
 
-        outputs: list[Union[Output, Exception]] = []
+        outputs: list[Output | Exception] = []
         for idx in range(len(inputs)):
             if idx in results_map:
                 outputs.append(results_map[idx])
@@ -283,7 +291,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
     def batch(
         self,
         inputs: list[Input],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
         **kwargs: Any,
@@ -298,34 +306,42 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
         run_manager: list["AsyncCallbackManagerForChainRun"],
         config: list[RunnableConfig],
         **kwargs: Any,
-    ) -> list[Union[Output, Exception]]:
+    ) -> list[Output | Exception]:
         results_map: dict[int, Output] = {}
 
-        def pending(iterable: list[U]) -> list[U]:
-            return [item for idx, item in enumerate(iterable) if idx not in results_map]
-
         not_set: list[Output] = []
         result = not_set
         try:
             async for attempt in self._async_retrying():
                 with attempt:
-                    # Get the results of the inputs that have not succeeded yet.
+                    # Retry for inputs that have not yet succeeded
+                    # Determine which original indices remain.
+                    remaining_indices = [
+                        i for i in range(len(inputs)) if i not in results_map
+                    ]
+                    if not remaining_indices:
+                        break
+                    pending_inputs = [inputs[i] for i in remaining_indices]
+                    pending_configs = [config[i] for i in remaining_indices]
+                    pending_run_managers = [run_manager[i] for i in remaining_indices]
                     result = await super().abatch(
-                        pending(inputs),
+                        pending_inputs,
                         self._patch_config_list(
-                            pending(config), pending(run_manager), attempt.retry_state
+                            pending_configs, pending_run_managers, attempt.retry_state
                         ),
                         return_exceptions=True,
                         **kwargs,
                     )
-                    # Register the results of the inputs that have succeeded.
+                    # Register the results of the inputs that have succeeded, mapping
+                    # back to their original indices.
                     first_exception = None
-                    for i, r in enumerate(result):
+                    for offset, r in enumerate(result):
                         if isinstance(r, Exception):
                             if not first_exception:
                                 first_exception = r
                             continue
-                        results_map[i] = r
+                        orig_idx = remaining_indices[offset]
+                        results_map[orig_idx] = r
                     # If any exception occurred, raise it, to retry the failed ones
                     if first_exception:
                         raise first_exception
@@ -338,7 +354,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
             if result is not_set:
                 result = cast("list[Output]", [e] * len(inputs))
 
-        outputs: list[Union[Output, Exception]] = []
+        outputs: list[Output | Exception] = []
         for idx in range(len(inputs)):
             if idx in results_map:
                 outputs.append(results_map[idx])
@@ -350,7 +366,7 @@ class RunnableRetry(RunnableBindingBase[Input, Output]):
     async def abatch(
         self,
         inputs: list[Input],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
         **kwargs: Any,

langchain_core/runnables/router.py

@@ -2,13 +2,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
-    Optional,
-    Union,
     cast,
 )
 
@@ -38,15 +35,12 @@ if TYPE_CHECKING:
 
 
 class RouterInput(TypedDict):
-    """Router input.
-
-    Attributes:
-        key: The key to route on.
-        input: The input to pass to the selected Runnable.
-    """
+    """Router input."""
 
     key: str
+    """The key to route on."""
     input: Any
+    """The input to pass to the selected Runnable."""
 
 
 class RouterRunnable(RunnableSerializable[RouterInput, Output]):
@@ -55,18 +49,16 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     Returns the output of the selected Runnable.
 
     Example:
+        ```python
+        from langchain_core.runnables.router import RouterRunnable
+        from langchain_core.runnables import RunnableLambda
 
-        .. code-block:: python
-
-            from langchain_core.runnables.router import RouterRunnable
-            from langchain_core.runnables import RunnableLambda
-
-            add = RunnableLambda(func=lambda x: x + 1)
-            square = RunnableLambda(func=lambda x: x**2)
-
-            router = RouterRunnable(runnables={"add": add, "square": square})
-            router.invoke({"key": "square", "input": 3})
+        add = RunnableLambda(func=lambda x: x + 1)
+        square = RunnableLambda(func=lambda x: x**2)
 
+        router = RouterRunnable(runnables={"add": add, "square": square})
+        router.invoke({"key": "square", "input": 3})
+        ```
     """
 
     runnables: Mapping[str, Runnable[Any, Output]]
@@ -80,14 +72,14 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
 
     def __init__(
         self,
-        runnables: Mapping[str, Union[Runnable[Any, Output], Callable[[Any], Output]]],
+        runnables: Mapping[str, Runnable[Any, Output] | Callable[[Any], Output]],
     ) -> None:
         """Create a RouterRunnable.
 
         Args:
             runnables: A mapping of keys to Runnables.
         """
-        super().__init__(  # type: ignore[call-arg]
+        super().__init__(
             runnables={key: coerce_to_runnable(r) for key, r in runnables.items()}
         )
 
@@ -98,18 +90,22 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     @classmethod
     @override
     def is_lc_serializable(cls) -> bool:
-        """Return whether this class is serializable."""
+        """Return True as this class is serializable."""
         return True
 
     @classmethod
     @override
     def get_lc_namespace(cls) -> list[str]:
-        """Get the namespace of the langchain object."""
+        """Get the namespace of the LangChain object.
+
+        Returns:
+            `["langchain", "schema", "runnable"]`
+        """
         return ["langchain", "schema", "runnable"]
 
     @override
     def invoke(
-        self, input: RouterInput, config: Optional[RunnableConfig] = None, **kwargs: Any
+        self, input: RouterInput, config: RunnableConfig | None = None, **kwargs: Any
     ) -> Output:
         key = input["key"]
         actual_input = input["input"]
@@ -124,8 +120,8 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     async def ainvoke(
         self,
         input: RouterInput,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> Output:
         key = input["key"]
         actual_input = input["input"]
@@ -140,10 +136,10 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     def batch(
         self,
         inputs: list[RouterInput],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
-        **kwargs: Optional[Any],
+        **kwargs: Any | None,
     ) -> list[Output]:
         if not inputs:
             return []
@@ -156,7 +152,7 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
 
         def invoke(
             runnable: Runnable, input_: Input, config: RunnableConfig
-        ) -> Union[Output, Exception]:
+        ) -> Output | Exception:
             if return_exceptions:
                 try:
                     return runnable.invoke(input_, config, **kwargs)
@@ -177,10 +173,10 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     async def abatch(
         self,
         inputs: list[RouterInput],
-        config: Optional[Union[RunnableConfig, list[RunnableConfig]]] = None,
+        config: RunnableConfig | list[RunnableConfig] | None = None,
         *,
         return_exceptions: bool = False,
-        **kwargs: Optional[Any],
+        **kwargs: Any | None,
     ) -> list[Output]:
         if not inputs:
             return []
@@ -193,7 +189,7 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
 
         async def ainvoke(
             runnable: Runnable, input_: Input, config: RunnableConfig
-        ) -> Union[Output, Exception]:
+        ) -> Output | Exception:
             if return_exceptions:
                 try:
                     return await runnable.ainvoke(input_, config, **kwargs)
@@ -213,8 +209,8 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     def stream(
         self,
         input: RouterInput,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> Iterator[Output]:
         key = input["key"]
         actual_input = input["input"]
@@ -229,8 +225,8 @@ class RouterRunnable(RunnableSerializable[RouterInput, Output]):
     async def astream(
         self,
         input: RouterInput,
-        config: Optional[RunnableConfig] = None,
-        **kwargs: Optional[Any],
+        config: RunnableConfig | None = None,
+        **kwargs: Any | None,
     ) -> AsyncIterator[Output]:
         key = input["key"]
         actual_input = input["input"]