Finish

libs/core/langchain_core/runnables/base.py

@@ -537,7 +537,7 @@ def batch_as_completed(
         *,
         return_exceptions: bool = False,
         **kwargs: Optional[Any],
-    ) -> Iterator[Tuple[Input, Output]]:
+    ) -> Iterator[Tuple[int, Output]]:
         """Run invoke in parallel on a list of inputs,
         yielding results as they complete."""
 
@@ -546,7 +546,9 @@ def batch_as_completed(
 
         configs = get_config_list(config, len(inputs))
 
-        def invoke(input: Input, config: RunnableConfig) -> Union[Output, Exception]:
+        def invoke(
+            i: int, input: Input, config: RunnableConfig
+        ) -> Union[Tuple[int, Output], Exception]:
             if return_exceptions:
                 try:
                     out = self.invoke(input, config, **kwargs)
@@ -555,16 +557,16 @@ def invoke(input: Input, config: RunnableConfig) -> Union[Output, Exception]:
             else:
                 out = self.invoke(input, config, **kwargs)
 
-            return (input, out)
+            return (i, out)
 
         if len(inputs) == 1:
-            yield (inputs[0], invoke(inputs[0], configs[0]))
+            yield (0, invoke(inputs[0], configs[0]))
             return
 
         with get_executor_for_config(configs[0]) as executor:
             futures = [
-                executor.submit(invoke, input, config)
-                for input, config in zip(inputs, configs)
+                executor.submit(invoke, i, input, config)
+                for i, (input, config) in enumerate(zip(inputs, configs))
             ]
 
             try:
@@ -617,7 +619,7 @@ async def abatch_as_completed(
         *,
         return_exceptions: bool = False,
         **kwargs: Optional[Any],
-    ) -> AsyncIterator[Tuple[Input, Output]]:
+    ) -> AsyncIterator[Tuple[int, Output]]:
         """Run ainvoke in parallel on a list of inputs,
         yielding results as they complete."""
 
@@ -627,8 +629,8 @@ async def abatch_as_completed(
         configs = get_config_list(config, len(inputs))
 
         async def ainvoke(
-            input: Input, config: RunnableConfig
-        ) -> Union[Output, Exception]:
+            i: int, input: Input, config: RunnableConfig
+        ) -> Union[Tuple[int, Output], Exception]:
             if return_exceptions:
                 try:
                     out = await self.ainvoke(input, config, **kwargs)
@@ -637,9 +639,9 @@ async def ainvoke(
             else:
                 out = await self.ainvoke(input, config, **kwargs)
 
-            return (input, out)
+            return (i, out)
 
-        coros = map(ainvoke, inputs, configs)
+        coros = map(ainvoke, range(len(inputs)), inputs, configs)
 
         for coro in asyncio.as_completed(coros):
             yield await coro
@@ -4229,6 +4231,51 @@ async def abatch(
             **{**self.kwargs, **kwargs},
         )
 
+    def batch_as_completed(
+        self,
+        inputs: List[Input],
+        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
+        *,
+        return_exceptions: bool = False,
+        **kwargs: Optional[Any],
+    ) -> Iterator[Tuple[int, Output]]:
+        if isinstance(config, list):
+            configs = cast(
+                List[RunnableConfig],
+                [self._merge_configs(conf) for conf in config],
+            )
+        else:
+            configs = [self._merge_configs(config) for _ in range(len(inputs))]
+        yield from self.bound.batch_as_completed(
+            inputs,
+            configs,
+            return_exceptions=return_exceptions,
+            **{**self.kwargs, **kwargs},
+        )
+
+    async def abatch_as_completed(
+        self,
+        inputs: List[Input],
+        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
+        *,
+        return_exceptions: bool = False,
+        **kwargs: Optional[Any],
+    ) -> AsyncIterator[Tuple[int, Output]]:
+        if isinstance(config, list):
+            configs = cast(
+                List[RunnableConfig],
+                [self._merge_configs(conf) for conf in config],
+            )
+        else:
+            configs = [self._merge_configs(config) for _ in range(len(inputs))]
+        async for item in self.bound.abatch_as_completed(
+            inputs,
+            configs,
+            return_exceptions=return_exceptions,
+            **{**self.kwargs, **kwargs},
+        ):
+            yield item
+
     def stream(
         self,
         input: Input,

libs/core/tests/unit_tests/runnables/test_runnable.py

@@ -1428,6 +1428,30 @@ async def test_with_config(mocker: MockerFixture) -> None:
 
     spy.reset_mock()
 
+    assert sorted(
+        c
+        for c in fake.with_config(recursion_limit=5).batch_as_completed(
+            ["hello", "wooorld"],
+            [dict(tags=["a-tag"]), dict(metadata={"key": "value"})],
+        )
+    ) == [(0, 5), (1, 7)]
+
+    assert len(spy.call_args_list) == 2
+    for i, call in enumerate(
+        sorted(spy.call_args_list, key=lambda x: 0 if x.args[0] == "hello" else 1)
+    ):
+        assert call.args[0] == ("hello" if i == 0 else "wooorld")
+        if i == 0:
+            assert call.args[1].get("recursion_limit") == 5
+            assert call.args[1].get("tags") == ["a-tag"]
+            assert call.args[1].get("metadata") == {}
+        else:
+            assert call.args[1].get("recursion_limit") == 5
+            assert call.args[1].get("tags") == []
+            assert call.args[1].get("metadata") == {"key": "value"}
+
+    spy.reset_mock()
+
     assert fake.with_config(metadata={"a": "b"}).batch(
         ["hello", "wooorld"], dict(tags=["a-tag"])
     ) == [5, 7]
@@ -1438,6 +1462,15 @@ async def test_with_config(mocker: MockerFixture) -> None:
         assert call.args[1].get("metadata") == {"a": "b"}
     spy.reset_mock()
 
+    assert sorted(
+        c for c in fake.batch_as_completed(["hello", "wooorld"], dict(tags=["a-tag"]))
+    ) == [(0, 5), (1, 7)]
+    assert len(spy.call_args_list) == 2
+    for i, call in enumerate(spy.call_args_list):
+        assert call.args[0] == ("hello" if i == 0 else "wooorld")
+        assert call.args[1].get("tags") == ["a-tag"]
+    spy.reset_mock()
+
     handler = ConsoleCallbackHandler()
     assert (
         await fake.with_config(metadata={"a": "b"}).ainvoke(
@@ -1484,6 +1517,40 @@ async def test_with_config(mocker: MockerFixture) -> None:
             ),
         ),
     ]
+    spy.reset_mock()
+
+    assert sorted(
+        [
+            c
+            async for c in fake.with_config(
+                recursion_limit=5, tags=["c"]
+            ).abatch_as_completed(["hello", "wooorld"], dict(metadata={"key": "value"}))
+        ]
+    ) == [
+        (0, 5),
+        (1, 7),
+    ]
+    assert len(spy.call_args_list) == 2
+    first_call = next(call for call in spy.call_args_list if call.args[0] == "hello")
+    assert first_call == mocker.call(
+        "hello",
+        dict(
+            metadata={"key": "value"},
+            tags=["c"],
+            callbacks=None,
+            recursion_limit=5,
+        ),
+    )
+    second_call = next(call for call in spy.call_args_list if call.args[0] == "wooorld")
+    assert second_call == mocker.call(
+        "wooorld",
+        dict(
+            metadata={"key": "value"},
+            tags=["c"],
+            callbacks=None,
+            recursion_limit=5,
+        ),
+    )
 
 
 async def test_default_method_implementations(mocker: MockerFixture) -> None: