test_seq_abatch_return_exceptions() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/runnables/test_runnable.py lines 4318–4456

async def test_seq_abatch_return_exceptions(mocker: MockerFixture) -> None:
    class ControlledExceptionRunnable(Runnable[str, str]):
        def __init__(self, fail_starts_with: str) -> None:
            self.fail_starts_with = fail_starts_with

        @override
        def invoke(
            self, input: Any, config: RunnableConfig | None = None, **kwargs: Any
        ) -> Any:
            raise NotImplementedError

        async def _abatch(
            self,
            inputs: list[str],
        ) -> list[str | Exception]:
            outputs: list[str | Exception] = []
            for value in inputs:
                if value.startswith(self.fail_starts_with):
                    outputs.append(
                        ValueError(
                            f"ControlledExceptionRunnable({self.fail_starts_with}) "
                            f"fail for {value}"
                        )
                    )
                else:
                    outputs.append(value + "a")
            return outputs

        async def abatch(
            self,
            inputs: list[str],
            config: RunnableConfig | list[RunnableConfig] | None = None,
            *,
            return_exceptions: bool = False,
            **kwargs: Any,
        ) -> list[str]:
            return await self._abatch_with_config(
                self._abatch,
                inputs,
                config,
                return_exceptions=return_exceptions,
                **kwargs,
            )

    chain = (
        ControlledExceptionRunnable("bux")
        | ControlledExceptionRunnable("bar")
        | ControlledExceptionRunnable("baz")
        | ControlledExceptionRunnable("foo")
    )

    assert isinstance(chain, RunnableSequence)

    # Test abatch
    with pytest.raises(
        ValueError, match=re.escape("ControlledExceptionRunnable(bar) fail for bara")
    ):
        await chain.abatch(["foo", "bar", "baz", "qux"])

    spy = mocker.spy(ControlledExceptionRunnable, "abatch")
    tracer = FakeTracer()
    inputs = ["foo", "bar", "baz", "qux"]
    outputs = await chain.abatch(
        inputs, {"callbacks": [tracer]}, return_exceptions=True
    )
    assert len(outputs) == 4
    assert isinstance(outputs[0], ValueError)
    assert isinstance(outputs[1], ValueError)
    assert isinstance(outputs[2], ValueError)
    assert outputs[3] == "quxaaaa"
    assert spy.call_count == 4
    inputs_to_batch = [c[0][1] for c in spy.call_args_list]
    assert inputs_to_batch == [
        # inputs to sequence step 0
        # same as inputs to sequence.batch()
        ["foo", "bar", "baz", "qux"],
        # inputs to sequence step 1
        # == outputs of sequence step 0 as no exceptions were raised
        ["fooa", "bara", "baza", "quxa"],
        # inputs to sequence step 2
        # 'bar' was dropped as it raised an exception in step 1