diff --git a/haystack/core/component/component.py b/haystack/core/component/component.py index 29453d0ba..08da251e4 100644 --- a/haystack/core/component/component.py +++ b/haystack/core/component/component.py @@ -512,13 +512,18 @@ class _Component: # no decorators here def run(self, value: int): return {"output_1": 1, "output_2": "2"} + + # also no decorators here + async def run_async(self, value: int): + return {"output_1": 1, "output_2": "2"} ``` """ - has_decorator = hasattr(instance.run, "_output_types_cache") - if has_decorator: + has_run_decorator = hasattr(instance.run, "_output_types_cache") + has_run_async_decorator = hasattr(instance, "run_async") and hasattr(instance.run_async, "_output_types_cache") + if has_run_decorator or has_run_async_decorator: raise ComponentError( - "Cannot call `set_output_types` on a component that already has " - "the 'output_types' decorator on its `run` method" + "Cannot call `set_output_types` on a component that already has the 'output_types' decorator on its " + "`run` or `run_async` methods." ) instance.__haystack_output__ = Sockets( diff --git a/releasenotes/notes/fix-set-output-types-async-check-2e5d578a3ccdcc59.yaml b/releasenotes/notes/fix-set-output-types-async-check-2e5d578a3ccdcc59.yaml new file mode 100644 index 000000000..ee5f05e16 --- /dev/null +++ b/releasenotes/notes/fix-set-output-types-async-check-2e5d578a3ccdcc59.yaml @@ -0,0 +1,4 @@ +--- +fixes: + - | + When calling `set_output_types` we now also check that the decorator @component.output_types is not present on the `run_async` method of a Component. Previously we only checked that the Component.run method did not possess the decorator. diff --git a/test/core/component/test_component.py b/test/core/component/test_component.py index ff4145830..8c58d1054 100644 --- a/test/core/component/test_component.py +++ b/test/core/component/test_component.py @@ -328,6 +328,23 @@ def test_output_types_decorator_and_set_output_types(): _ = MockComponent() +def test_output_types_decorator_and_set_output_types_async(): + @component + class MockComponent: + def __init__(self) -> None: + component.set_output_types(self, value=int) + + def run(self, value: int): + return {"value": 1} + + @component.output_types(value=int) + async def run_async(self, value: int): + return {"value": 1} + + with pytest.raises(ComponentError, match="Cannot call `set_output_types`"): + _ = MockComponent() + + def test_output_types_decorator_mismatch_run_async_run(): @component class MockComponent: