chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,576 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component import Component, InputSocket, OutputSocket, component
|
||||
from haystack.core.component.component import _hook_component_init
|
||||
from haystack.core.errors import ComponentError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
|
||||
|
||||
def test_correct_declaration():
|
||||
@component
|
||||
class MockComponent:
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
# Verifies also instantiation works with no issues
|
||||
assert MockComponent()
|
||||
assert component.registry["test_component.MockComponent"] == MockComponent
|
||||
assert isinstance(MockComponent(), Component)
|
||||
assert MockComponent().__haystack_supports_async__ is False # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_correct_declaration_with_async():
|
||||
@component
|
||||
class MockComponent:
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
async def run_async(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
# Verifies also instantiation works with no issues
|
||||
assert MockComponent()
|
||||
assert component.registry["test_component.MockComponent"] == MockComponent
|
||||
assert isinstance(MockComponent(), Component)
|
||||
assert MockComponent().__haystack_supports_async__ is True # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_correct_declaration_with_additional_readonly_property():
|
||||
@component
|
||||
class MockComponent:
|
||||
@property
|
||||
def store(self):
|
||||
return "test_store"
|
||||
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
# Verifies that instantiation works with no issues
|
||||
assert MockComponent()
|
||||
assert component.registry["test_component.MockComponent"] == MockComponent
|
||||
assert MockComponent().store == "test_store"
|
||||
|
||||
|
||||
def test_correct_declaration_with_additional_writable_property():
|
||||
@component
|
||||
class MockComponent:
|
||||
@property
|
||||
def store(self):
|
||||
return "test_store"
|
||||
|
||||
@store.setter
|
||||
def store(self, value):
|
||||
self._store = value
|
||||
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
# Verifies that instantiation works with no issues
|
||||
assert component.registry["test_component.MockComponent"] == MockComponent
|
||||
comp = MockComponent()
|
||||
comp.store = "test_store"
|
||||
assert comp.store == "test_store"
|
||||
|
||||
|
||||
def test_missing_run():
|
||||
with pytest.raises(ComponentError, match=r"must have a 'run\(\)' method"):
|
||||
|
||||
@component
|
||||
class MockComponent: # type: ignore[type-var]
|
||||
def another_method(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
|
||||
def test_async_run_not_async():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run_async(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=r"must be a coroutine"):
|
||||
_ = MockComponent()
|
||||
|
||||
|
||||
def test_async_run_not_coroutine():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self, value: int) -> dict[str, int]: # type: ignore[misc]
|
||||
yield {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=r"must be a coroutine"):
|
||||
_ = MockComponent()
|
||||
|
||||
|
||||
def test_parameters_mismatch_run_and_async_run():
|
||||
err_msg = r"Parameters of 'run' and 'run_async' methods must be the same"
|
||||
|
||||
@component
|
||||
class MockComponentMismatchingInputTypes:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self, value: str) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=err_msg):
|
||||
_ = MockComponentMismatchingInputTypes()
|
||||
|
||||
@component
|
||||
class MockComponentMismatchingInputs:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int, **kwargs: Any) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=err_msg):
|
||||
_ = MockComponentMismatchingInputs()
|
||||
|
||||
@component
|
||||
class MockComponentMismatchingInputOrder:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int, another: str) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self, another: str, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=err_msg):
|
||||
_ = MockComponentMismatchingInputOrder()
|
||||
|
||||
|
||||
def test_set_input_types():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, flag: bool):
|
||||
component.set_input_types(self, value=Any)
|
||||
if flag:
|
||||
component.set_input_type(self, name="another", type=str)
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self, **kwargs):
|
||||
return {"value": 1}
|
||||
|
||||
comp = MockComponent(False)
|
||||
assert comp.__haystack_input__._sockets_dict == {"value": InputSocket("value", Any)} # type: ignore[attr-defined]
|
||||
assert comp.run() == {"value": 1}
|
||||
|
||||
comp = MockComponent(True)
|
||||
assert comp.__haystack_input__._sockets_dict == { # type: ignore[attr-defined]
|
||||
"value": InputSocket("value", Any),
|
||||
"another": InputSocket("another", str),
|
||||
}
|
||||
assert comp.run() == {"value": 1}
|
||||
|
||||
|
||||
def test_set_input_types_no_kwarg():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, flag: bool):
|
||||
if flag:
|
||||
component.set_input_type(self, name="another", type=str)
|
||||
else:
|
||||
component.set_input_types(self, value=Any)
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self, fini: bool) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=r"doesn't have a kwargs parameter"):
|
||||
_ = MockComponent(False)
|
||||
|
||||
with pytest.raises(ComponentError, match=r"doesn't have a kwargs parameter"):
|
||||
_ = MockComponent(True)
|
||||
|
||||
|
||||
def test_set_input_types_overrides_run():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, state: bool):
|
||||
if state:
|
||||
component.set_input_type(self, name="fini", type=str)
|
||||
else:
|
||||
component.set_input_types(self, fini=Any)
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self, fini: bool, **kwargs: Any) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
err_msg = "cannot override the parameters of the 'run' method"
|
||||
with pytest.raises(ComponentError, match=err_msg):
|
||||
_ = MockComponent(False)
|
||||
|
||||
with pytest.raises(ComponentError, match=err_msg):
|
||||
_ = MockComponent(True)
|
||||
|
||||
|
||||
def test_set_input_types_postponed_annotations():
|
||||
# The component HelloUsingFutureAnnotations must live in a different module than the one where the test is defined,
|
||||
# so we can properly set up postponed evaluation of annotations using `from __future__ import annotations`.
|
||||
# For this reason, we define it in haystack.testing.sample_components.future_annotations and import it here.
|
||||
from haystack.testing.sample_components import HelloUsingFutureAnnotations
|
||||
|
||||
assert HelloUsingFutureAnnotations().__haystack_input__._sockets_dict == {"word": InputSocket("word", str)} # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_set_output_types():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self):
|
||||
component.set_output_types(self, value=int)
|
||||
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
comp = MockComponent()
|
||||
assert comp.__haystack_output__._sockets_dict == {"value": OutputSocket("value", int)} # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_output_types_decorator_with_compatible_type():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "MockComponent":
|
||||
return cls()
|
||||
|
||||
comp = MockComponent()
|
||||
assert comp.__haystack_output__._sockets_dict == {"value": OutputSocket("value", int)} # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_output_types_decorator_wrong_method():
|
||||
with pytest.raises(ComponentError):
|
||||
|
||||
@component
|
||||
class MockComponent:
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
def to_dict(self):
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
return cls()
|
||||
|
||||
|
||||
def test_output_types_decorator_and_set_output_types():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self) -> None:
|
||||
component.set_output_types(self, value=int)
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match="Cannot call `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) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self, value: int) -> dict[str, 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:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, str]:
|
||||
return {"value": "1"}
|
||||
|
||||
@component.output_types(value=str)
|
||||
async def run_async(self, value: int) -> dict[str, str]:
|
||||
return {"value": "1"}
|
||||
|
||||
with pytest.raises(ComponentError, match=r"Output type specifications .* must be the same"):
|
||||
_ = MockComponent()
|
||||
|
||||
|
||||
def test_output_types_decorator_missing_async_run():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
async def run_async(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
with pytest.raises(ComponentError, match=r"Output type specifications .* must be the same"):
|
||||
_ = MockComponent()
|
||||
|
||||
|
||||
def test_component_decorator_set_it_as_component():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": 1}
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "MockComponent":
|
||||
return cls()
|
||||
|
||||
comp = MockComponent()
|
||||
assert isinstance(comp, Component)
|
||||
|
||||
|
||||
def test_input_has_default_value():
|
||||
@component
|
||||
class MockComponent:
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int = 42) -> dict[str, int]:
|
||||
return {"value": value}
|
||||
|
||||
comp = MockComponent()
|
||||
assert comp.__haystack_input__._sockets_dict["value"].default_value == 42 # type: ignore[attr-defined]
|
||||
assert not comp.__haystack_input__._sockets_dict["value"].is_mandatory # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_keyword_only_args():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self):
|
||||
component.set_output_types(self, value=int)
|
||||
|
||||
def run(self, *, arg: int) -> dict[str, int]:
|
||||
return {"value": arg}
|
||||
|
||||
comp = MockComponent()
|
||||
component_inputs = {
|
||||
name: {"type": socket.type}
|
||||
for name, socket in comp.__haystack_input__._sockets_dict.items() # type: ignore[attr-defined]
|
||||
}
|
||||
assert component_inputs == {"arg": {"type": int}}
|
||||
|
||||
|
||||
def test_repr():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self):
|
||||
component.set_output_types(self, value=int)
|
||||
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": value}
|
||||
|
||||
comp = MockComponent()
|
||||
assert repr(comp) == f"{object.__repr__(comp)}\nInputs:\n - value: int\nOutputs:\n - value: int"
|
||||
|
||||
|
||||
def test_repr_added_to_pipeline():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self):
|
||||
component.set_output_types(self, value=int)
|
||||
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": value}
|
||||
|
||||
pipe = Pipeline()
|
||||
comp = MockComponent()
|
||||
pipe.add_component("my_component", comp)
|
||||
assert repr(comp) == f"{object.__repr__(comp)}\nmy_component\nInputs:\n - value: int\nOutputs:\n - value: int"
|
||||
|
||||
|
||||
def test_pre_init_hooking():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, pos_arg1, pos_arg2, pos_arg3=None, *, kwarg1=1, kwarg2="string"):
|
||||
self.pos_arg1 = pos_arg1
|
||||
self.pos_arg2 = pos_arg2
|
||||
self.pos_arg3 = pos_arg3
|
||||
self.kwarg1 = kwarg1
|
||||
self.kwarg2 = kwarg2
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
def pre_init_hook(component_class, init_params, expected_params):
|
||||
assert component_class == MockComponent
|
||||
assert init_params == expected_params
|
||||
|
||||
def pre_init_hook_modify(component_class, init_params, expected_params):
|
||||
assert component_class == MockComponent
|
||||
assert init_params == expected_params
|
||||
|
||||
init_params["pos_arg1"] = 2
|
||||
init_params["pos_arg2"] = 0
|
||||
init_params["pos_arg3"] = "modified"
|
||||
init_params["kwarg2"] = "modified string"
|
||||
|
||||
with _hook_component_init(partial(pre_init_hook, expected_params={"pos_arg1": 1, "pos_arg2": 2, "kwarg1": None})):
|
||||
_ = MockComponent(1, 2, kwarg1=None)
|
||||
|
||||
with _hook_component_init(partial(pre_init_hook, expected_params={"pos_arg1": 1, "pos_arg2": 2, "pos_arg3": 0.01})):
|
||||
_ = MockComponent(pos_arg1=1, pos_arg2=2, pos_arg3=0.01)
|
||||
|
||||
with _hook_component_init(
|
||||
partial(pre_init_hook_modify, expected_params={"pos_arg1": 0, "pos_arg2": 1, "pos_arg3": 0.01, "kwarg1": 0})
|
||||
):
|
||||
c = MockComponent(0, 1, pos_arg3=0.01, kwarg1=0)
|
||||
|
||||
assert c.pos_arg1 == 2
|
||||
assert c.pos_arg2 == 0
|
||||
assert c.pos_arg3 == "modified"
|
||||
assert c.kwarg1 == 0
|
||||
assert c.kwarg2 == "modified string"
|
||||
|
||||
|
||||
def test_pre_init_hooking_variadic_positional_args():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, *args, kwarg1=1, kwarg2="string"):
|
||||
self.args = args
|
||||
self.kwarg1 = kwarg1
|
||||
self.kwarg2 = kwarg2
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
def pre_init_hook(component_class, init_params, expected_params):
|
||||
assert component_class == MockComponent
|
||||
assert init_params == expected_params
|
||||
|
||||
c = MockComponent(1, 2, 3, kwarg1=None)
|
||||
assert c.args == (1, 2, 3)
|
||||
assert c.kwarg1 is None
|
||||
assert c.kwarg2 == "string"
|
||||
|
||||
with (
|
||||
pytest.raises(ComponentError),
|
||||
_hook_component_init(partial(pre_init_hook, expected_params={"args": (1, 2), "kwarg1": None})),
|
||||
):
|
||||
_ = MockComponent(1, 2, kwarg1=None)
|
||||
|
||||
|
||||
def test_pre_init_hooking_variadic_kwargs():
|
||||
@component
|
||||
class MockComponent:
|
||||
def __init__(self, pos_arg1, pos_arg2=None, **kwargs):
|
||||
self.pos_arg1 = pos_arg1
|
||||
self.pos_arg2 = pos_arg2
|
||||
self.kwargs = kwargs
|
||||
|
||||
@component.output_types(output_value=int)
|
||||
def run(self, input_value: int) -> dict[str, int]:
|
||||
return {"output_value": input_value}
|
||||
|
||||
def pre_init_hook(component_class, init_params, expected_params):
|
||||
assert component_class == MockComponent
|
||||
assert init_params == expected_params
|
||||
|
||||
with _hook_component_init(
|
||||
partial(pre_init_hook, expected_params={"pos_arg1": 1, "kwarg1": None, "kwarg2": 10, "kwarg3": "string"})
|
||||
):
|
||||
c = MockComponent(1, kwarg1=None, kwarg2=10, kwarg3="string")
|
||||
assert c.pos_arg1 == 1
|
||||
assert c.pos_arg2 is None
|
||||
assert c.kwargs == {"kwarg1": None, "kwarg2": 10, "kwarg3": "string"}
|
||||
|
||||
def pre_init_hook_modify(component_class, init_params, expected_params):
|
||||
assert component_class == MockComponent
|
||||
assert init_params == expected_params
|
||||
|
||||
init_params["pos_arg1"] = 2
|
||||
init_params["pos_arg2"] = 0
|
||||
init_params["some_kwarg"] = "modified string"
|
||||
|
||||
with _hook_component_init(
|
||||
partial(
|
||||
pre_init_hook_modify,
|
||||
expected_params={"pos_arg1": 0, "pos_arg2": 1, "kwarg1": 999, "some_kwarg": "some_value"},
|
||||
)
|
||||
):
|
||||
c = MockComponent(0, 1, kwarg1=999, some_kwarg="some_value")
|
||||
|
||||
assert c.pos_arg1 == 2
|
||||
assert c.pos_arg2 == 0
|
||||
assert c.kwargs == {"kwarg1": 999, "some_kwarg": "modified string"}
|
||||
@@ -0,0 +1,94 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component import component
|
||||
from haystack.core.component.component import ComponentError
|
||||
|
||||
|
||||
@component
|
||||
class ValidComponent:
|
||||
def run(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
|
||||
@component
|
||||
class DifferentParamNameComponent:
|
||||
def run(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, input_text: str) -> dict[str, Any]:
|
||||
return {"result": input_text}
|
||||
|
||||
|
||||
@component
|
||||
class DifferentParamTypeComponent:
|
||||
def run(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, text: list[str]) -> dict[str, Any]:
|
||||
return {"result": text[0]}
|
||||
|
||||
|
||||
@component
|
||||
class DifferentDefaultValueComponent:
|
||||
def run(self, text: str = "default") -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, text: str = "different") -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
|
||||
@component
|
||||
class DifferentParamKindComponent:
|
||||
def run(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, *, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
|
||||
@component
|
||||
class DifferentParamCountComponent:
|
||||
def run(self, text: str) -> dict[str, Any]:
|
||||
return {"result": text}
|
||||
|
||||
async def run_async(self, text: str, extra: str) -> dict[str, Any]:
|
||||
return {"result": text + extra}
|
||||
|
||||
|
||||
def test_valid_signatures():
|
||||
component = ValidComponent()
|
||||
assert component.run("test") == {"result": "test"}
|
||||
|
||||
|
||||
def test_different_param_names():
|
||||
with pytest.raises(ComponentError, match="name mismatch"):
|
||||
DifferentParamNameComponent()
|
||||
|
||||
|
||||
def test_different_param_types():
|
||||
with pytest.raises(ComponentError, match="type mismatch"):
|
||||
DifferentParamTypeComponent()
|
||||
|
||||
|
||||
def test_different_default_values():
|
||||
with pytest.raises(ComponentError, match="default value mismatch"):
|
||||
DifferentDefaultValueComponent()
|
||||
|
||||
|
||||
def test_different_param_kinds():
|
||||
with pytest.raises(ComponentError, match=r"kind \(POSITIONAL, KEYWORD, etc\.\) mismatch: "):
|
||||
DifferentParamKindComponent()
|
||||
|
||||
|
||||
def test_different_param_count():
|
||||
with pytest.raises(ComponentError, match="Different number of parameters"):
|
||||
DifferentParamCountComponent()
|
||||
@@ -0,0 +1,101 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component.sockets import InputSocket, OutputSocket, Sockets
|
||||
from haystack.testing.factory import component_class
|
||||
|
||||
|
||||
class TestSockets:
|
||||
def test_init(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
sockets: dict[str, InputSocket | OutputSocket] = {
|
||||
"input_1": InputSocket("input_1", int),
|
||||
"input_2": InputSocket("input_2", int),
|
||||
}
|
||||
io = Sockets(component=comp, sockets_dict=sockets, sockets_io_type=InputSocket)
|
||||
assert io._component == comp
|
||||
assert "input_1" in io.__dict__
|
||||
assert io.__dict__["input_1"] == comp.__haystack_input__._sockets_dict["input_1"] # type: ignore[attr-defined]
|
||||
assert "input_2" in io.__dict__
|
||||
assert io.__dict__["input_2"] == comp.__haystack_input__._sockets_dict["input_2"] # type: ignore[attr-defined]
|
||||
|
||||
def test_init_with_empty_sockets(self):
|
||||
comp = component_class("SomeComponent")()
|
||||
io = Sockets(component=comp, sockets_dict={}, sockets_io_type=InputSocket)
|
||||
|
||||
assert io._component == comp
|
||||
assert io._sockets_dict == {}
|
||||
|
||||
def test_getattribute(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
io = Sockets(
|
||||
component=comp,
|
||||
sockets_dict=comp.__haystack_input__._sockets_dict, # type: ignore[attr-defined]
|
||||
sockets_io_type=InputSocket,
|
||||
)
|
||||
|
||||
assert io.input_1 == comp.__haystack_input__._sockets_dict["input_1"] # type: ignore[attr-defined]
|
||||
assert io.input_2 == comp.__haystack_input__._sockets_dict["input_2"] # type: ignore[attr-defined]
|
||||
|
||||
def test_getattribute_non_existing_socket(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
io = Sockets(
|
||||
component=comp,
|
||||
sockets_dict=comp.__haystack_input__._sockets_dict, # type: ignore[attr-defined]
|
||||
sockets_io_type=InputSocket,
|
||||
)
|
||||
|
||||
with pytest.raises(AttributeError):
|
||||
io.input_3
|
||||
|
||||
def test_repr(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
io = Sockets(
|
||||
component=comp,
|
||||
sockets_dict=comp.__haystack_input__._sockets_dict, # type: ignore[attr-defined]
|
||||
sockets_io_type=InputSocket,
|
||||
)
|
||||
res = repr(io)
|
||||
assert res == "Inputs:\n - input_1: int\n - input_2: int"
|
||||
|
||||
def test_get(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
io = Sockets(
|
||||
component=comp,
|
||||
sockets_dict=comp.__haystack_input__._sockets_dict, # type: ignore[attr-defined]
|
||||
sockets_io_type=InputSocket,
|
||||
)
|
||||
|
||||
assert io.get("input_1") == comp.__haystack_input__._sockets_dict["input_1"] # type: ignore[attr-defined]
|
||||
assert io.get("input_2") == comp.__haystack_input__._sockets_dict["input_2"] # type: ignore[attr-defined]
|
||||
assert io.get("invalid") is None
|
||||
assert io.get("invalid", InputSocket("input_2", int)) == InputSocket("input_2", int)
|
||||
|
||||
def test_contains(self):
|
||||
comp = component_class("SomeComponent", input_types={"input_1": int, "input_2": int})()
|
||||
io = Sockets(
|
||||
component=comp,
|
||||
sockets_dict=comp.__haystack_input__._sockets_dict, # type: ignore[attr-defined]
|
||||
sockets_io_type=InputSocket,
|
||||
)
|
||||
|
||||
assert "input_1" in io
|
||||
assert "input_2" in io
|
||||
assert "invalid" not in io
|
||||
|
||||
|
||||
def test_input_socket_variadic_without_type_arg_raises_component_error():
|
||||
"""Bare Variadic (no type argument) should raise ComponentError, not IndexError."""
|
||||
from collections.abc import Iterable
|
||||
from typing import Annotated
|
||||
|
||||
from haystack.core.component.types import HAYSTACK_VARIADIC_ANNOTATION
|
||||
from haystack.core.errors import ComponentError
|
||||
|
||||
bare_variadic = Annotated[Iterable, HAYSTACK_VARIADIC_ANNOTATION]
|
||||
|
||||
with pytest.raises(ComponentError, match="must have a type argument"):
|
||||
InputSocket("inputs", bare_variadic)
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
@@ -0,0 +1,56 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.pipeline.breakpoint import HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, load_pipeline_snapshot
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_snapshot_saving(monkeypatch):
|
||||
"""Enable snapshot file saving for these integration tests."""
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def output_directory(tmp_path):
|
||||
"""Provide a temporary directory for snapshot files."""
|
||||
return tmp_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def load_and_resume_pipeline_snapshot():
|
||||
"""Fixture that returns a function to load and resume a pipeline from a snapshot."""
|
||||
|
||||
def _resume(pipeline: Pipeline, output_directory: Path, component_name: str, data: dict | None = None) -> dict:
|
||||
"""
|
||||
Utility function to load and resume pipeline snapshot from a breakpoint file.
|
||||
|
||||
:param pipeline: The pipeline instance to resume
|
||||
:param output_directory: Directory containing the breakpoint files
|
||||
:param component_name: Component name to look for in breakpoint files
|
||||
:param data: Data to pass to the pipeline run (defaults to empty dict)
|
||||
|
||||
:returns:
|
||||
Dict containing the pipeline run results
|
||||
|
||||
:raises:
|
||||
ValueError: If no breakpoint file is found for the given component
|
||||
"""
|
||||
data = data or {}
|
||||
all_files = list(output_directory.glob("*"))
|
||||
|
||||
for full_path in all_files:
|
||||
f_name = Path(full_path).name
|
||||
if str(f_name).startswith(component_name):
|
||||
pipeline_snapshot = load_pipeline_snapshot(full_path)
|
||||
return pipeline.run(data=data, pipeline_snapshot=pipeline_snapshot)
|
||||
|
||||
msg = f"No files found for {component_name} in {output_directory}."
|
||||
raise ValueError(msg)
|
||||
|
||||
return _resume
|
||||
@@ -0,0 +1,83 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import component
|
||||
from haystack.components.builders.answer_builder import AnswerBuilder
|
||||
from haystack.components.joiners import AnswerJoiner
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.core.pipeline.pipeline import Pipeline
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
|
||||
|
||||
@component
|
||||
class FakeChatGenerator:
|
||||
def __init__(self, content: str, model_name: str):
|
||||
self.content = content
|
||||
self.model_name = model_name
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||||
return {"replies": [ChatMessage.from_assistant(self.content)]}
|
||||
|
||||
|
||||
class TestPipelineBreakpoints:
|
||||
@pytest.fixture
|
||||
def answer_join_pipeline(self):
|
||||
"""Creates a pipeline with fake components."""
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("gpt-4o", FakeChatGenerator("GPT-4 response", "gpt-4o"))
|
||||
pipeline.add_component("gpt-3", FakeChatGenerator("GPT-3 response", "gpt-3.5-turbo"))
|
||||
pipeline.add_component("answer_builder_a", AnswerBuilder())
|
||||
pipeline.add_component("answer_builder_b", AnswerBuilder())
|
||||
pipeline.add_component("answer_joiner", AnswerJoiner())
|
||||
pipeline.connect("gpt-4o.replies", "answer_builder_a")
|
||||
pipeline.connect("gpt-3.replies", "answer_builder_b")
|
||||
pipeline.connect("answer_builder_a.answers", "answer_joiner")
|
||||
pipeline.connect("answer_builder_b.answers", "answer_joiner")
|
||||
|
||||
return pipeline
|
||||
|
||||
BREAKPOINT_COMPONENTS = ["gpt-4o", "gpt-3", "answer_builder_a", "answer_builder_b", "answer_joiner"]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_pipeline_breakpoints_answer_joiner(
|
||||
self, answer_join_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
"""
|
||||
Test that an answer joiner pipeline can be executed with breakpoints at each component.
|
||||
"""
|
||||
query = "What's Natural Language Processing?"
|
||||
messages = [
|
||||
ChatMessage.from_system("You are a helpful, respectful and honest assistant. Be super concise."),
|
||||
ChatMessage.from_user(query),
|
||||
]
|
||||
data = {
|
||||
"gpt-4o": {"messages": messages},
|
||||
"gpt-3": {"messages": messages},
|
||||
"answer_builder_a": {"query": query},
|
||||
"answer_builder_b": {"query": query},
|
||||
}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = answer_join_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=answer_join_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
assert result["answer_joiner"]
|
||||
assert len(result["answer_joiner"]["answers"]) == 2
|
||||
assert "GPT-4 response" in [a.data for a in result["answer_joiner"]["answers"]]
|
||||
assert "GPT-3 response" in [a.data for a in result["answer_joiner"]["answers"]]
|
||||
@@ -0,0 +1,90 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import component
|
||||
from haystack.components.converters import OutputAdapter
|
||||
from haystack.components.joiners import BranchJoiner
|
||||
from haystack.components.validators import JsonSchemaValidator
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.core.pipeline.pipeline import Pipeline
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
|
||||
|
||||
@component
|
||||
class FakeChatGenerator:
|
||||
def __init__(self, content: str):
|
||||
self.content = content
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(
|
||||
self, messages: list[ChatMessage], generation_kwargs: dict | None = None, **kwargs: Any
|
||||
) -> dict[str, list[ChatMessage]]:
|
||||
return {"replies": [ChatMessage.from_assistant(self.content)]}
|
||||
|
||||
|
||||
class TestPipelineBreakpoints:
|
||||
@pytest.fixture
|
||||
def branch_joiner_pipeline(self):
|
||||
person_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"first_name": {"type": "string", "pattern": "^[A-Z][a-z]+$"},
|
||||
"last_name": {"type": "string", "pattern": "^[A-Z][a-z]+$"},
|
||||
"nationality": {"type": "string", "enum": ["Italian", "Portuguese", "American"]},
|
||||
},
|
||||
"required": ["first_name", "last_name", "nationality"],
|
||||
}
|
||||
|
||||
content = '{"first_name": "Peter", "last_name": "Parker", "nationality": "American"}'
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("joiner", BranchJoiner(list[ChatMessage]))
|
||||
pipe.add_component("fc_llm", FakeChatGenerator(content))
|
||||
pipe.add_component("validator", JsonSchemaValidator(json_schema=person_schema))
|
||||
pipe.add_component("adapter", OutputAdapter("{{chat_message}}", list[ChatMessage], unsafe=True))
|
||||
|
||||
pipe.connect("adapter", "joiner")
|
||||
pipe.connect("joiner", "fc_llm")
|
||||
pipe.connect("fc_llm.replies", "validator.messages")
|
||||
pipe.connect("validator.validation_error", "joiner")
|
||||
|
||||
return pipe
|
||||
|
||||
BREAKPOINT_COMPONENTS = ["joiner", "fc_llm", "validator", "adapter"]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_pipeline_breakpoints_branch_joiner(
|
||||
self, branch_joiner_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
data = {
|
||||
"fc_llm": {"generation_kwargs": {"response_format": {"type": "json_object"}}},
|
||||
"adapter": {"chat_message": [ChatMessage.from_user("Create JSON from Peter Parker")]},
|
||||
}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = branch_joiner_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=branch_joiner_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
assert result["validator"]
|
||||
valid_json = json.loads(result["validator"]["validated"][0].text)
|
||||
assert valid_json["first_name"] == "Peter"
|
||||
assert valid_json["last_name"] == "Parker"
|
||||
assert valid_json["nationality"] == "American"
|
||||
@@ -0,0 +1,86 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Pipeline, component
|
||||
from haystack.components.builders import ChatPromptBuilder
|
||||
from haystack.components.joiners import ListJoiner
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
|
||||
|
||||
@component
|
||||
class FakeChatGenerator:
|
||||
def __init__(self, response: str):
|
||||
self.response = response
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage], **kwargs: Any) -> dict[str, list[ChatMessage]]:
|
||||
return {"replies": [ChatMessage.from_assistant(self.response)]}
|
||||
|
||||
|
||||
class TestPipelineBreakpoints:
|
||||
@pytest.fixture
|
||||
def list_joiner_pipeline(self):
|
||||
user_message = [ChatMessage.from_user("Give a brief answer the following question: {{query}}")]
|
||||
|
||||
feedback_prompt = """
|
||||
You are given a question and an answer.
|
||||
Your task is to provide a score and a brief feedback on the answer.
|
||||
Question: {{query}}
|
||||
Answer: {{response}}
|
||||
"""
|
||||
feedback_message = [ChatMessage.from_system(feedback_prompt)]
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("prompt_builder", ChatPromptBuilder(template=user_message, required_variables=None))
|
||||
pipe.add_component("llm", FakeChatGenerator("Nuclear physics is the study of atomic nuclei."))
|
||||
pipe.add_component(
|
||||
"feedback_prompt_builder", ChatPromptBuilder(template=feedback_message, required_variables=None)
|
||||
)
|
||||
pipe.add_component("feedback_llm", FakeChatGenerator("Score: 8/10. Concise and accurate."))
|
||||
pipe.add_component("list_joiner", ListJoiner(list[ChatMessage]))
|
||||
|
||||
pipe.connect("prompt_builder.prompt", "llm.messages")
|
||||
pipe.connect("prompt_builder.prompt", "list_joiner")
|
||||
pipe.connect("llm.replies", "list_joiner")
|
||||
pipe.connect("llm.replies", "feedback_prompt_builder.response")
|
||||
pipe.connect("feedback_prompt_builder.prompt", "feedback_llm.messages")
|
||||
pipe.connect("feedback_llm.replies", "list_joiner")
|
||||
|
||||
return pipe
|
||||
|
||||
BREAKPOINT_COMPONENTS = ["prompt_builder", "llm", "feedback_prompt_builder", "feedback_llm", "list_joiner"]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_list_joiner_pipeline(
|
||||
self, list_joiner_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
query = "What is nuclear physics?"
|
||||
data = {
|
||||
"prompt_builder": {"template_variables": {"query": query}},
|
||||
"feedback_prompt_builder": {"template_variables": {"query": query}},
|
||||
}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = list_joiner_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=list_joiner_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
assert result["list_joiner"]
|
||||
assert len(result["list_joiner"]["values"]) == 3
|
||||
@@ -0,0 +1,145 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from haystack import component
|
||||
from haystack.components.builders import ChatPromptBuilder
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.core.pipeline.pipeline import Pipeline
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
|
||||
|
||||
@component
|
||||
class OutputValidator:
|
||||
def __init__(self, pydantic_model: Any):
|
||||
self.pydantic_model = pydantic_model
|
||||
self.iteration_counter = 0
|
||||
|
||||
@component.output_types(valid_replies=list[ChatMessage], invalid_replies=list[ChatMessage], error_message=str)
|
||||
def run(self, replies: list[ChatMessage]) -> dict[str, list[ChatMessage] | str]:
|
||||
self.iteration_counter += 1
|
||||
try:
|
||||
assert replies[0].text is not None
|
||||
output_dict = json.loads(replies[0].text)
|
||||
self.pydantic_model.model_validate(output_dict)
|
||||
return {"valid_replies": replies}
|
||||
except (ValueError, ValidationError) as e:
|
||||
return {"invalid_replies": replies, "error_message": str(e)}
|
||||
|
||||
|
||||
@component
|
||||
class FakeChatGenerator:
|
||||
def __init__(self, response: str):
|
||||
self.response = response
|
||||
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||||
return {"replies": [ChatMessage.from_assistant(self.response)]}
|
||||
|
||||
|
||||
class City(BaseModel):
|
||||
name: str
|
||||
country: str
|
||||
population: int
|
||||
|
||||
|
||||
class CitiesData(BaseModel):
|
||||
cities: list[City]
|
||||
|
||||
|
||||
class TestPipelineBreakpointsLoops:
|
||||
"""
|
||||
This class contains tests for pipelines with validation loops and breakpoints.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def validation_loop_pipeline(self):
|
||||
"""Create a pipeline with validation loops for testing."""
|
||||
prompt_template = [
|
||||
ChatMessage.from_user(
|
||||
"""
|
||||
Create a JSON object from the information present in this passage: {{passage}}.
|
||||
Only use information that is present in the passage. Follow this JSON schema, but only return the
|
||||
actual instances without any additional schema definition:
|
||||
{{schema}}
|
||||
Make sure your response is a dict and not a list.
|
||||
{% if invalid_replies and error_message %}
|
||||
You already created the following output in a previous attempt: {{invalid_replies}}
|
||||
However, this doesn't comply with the format requirements from above and triggered this
|
||||
Python exception: {{error_message}}
|
||||
Correct the output and try again. Just return the corrected output without any extra explanations.
|
||||
{% endif %}
|
||||
"""
|
||||
)
|
||||
]
|
||||
|
||||
response_json = json.dumps(
|
||||
{
|
||||
"cities": [
|
||||
{"name": "Berlin", "country": "Germany", "population": 3850809},
|
||||
{"name": "Paris", "country": "France", "population": 2161000},
|
||||
{"name": "Lisbon", "country": "Portugal", "population": 504718},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
pipeline = Pipeline(max_runs_per_component=5)
|
||||
pipeline.add_component(
|
||||
instance=ChatPromptBuilder(template=prompt_template, required_variables=["passage", "schema"]),
|
||||
name="prompt_builder",
|
||||
)
|
||||
pipeline.add_component(instance=FakeChatGenerator(response=response_json), name="llm")
|
||||
pipeline.add_component(instance=OutputValidator(pydantic_model=CitiesData), name="output_validator")
|
||||
|
||||
pipeline.connect("prompt_builder.prompt", "llm.messages")
|
||||
pipeline.connect("llm.replies", "output_validator")
|
||||
pipeline.connect("output_validator.invalid_replies", "prompt_builder.invalid_replies")
|
||||
pipeline.connect("output_validator.error_message", "prompt_builder.error_message")
|
||||
|
||||
return pipeline
|
||||
|
||||
BREAKPOINT_COMPONENTS = ["prompt_builder", "llm", "output_validator"]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_pipeline_breakpoints_validation_loop(
|
||||
self, validation_loop_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
"""
|
||||
Test that a pipeline with validation loops can be executed with breakpoints at each component.
|
||||
"""
|
||||
data = {"prompt_builder": {"passage": "Berlin, Paris, Lisbon...", "schema": "CitiesData schema"}}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = validation_loop_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=validation_loop_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert "output_validator" in result
|
||||
assert "valid_replies" in result["output_validator"]
|
||||
valid_reply = result["output_validator"]["valid_replies"][0].text
|
||||
valid_json = json.loads(valid_reply)
|
||||
assert "cities" in valid_json
|
||||
assert len(valid_json["cities"]) == 3
|
||||
cities_data = CitiesData.model_validate(valid_json)
|
||||
assert len(cities_data.cities) == 3
|
||||
assert cities_data.cities[0].name == "Berlin"
|
||||
assert cities_data.cities[1].name == "Paris"
|
||||
assert cities_data.cities[2].name == "Lisbon"
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from random import random
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document, Pipeline, component
|
||||
from haystack.components.builders.answer_builder import AnswerBuilder
|
||||
from haystack.components.builders.prompt_builder import PromptBuilder
|
||||
from haystack.components.joiners import DocumentJoiner
|
||||
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
|
||||
|
||||
@component
|
||||
class FakeEmbedder:
|
||||
@component.output_types(documents=list[Document], embedding=list[float])
|
||||
def run(self, text: str) -> dict[str, list[Document] | list[float]]:
|
||||
return {"embedding": [random() for _ in range(100)]}
|
||||
|
||||
|
||||
@component
|
||||
class FakeRanker:
|
||||
@component.output_types(documents=list[Document])
|
||||
def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
|
||||
for i, doc in enumerate(documents):
|
||||
doc.score = 1.0 / (i + 1)
|
||||
return {"documents": sorted(documents, key=lambda x: x.score or 0, reverse=True)[:top_k]}
|
||||
|
||||
|
||||
@component
|
||||
class FakeGenerator:
|
||||
@component.output_types(replies=list[str], meta=list[dict[str, Any]])
|
||||
def run(self, prompt: str) -> dict[str, list[str | dict[str, Any]]]:
|
||||
return {"replies": ["Mark lives in Berlin."], "meta": [{"model": "fake"}]}
|
||||
|
||||
|
||||
class TestPipelineBreakpoints:
|
||||
"""
|
||||
This class contains tests for pipelines with breakpoints.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def document_store(self):
|
||||
"""Create and populate a document store for testing."""
|
||||
documents = [
|
||||
Document(content="My name is Jean and I live in Paris."),
|
||||
Document(content="My name is Mark and I live in Berlin."),
|
||||
Document(content="My name is Giorgio and I live in Rome."),
|
||||
]
|
||||
ds = InMemoryDocumentStore()
|
||||
# Add embeddings
|
||||
for doc in documents:
|
||||
doc.embedding = [random() for _ in range(100)]
|
||||
ds.write_documents(documents)
|
||||
return ds
|
||||
|
||||
@pytest.fixture
|
||||
def hybrid_rag_pipeline(self, document_store):
|
||||
"""Create a hybrid RAG pipeline for testing."""
|
||||
prompt_template = "Documents: {% for doc in documents %}{{ doc.content }}{% endfor %} Question: {{question}}"
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("bm25_retriever", InMemoryBM25Retriever(document_store=document_store))
|
||||
pipeline.add_component("query_embedder", FakeEmbedder())
|
||||
pipeline.add_component("embedding_retriever", InMemoryEmbeddingRetriever(document_store=document_store))
|
||||
pipeline.add_component("doc_joiner", DocumentJoiner())
|
||||
pipeline.add_component("ranker", FakeRanker())
|
||||
pipeline.add_component("prompt_builder", PromptBuilder(template=prompt_template))
|
||||
pipeline.add_component("llm", FakeGenerator())
|
||||
pipeline.add_component("answer_builder", AnswerBuilder())
|
||||
|
||||
pipeline.connect("query_embedder.embedding", "embedding_retriever.query_embedding")
|
||||
pipeline.connect("embedding_retriever", "doc_joiner.documents")
|
||||
pipeline.connect("bm25_retriever", "doc_joiner.documents")
|
||||
pipeline.connect("doc_joiner", "ranker.documents")
|
||||
pipeline.connect("ranker", "prompt_builder.documents")
|
||||
pipeline.connect("prompt_builder", "llm")
|
||||
pipeline.connect("llm.replies", "answer_builder.replies")
|
||||
pipeline.connect("llm.meta", "answer_builder.meta")
|
||||
pipeline.connect("doc_joiner", "answer_builder.documents")
|
||||
|
||||
return pipeline
|
||||
|
||||
BREAKPOINT_COMPONENTS = [
|
||||
"bm25_retriever",
|
||||
"query_embedder",
|
||||
"embedding_retriever",
|
||||
"doc_joiner",
|
||||
"ranker",
|
||||
"prompt_builder",
|
||||
"llm",
|
||||
"answer_builder",
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_pipeline_breakpoints_hybrid_rag(
|
||||
self, hybrid_rag_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
"""
|
||||
Test that a hybrid RAG pipeline can be executed with breakpoints at each component.
|
||||
"""
|
||||
question = "Where does Mark live?"
|
||||
data = {
|
||||
"query_embedder": {"text": question},
|
||||
"bm25_retriever": {"query": question},
|
||||
"ranker": {"query": question, "top_k": 5},
|
||||
"prompt_builder": {"question": question},
|
||||
"answer_builder": {"query": question},
|
||||
}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = hybrid_rag_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=hybrid_rag_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
assert "answer_builder" in result
|
||||
assert result["answer_builder"]["answers"][0].data == "Mark lives in Berlin."
|
||||
assert len(result["answer_builder"]["answers"]) == 1
|
||||
assert result["answer_builder"]["answers"][0].meta["model"] == "fake"
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
|
||||
from haystack.components.converters import OutputAdapter
|
||||
from haystack.components.joiners import StringJoiner
|
||||
from haystack.core.errors import BreakpointException
|
||||
from haystack.core.pipeline.pipeline import Pipeline
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint
|
||||
|
||||
|
||||
class TestPipelineBreakpoints:
|
||||
@pytest.fixture
|
||||
def string_joiner_pipeline(self):
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component(
|
||||
"prompt_builder_1", ChatPromptBuilder(template=[ChatMessage.from_user("Builder 1: {{query}}")])
|
||||
)
|
||||
pipeline.add_component(
|
||||
"prompt_builder_2", ChatPromptBuilder(template=[ChatMessage.from_user("Builder 2: {{query}}")])
|
||||
)
|
||||
pipeline.add_component("adapter_1", OutputAdapter("{{messages[0].text}}", output_type=str))
|
||||
pipeline.add_component("adapter_2", OutputAdapter("{{messages[0].text}}", output_type=str))
|
||||
pipeline.add_component("string_joiner", StringJoiner())
|
||||
|
||||
pipeline.connect("prompt_builder_1.prompt", "adapter_1.messages")
|
||||
pipeline.connect("prompt_builder_2.prompt", "adapter_2.messages")
|
||||
pipeline.connect("adapter_1", "string_joiner.strings")
|
||||
pipeline.connect("adapter_2", "string_joiner.strings")
|
||||
|
||||
return pipeline
|
||||
|
||||
BREAKPOINT_COMPONENTS = ["prompt_builder_1", "prompt_builder_2", "adapter_1", "adapter_2", "string_joiner"]
|
||||
|
||||
@pytest.mark.parametrize("component", BREAKPOINT_COMPONENTS, ids=BREAKPOINT_COMPONENTS)
|
||||
@pytest.mark.integration
|
||||
def test_string_joiner_pipeline(
|
||||
self, string_joiner_pipeline, output_directory, component, load_and_resume_pipeline_snapshot
|
||||
):
|
||||
string_1 = "What's Natural Language Processing?"
|
||||
string_2 = "What is life?"
|
||||
data = {"prompt_builder_1": {"query": string_1}, "prompt_builder_2": {"query": string_2}}
|
||||
|
||||
# Create a Breakpoint on-the-fly using the shared output directory
|
||||
break_point = Breakpoint(component_name=component, visit_count=0, snapshot_file_path=str(output_directory))
|
||||
|
||||
try:
|
||||
_ = string_joiner_pipeline.run(data, break_point=break_point)
|
||||
except BreakpointException:
|
||||
pass
|
||||
|
||||
result = load_and_resume_pipeline_snapshot(
|
||||
pipeline=string_joiner_pipeline,
|
||||
output_directory=output_directory,
|
||||
component_name=break_point.component_name,
|
||||
data=data,
|
||||
)
|
||||
assert result["string_joiner"]
|
||||
assert "Builder 1: What's Natural Language Processing?" in result["string_joiner"]["strings"]
|
||||
assert "Builder 2: What is life?" in result["string_joiner"]["strings"]
|
||||
@@ -0,0 +1,135 @@
|
||||
# `Pipeline.run()` behavioural tests
|
||||
|
||||
This module contains all behavioural tests for `Pipeline.run()`.
|
||||
|
||||
`pipeline_run.feature` contains the definition of the tests using a subset of the [Gherkin language](https://cucumber.io/docs/gherkin/). It's not the full language because we're using `pytest-bdd` and it doesn't implement it in full, but it's good enough for our use case. For more info see the [project `README.md`](https://github.com/pytest-dev/pytest-bdd).
|
||||
|
||||
There are two cases covered by these tests:
|
||||
|
||||
1. `Pipeline.run()` returns some output
|
||||
2. `Pipeline.run()` raises an exception
|
||||
|
||||
### Correct Pipeline
|
||||
|
||||
In the first case to add a new test you need add a new entry in the `Examples` of the `Running a correct Pipeline` scenario outline and create the corresponding step that creates the `Pipeline` you need to test.
|
||||
|
||||
For example to add a test for a linear `Pipeline` I add a new `that is linear` kind in `pipeline_run.feature`.
|
||||
|
||||
```gherkin
|
||||
Scenario Outline: Running a correct Pipeline
|
||||
Given a pipeline <kind>
|
||||
When I run the Pipeline
|
||||
Then it should return the expected result
|
||||
|
||||
Examples:
|
||||
| kind |
|
||||
| that has no components |
|
||||
| that is linear |
|
||||
```
|
||||
|
||||
Then define a new `pipeline_that_is_linear` function in `test_run.py`.
|
||||
The function must be decorated with `@given` and return a tuple containing the `Pipeline` instance and a list of `PipelineRunData` instances.
|
||||
`PipelineRunData` is a dataclass that stores all the information necessary to verify the `Pipeline` ran as expected.
|
||||
The `@given` arguments must be the full step name, `"a pipeline that is linear"` in this case, and `target_fixture` must be set to `"pipeline_data"`.
|
||||
|
||||
```python
|
||||
@given("a pipeline that is linear", target_fixture="pipeline_data")
|
||||
def pipeline_that_is_linear():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||||
pipeline.add_component("second_addition", AddFixedValue())
|
||||
pipeline.add_component("double", Double())
|
||||
pipeline.connect("first_addition", "double")
|
||||
pipeline.connect("double", "second_addition")
|
||||
|
||||
return (
|
||||
pipeline,
|
||||
[
|
||||
PipelineRunData(
|
||||
inputs={"first_addition": {"value": 1}},
|
||||
expected_outputs={"second_addition": {"result": 7}},
|
||||
expected_run_order=["first_addition", "double", "second_addition"],
|
||||
)
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
Some kinds of `Pipeline`s require multiple runs to verify they work correctly, for example those with multiple branches.
|
||||
For this reason we can return a list of `PipelineRunData`, we'll run the `Pipeline` for each instance.
|
||||
For example, we could test two different runs of the same pipeline like this:
|
||||
|
||||
```python
|
||||
@given("a pipeline that is linear", target_fixture="pipeline_data")
|
||||
def pipeline_that_is_linear():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("first_addition", AddFixedValue(add=2))
|
||||
pipeline.add_component("second_addition", AddFixedValue())
|
||||
pipeline.add_component("double", Double())
|
||||
pipeline.connect("first_addition", "double")
|
||||
pipeline.connect("double", "second_addition")
|
||||
|
||||
return (
|
||||
pipeline,
|
||||
[
|
||||
PipelineRunData(
|
||||
inputs={"first_addition": {"value": 1}},
|
||||
include_outputs_from=set(),
|
||||
expected_outputs={"second_addition": {"result": 7}},
|
||||
expected_run_order=["first_addition", "double", "second_addition"],
|
||||
),
|
||||
PipelineRunData(
|
||||
inputs={"first_addition": {"value": 100}},
|
||||
include_outputs_from=set(),
|
||||
expected_outputs={"first_addition": {"value": 206}},
|
||||
expected_run_order=["first_addition", "double", "second_addition"],
|
||||
),
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
### Bad Pipeline
|
||||
|
||||
The second case is similar to the first one.
|
||||
In this case we test that a `Pipeline` with an infinite loop raises `PipelineMaxLoops`.
|
||||
|
||||
```gherkin
|
||||
Scenario Outline: Running a bad Pipeline
|
||||
Given a pipeline <kind>
|
||||
When I run the Pipeline
|
||||
Then it must have raised <exception>
|
||||
|
||||
Examples:
|
||||
| kind | exception |
|
||||
| that has an infinite loop | PipelineMaxLoops |
|
||||
```
|
||||
|
||||
In a similar way as first case we need to defined a new `pipeline_that_has_an_infinite_loop` function in `test_run.py`, with some small differences.
|
||||
The only difference from the first case is the last value returned by the function, we just omit the expected outputs and the expected run order.
|
||||
|
||||
```python
|
||||
@given("a pipeline that has an infinite loop", target_fixture="pipeline_data")
|
||||
def pipeline_that_has_an_infinite_loop():
|
||||
def custom_init(self):
|
||||
component.set_input_type(self, "x", int)
|
||||
component.set_input_type(self, "y", int, 1)
|
||||
component.set_output_types(self, a=int, b=int)
|
||||
|
||||
FakeComponent = component_class("FakeComponent", output={"a": 1, "b": 1}, extra_fields={"__init__": custom_init})
|
||||
pipe.add_component("first", FakeComponent())
|
||||
pipe.add_component("second", FakeComponent())
|
||||
pipe.connect("first.a", "second.x")
|
||||
pipe.connect("second.b", "first.y")
|
||||
return pipe, [PipelineRunData({"first": {"x": 1}})]
|
||||
```
|
||||
|
||||
## Why?
|
||||
|
||||
As the time of writing, tests that invoke `Pipeline.run()` are scattered between different files with very little clarity on what they are intended to test - the only indicators are the name of each test itself and the name of their parent module. This makes it difficult to understand which behaviours are being tested, if they are tested redundantly or if they work correctly.
|
||||
|
||||
The introduction of the Gherkin file allows for a single "source of truth" that enumerates (ideally, in an exhaustive manner) all the behaviours of the pipeline execution logic that we wish to test. This intermediate mapping of behaviours to actual test cases is meant to provide an overview of the latter and reduce the cognitive overhead of understanding them. When writing new tests, we now "tag" them with a specific behavioural parameter that's specified in a Gherkin scenario.
|
||||
|
||||
This tag and behavioural parameter mapping is meant to be 1 to 1, meaning each "Given" step must map to one and only one function. If multiple function are marked with `@given("step name")` the last declaration will override all the previous ones. So it's important to verify that there are no other existing steps with the same name when adding a new one.
|
||||
|
||||
While one could functionally do the same with well-defined test names and detailed comments on what is being tested, it would still lack the overview that the above approach provides. It's also extensible in that new scenarios with different behaviours can be introduced easily (e.g: for `async` pipeline execution logic).
|
||||
|
||||
Apart from the above, the newly introduced harness ensures that all behavioural pipeline tests return a structured result, which simplifies checking of side-effects.
|
||||
@@ -0,0 +1,188 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pandas import DataFrame
|
||||
from pytest_bdd import parsers, then, when
|
||||
|
||||
from haystack import Pipeline, component
|
||||
from test.tracing.utils import SpyingTracer
|
||||
|
||||
|
||||
@pytest.fixture(params=["sync", "async"])
|
||||
def pipeline_run_mode(request):
|
||||
"""
|
||||
Parametrizes each scenario so it runs once through `Pipeline.run` (sync) and once through
|
||||
`Pipeline.run_async` (async), exercising both execution engines.
|
||||
"""
|
||||
return request.param
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineRunData:
|
||||
"""
|
||||
Holds the inputs and expected outputs for a single Pipeline run.
|
||||
"""
|
||||
|
||||
inputs: dict[str, Any]
|
||||
include_outputs_from: set[str] = field(default_factory=set)
|
||||
expected_outputs: dict[str, Any] = field(default_factory=dict)
|
||||
expected_component_calls: dict[tuple[str, int], dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _PipelineResult:
|
||||
"""
|
||||
Holds the outputs and the run order of a single Pipeline run.
|
||||
"""
|
||||
|
||||
outputs: dict[str, Any]
|
||||
component_calls: dict[tuple[str, int], dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@when("I run the Pipeline", target_fixture="pipeline_result")
|
||||
def run_pipeline(
|
||||
pipeline_data: tuple[Pipeline, list[PipelineRunData]], spying_tracer: SpyingTracer, pipeline_run_mode: str
|
||||
) -> list[tuple[_PipelineResult, PipelineRunData]] | Exception:
|
||||
if pipeline_run_mode == "async":
|
||||
return run_async_pipeline(pipeline_data, spying_tracer)
|
||||
return run_sync_pipeline(pipeline_data, spying_tracer)
|
||||
|
||||
|
||||
def run_async_pipeline(
|
||||
pipeline_data: tuple[Pipeline, list[PipelineRunData]], spying_tracer: SpyingTracer
|
||||
) -> list[tuple[_PipelineResult, PipelineRunData]] | Exception:
|
||||
"""
|
||||
Attempts to run a pipeline with the given inputs.
|
||||
`pipeline_data` is a tuple that must contain:
|
||||
* A Pipeline instance
|
||||
* The data to run the pipeline with
|
||||
|
||||
If successful returns a tuple of the run outputs and the expected outputs.
|
||||
In case an exceptions is raised returns that.
|
||||
"""
|
||||
pipeline, pipeline_run_data = pipeline_data[0], pipeline_data[1]
|
||||
|
||||
results: list[_PipelineResult] = []
|
||||
|
||||
async def run_inner(data, include_outputs_from):
|
||||
"""Wrapper function to call pipeline.run_async method with required params."""
|
||||
return await pipeline.run_async(data=data.inputs, include_outputs_from=include_outputs_from)
|
||||
|
||||
for data in pipeline_run_data:
|
||||
try:
|
||||
outputs = asyncio.run(run_inner(data, data.include_outputs_from))
|
||||
|
||||
component_calls = {
|
||||
(span.tags["haystack.component.name"], span.tags["haystack.component.visits"]): span.tags[
|
||||
"haystack.component.input"
|
||||
]
|
||||
for span in spying_tracer.spans
|
||||
if "haystack.component.name" in span.tags and "haystack.component.visits" in span.tags
|
||||
}
|
||||
results.append(_PipelineResult(outputs=outputs, component_calls=component_calls))
|
||||
spying_tracer.spans.clear()
|
||||
except Exception as e:
|
||||
return e
|
||||
|
||||
return list(zip(results, pipeline_run_data, strict=True))
|
||||
|
||||
|
||||
def run_sync_pipeline(
|
||||
pipeline_data: tuple[Pipeline, list[PipelineRunData]], spying_tracer: SpyingTracer
|
||||
) -> list[tuple[_PipelineResult, PipelineRunData]] | Exception:
|
||||
"""
|
||||
Attempts to run a pipeline with the given inputs.
|
||||
`pipeline_data` is a tuple that must contain:
|
||||
* A Pipeline instance
|
||||
* The data to run the pipeline with
|
||||
|
||||
If successful returns a tuple of the run outputs and the expected outputs.
|
||||
In case an exceptions is raised returns that.
|
||||
"""
|
||||
pipeline, pipeline_run_data = pipeline_data[0], pipeline_data[1]
|
||||
|
||||
results: list[_PipelineResult] = []
|
||||
|
||||
for data in pipeline_run_data:
|
||||
try:
|
||||
outputs = pipeline.run(data=data.inputs, include_outputs_from=data.include_outputs_from)
|
||||
|
||||
component_calls = {
|
||||
(span.tags["haystack.component.name"], span.tags["haystack.component.visits"]): span.tags[
|
||||
"haystack.component.input"
|
||||
]
|
||||
for span in spying_tracer.spans
|
||||
if "haystack.component.name" in span.tags and "haystack.component.visits" in span.tags
|
||||
}
|
||||
results.append(_PipelineResult(outputs=outputs, component_calls=component_calls))
|
||||
spying_tracer.spans.clear()
|
||||
except Exception as e:
|
||||
return e
|
||||
return list(zip(results, pipeline_run_data, strict=True))
|
||||
|
||||
|
||||
@then("it should return the expected result")
|
||||
def check_pipeline_result(pipeline_result: list[tuple[_PipelineResult, PipelineRunData]]) -> None:
|
||||
for res, data in pipeline_result:
|
||||
compare_outputs_with_dataframes(res.outputs, data.expected_outputs)
|
||||
|
||||
|
||||
@then("components are called with the expected inputs")
|
||||
def check_component_calls(pipeline_result: list[tuple[_PipelineResult, PipelineRunData]]) -> None:
|
||||
for res, data in pipeline_result:
|
||||
assert compare_outputs_with_dataframes(res.component_calls, data.expected_component_calls)
|
||||
|
||||
|
||||
@then(parsers.parse("it must have raised {exception_class_name}"))
|
||||
def check_pipeline_raised(pipeline_result: Exception, exception_class_name: str) -> None:
|
||||
assert pipeline_result.__class__.__name__ == exception_class_name
|
||||
|
||||
|
||||
def compare_outputs_with_dataframes(actual: dict, expected: dict) -> bool:
|
||||
"""
|
||||
Compare two component_calls or pipeline outputs dictionaries where values may contain DataFrames.
|
||||
"""
|
||||
assert actual.keys() == expected.keys()
|
||||
|
||||
for key in actual:
|
||||
actual_data = actual[key]
|
||||
expected_data = expected[key]
|
||||
|
||||
assert actual_data.keys() == expected_data.keys()
|
||||
|
||||
for data_key in actual_data:
|
||||
actual_value = actual_data[data_key]
|
||||
expected_value = expected_data[data_key]
|
||||
|
||||
if isinstance(actual_value, DataFrame) and isinstance(expected_value, DataFrame):
|
||||
assert actual_value.equals(expected_value)
|
||||
else:
|
||||
# We do expected_value first so ANY can be used in expected outputs
|
||||
assert expected_value == actual_value
|
||||
|
||||
return True
|
||||
|
||||
|
||||
@component
|
||||
class FixedGenerator:
|
||||
def __init__(self, replies):
|
||||
self.replies = replies
|
||||
self.idx = 0
|
||||
|
||||
@component.output_types(replies=list[str])
|
||||
def run(self, prompt: str) -> dict[str, Any]:
|
||||
if self.idx < len(self.replies):
|
||||
replies = [self.replies[self.idx]]
|
||||
self.idx += 1
|
||||
else:
|
||||
self.idx = 0
|
||||
replies = [self.replies[self.idx]]
|
||||
self.idx += 1
|
||||
|
||||
return {"replies": replies}
|
||||
@@ -0,0 +1,76 @@
|
||||
Feature: Pipeline running
|
||||
|
||||
Scenario Outline: Running a correct Pipeline
|
||||
Given a pipeline <kind>
|
||||
When I run the Pipeline
|
||||
Then components are called with the expected inputs
|
||||
And it should return the expected result
|
||||
|
||||
Examples:
|
||||
| kind |
|
||||
| that has no components |
|
||||
| that is linear |
|
||||
| that is really complex with lots of components, forks, and loops |
|
||||
| that has a single component with a default input |
|
||||
| that has two loops of identical lengths |
|
||||
| that has two loops of different lengths |
|
||||
| that has a single loop with two conditional branches |
|
||||
| that has a component with dynamic inputs defined in init |
|
||||
| that has two branches that don't merge |
|
||||
| that has three branches that don't merge |
|
||||
| that has two branches that merge |
|
||||
| that has different combinations of branches that merge and do not merge |
|
||||
| that has two branches, one of which loops back |
|
||||
| that has a component with mutable input |
|
||||
| that has a component with mutable output sent to multiple inputs |
|
||||
| that has a greedy and variadic component after a component with default input |
|
||||
| that has a component with only default inputs |
|
||||
| that has a component with only default inputs as first to run and receives inputs from a loop |
|
||||
| that has multiple branches that merge into a component with a single variadic input |
|
||||
| that has multiple branches of different lengths that merge into a component with a single variadic input |
|
||||
| that is linear and returns intermediate outputs |
|
||||
| that has a loop and returns intermediate outputs from it |
|
||||
| that is linear and returns intermediate outputs from multiple sockets |
|
||||
| that has a component with default inputs that doesn't receive anything from its sender |
|
||||
| that has a component with default inputs that doesn't receive anything from its sender but receives input from user |
|
||||
| that has a loop and a component with default inputs that doesn't receive anything from its sender but receives input from user |
|
||||
| that has multiple components with only default inputs and are added in a different order from the order of execution |
|
||||
| that is linear with conditional branching and multiple joins |
|
||||
| that is a simple agent |
|
||||
| that has a variadic component that receives partial inputs |
|
||||
| that has a variadic component that receives partial inputs in a different order |
|
||||
| that has an answer joiner variadic component |
|
||||
| that is linear and a component in the middle receives optional input from other components and input from the user |
|
||||
| that has a loop in the middle |
|
||||
| that has variadic component that receives a conditional input |
|
||||
| that has a string variadic component |
|
||||
| that is an agent that can use RAG |
|
||||
| that has a feedback loop |
|
||||
| created in a non-standard order that has a loop |
|
||||
| that has an agent with a feedback cycle |
|
||||
| that passes outputs that are consumed in cycle to outside the cycle |
|
||||
| with a component that has dynamic default inputs |
|
||||
| with a component that has variadic dynamic default inputs |
|
||||
| that is a file conversion pipeline with two joiners |
|
||||
| that is a file conversion pipeline with three joiners |
|
||||
| that is a file conversion pipeline with three joiners and a loop |
|
||||
| that has components returning dataframes |
|
||||
| where a single component connects multiple sockets to the same receiver socket |
|
||||
| where a component in a cycle provides inputs for a component outside the cycle in one iteration and no input in another iteration |
|
||||
| that is blocked because not enough component inputs |
|
||||
| that is a file conversion pipeline with three auto joiners |
|
||||
| that has an auto joiner that takes in user inputs |
|
||||
| that performs automatic conversion between list of ChatMessage and str |
|
||||
| that performs automatic conversion wrapping ChatMessage for a Union receiver |
|
||||
|
||||
Scenario Outline: Running a bad Pipeline
|
||||
Given a pipeline <kind>
|
||||
When I run the Pipeline
|
||||
Then it must have raised <exception>
|
||||
|
||||
Examples:
|
||||
| kind | exception |
|
||||
| that has an infinite loop | PipelineMaxComponentRuns |
|
||||
| that has a component that doesn't return a dictionary | PipelineRuntimeError |
|
||||
| that has a cycle that would get it stuck | PipelineComponentsBlockedError |
|
||||
| that fails automatic conversion between list of ChatMessage and str | PipelineRuntimeError |
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,552 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Document, Pipeline, component
|
||||
from haystack.components.joiners import BranchJoiner
|
||||
from haystack.core.errors import PipelineRuntimeError
|
||||
|
||||
_test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar("_test_context_var", default="unset")
|
||||
|
||||
|
||||
def test_async_pipeline_reentrance(waiting_component, spying_tracer):
|
||||
pp = Pipeline()
|
||||
pp.add_component("wait", waiting_component())
|
||||
|
||||
run_data = [{"wait_for": 0.001}, {"wait_for": 0.002}]
|
||||
|
||||
async def run_all():
|
||||
# Create concurrent tasks for each pipeline run
|
||||
tasks = [pp.run_async(data) for data in run_data]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
asyncio.run(run_all())
|
||||
component_spans = [sp for sp in spying_tracer.spans if sp.operation_name == "haystack.component.run"]
|
||||
assert len(component_spans) == 2
|
||||
for span in component_spans:
|
||||
assert span.tags["haystack.component.visits"] == 1
|
||||
|
||||
|
||||
def test_run_in_sync_context(waiting_component):
|
||||
pp = Pipeline()
|
||||
pp.add_component("wait", waiting_component())
|
||||
|
||||
result = pp.run({"wait_for": 0.001})
|
||||
|
||||
assert result == {"wait": {"waited_for": 0.001}}
|
||||
|
||||
|
||||
def test_run_async_with_invalid_concurrency_limit():
|
||||
pp = Pipeline()
|
||||
with pytest.raises(ValueError, match="concurrency_limit must be greater than or equal to 1"):
|
||||
asyncio.run(pp.run_async({}, concurrency_limit=0))
|
||||
|
||||
|
||||
def test_component_with_empty_dict_as_output_appears_in_results():
|
||||
"""Test that components that return an empty dict as output appear in results as an empty dict"""
|
||||
|
||||
@component
|
||||
class Producer:
|
||||
def __init__(self, prefix: str):
|
||||
self.prefix = prefix
|
||||
|
||||
@component.output_types(value=str | None)
|
||||
def run(self, text: str | None) -> dict[str, str | None]:
|
||||
return {"value": f"{self.prefix}: {text}"}
|
||||
|
||||
@component.output_types(value=str | None)
|
||||
async def run_async(self, text: str | None) -> dict[str, str | None]:
|
||||
return {"value": f"{self.prefix}: {text}"}
|
||||
|
||||
@component
|
||||
class EmptyProcessor:
|
||||
@component.output_types()
|
||||
def run(self, sources: list[str]) -> dict:
|
||||
# Returns empty dict when sources is empty
|
||||
return {}
|
||||
|
||||
@component.output_types()
|
||||
async def run_async(self, sources: list[str]) -> dict:
|
||||
# Returns empty dict when sources is empty
|
||||
return {}
|
||||
|
||||
@component
|
||||
class Combiner:
|
||||
@component.output_types(combined=str)
|
||||
def run(self, input_a: str | None, input_b: str | None) -> dict[str, str]:
|
||||
if input_a is None:
|
||||
input_a = ""
|
||||
if input_b is None:
|
||||
input_b = ""
|
||||
return {"combined": f"{input_a} | {input_b}"}
|
||||
|
||||
@component.output_types(combined=str)
|
||||
async def run_async(self, input_a: str | None, input_b: str | None) -> dict[str, str]:
|
||||
if input_a is None:
|
||||
input_a = ""
|
||||
if input_b is None:
|
||||
input_b = ""
|
||||
return {"combined": f"{input_a} | {input_b}"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("producer_a", Producer("A"))
|
||||
pp.add_component("producer_b", Producer("B"))
|
||||
pp.add_component("empty_processor", EmptyProcessor())
|
||||
pp.add_component("combiner", Combiner())
|
||||
|
||||
pp.connect("producer_a.value", "combiner.input_a")
|
||||
pp.connect("producer_b.value", "combiner.input_b")
|
||||
|
||||
result = asyncio.run(
|
||||
pp.run_async(
|
||||
{"producer_a": {"text": "hello"}, "producer_b": {"text": "world"}, "empty_processor": {"sources": []}},
|
||||
include_outputs_from={"producer_a", "empty_processor", "combiner"},
|
||||
)
|
||||
)
|
||||
|
||||
# Producer A should appear in results because it's in include_outputs_from
|
||||
assert "producer_a" in result
|
||||
assert result["producer_a"] == {"value": "A: hello"}
|
||||
# Producer B should NOT appear since it's not in include_outputs_from
|
||||
assert "producer_b" not in result
|
||||
# Combiner should appear in results
|
||||
assert "combiner" in result
|
||||
assert result["combiner"] == {"combined": "A: hello | B: world"}
|
||||
# Empty processor should appear in results even though it returns an empty dict
|
||||
# because it's in include_outputs_from
|
||||
assert "empty_processor" in result
|
||||
assert result["empty_processor"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test__run_component_async_warns_on_extra_output_keys(caplog):
|
||||
"""Test that a warning is raised when a component returns undeclared output keys."""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class ExtraKeyComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"output": value, "extra_key": "unexpected"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("extra", ExtraKeyComponent())
|
||||
|
||||
await pp._run_component_async(
|
||||
component_name="extra",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("extra", 0),
|
||||
component_inputs={"value": "test"},
|
||||
component_visits={"extra": 0},
|
||||
)
|
||||
assert "returned output keys" in caplog.text
|
||||
assert "extra_key" in caplog.text
|
||||
assert "not declared" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test__run_component_async_no_warning_on_correct_output_keys(caplog):
|
||||
"""Test that no warning is raised when a component returns the correct output keys."""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class CorrectComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"output": value}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("correct", CorrectComponent())
|
||||
|
||||
await pp._run_component_async(
|
||||
component_name="correct",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("correct", 0),
|
||||
component_inputs={"value": "test"},
|
||||
component_visits={"correct": 0},
|
||||
)
|
||||
assert "returned output keys" not in caplog.text
|
||||
assert "did not produce output keys" not in caplog.text
|
||||
|
||||
|
||||
def test_async_pipeline_is_possibly_blocked_warning_message(caplog):
|
||||
"""
|
||||
Test that the pipeline raises a warning when it is possibly blocked due to missing inputs.
|
||||
|
||||
The situation below looks a little contrived, but it has happened in practice that users create pipelines
|
||||
and accidentally made a mistake in their component code.
|
||||
"""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class MisconfiguredComponent:
|
||||
# Here we purposely declare other_output which is not actually returned by the run() method
|
||||
@component.output_types(output=str, other_output=str)
|
||||
def run(self, required_input: str) -> dict[str, str]:
|
||||
return {"output": "test"}
|
||||
|
||||
@component
|
||||
class SimpleComponentTwoInputs:
|
||||
@component.output_types(output=str)
|
||||
def run(self, required_input: str, second_required_input: str) -> dict[str, str]:
|
||||
return {"output": "test"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", MisconfiguredComponent())
|
||||
pp.add_component("second", SimpleComponentTwoInputs())
|
||||
|
||||
# NOTE: We connect both outputs from the first component to the second component, but the first component
|
||||
# doesn't actually produce other_output, so the second component will be blocked due to missing input.
|
||||
pp.connect("first.output", "second.required_input")
|
||||
pp.connect("first.other_output", "second.second_required_input")
|
||||
|
||||
asyncio.run(pp.run_async({"first": {"required_input": "test"}}))
|
||||
assert "Cannot run pipeline - the pipeline appears to be blocked." in caplog.text
|
||||
assert " - 'second' (SimpleComponentTwoInputs)" in caplog.text
|
||||
|
||||
|
||||
def test_async_pipeline_ensure_inputs_are_deep_copied():
|
||||
"""
|
||||
Test to ensure that async pipeline deep copies the inputs before passing them to components.
|
||||
|
||||
This is important to prevent unintended side effects when components modify their inputs especially when
|
||||
the output from one component is passed to multiple other components.
|
||||
|
||||
Some other notes about how this situation can arise in practice:
|
||||
- When a component returns a mutable object (like a Document) and that output is passed to multiple other
|
||||
components.
|
||||
- This doesn't happen when using output types like strings or integers, because they are not shared by
|
||||
reference so we will only commonly see this for objects like our dataclasses.
|
||||
"""
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(output=Document)
|
||||
def run(self, document: Document) -> dict[str, Document]:
|
||||
# Creates a new document to avoid modifying in place
|
||||
new_document = Document(content=document.content)
|
||||
return {"output": new_document}
|
||||
|
||||
@component
|
||||
class ModifyingComponent:
|
||||
@component.output_types(output=Document)
|
||||
def run(self, document: Document) -> dict[str, Document]:
|
||||
return {"output": replace(document, content="modified")}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", SimpleComponent())
|
||||
pp.add_component("modifier", ModifyingComponent())
|
||||
# It's important that the following component has a name lower down the alphabetical order than "modifier",
|
||||
# since the pipeline runs components in a first-in-first-out manner based on ordered_component_names which is
|
||||
# sorted alphabetically.
|
||||
pp.add_component("second", SimpleComponent())
|
||||
|
||||
pp.connect("first.output", "modifier.document")
|
||||
pp.connect("first.output", "second.document")
|
||||
|
||||
result = asyncio.run(pp.run_async({"first": {"document": Document(content="original")}}))
|
||||
|
||||
assert result["modifier"]["output"].content == "modified"
|
||||
# Without deep copying the inputs, the second component would also see the modified document and produce
|
||||
# "modified" instead of "original"
|
||||
assert result["second"]["output"].content == "original"
|
||||
|
||||
|
||||
def test_async_pipeline_does_not_corrupt_outputs():
|
||||
"""
|
||||
Test that a component's output collected via include_outputs_from is not corrupted when a downstream
|
||||
component receives and mutates the same data in-place.
|
||||
"""
|
||||
|
||||
@component
|
||||
class Producer:
|
||||
@component.output_types(doc=Document)
|
||||
def run(self) -> dict:
|
||||
return {"doc": Document(content="original")}
|
||||
|
||||
@component
|
||||
class Mutator:
|
||||
@component.output_types(doc=Document)
|
||||
def run(self, doc: Document) -> dict:
|
||||
return {"doc": replace(doc, content="mutated")}
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("producer", Producer())
|
||||
pipe.add_component("mutator", Mutator())
|
||||
pipe.connect("producer.doc", "mutator.doc")
|
||||
|
||||
result = asyncio.run(pipe.run_async({}, include_outputs_from={"producer"}))
|
||||
|
||||
assert result["producer"]["doc"].content == "original"
|
||||
assert result["mutator"]["doc"].content == "mutated"
|
||||
|
||||
|
||||
@component
|
||||
class _Doubler:
|
||||
"""Minimal component used to exercise the isolation helper."""
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": value * 2}
|
||||
|
||||
|
||||
def _build_isolation_state(pipeline: Pipeline, data: dict) -> dict:
|
||||
"""
|
||||
Build the ephemeral run state that `_run_component_in_isolation` expects.
|
||||
|
||||
Mirrors the setup `run_async_generator` performs before the scheduling loop.
|
||||
"""
|
||||
inputs = pipeline._convert_to_internal_format(pipeline._prepare_component_input_data(data))
|
||||
names = sorted(pipeline.graph.nodes.keys())
|
||||
return {
|
||||
"inputs": inputs,
|
||||
"pipeline_outputs": {},
|
||||
"component_visits": dict.fromkeys(names, 0),
|
||||
"running_tasks": {},
|
||||
"scheduled_components": set(),
|
||||
"cached_receivers": {name: pipeline._find_receivers_from(name) for name in names},
|
||||
"include_outputs_from": set(),
|
||||
"parent_span": None,
|
||||
}
|
||||
|
||||
|
||||
class TestRunComponentInIsolation:
|
||||
@pytest.mark.asyncio
|
||||
async def test_runs_component_and_yields_output(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("doubler", _Doubler())
|
||||
state = _build_isolation_state(pp, {"doubler": {"value": 3}})
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="doubler", **state)]
|
||||
|
||||
assert results == [{"doubler": {"value": 6}}]
|
||||
assert state["pipeline_outputs"] == {"doubler": {"value": 6}}
|
||||
assert state["component_visits"]["doubler"] == 1
|
||||
# The component is added to and removed from scheduled_components over the course of the run.
|
||||
assert state["scheduled_components"] == set()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runs_greedy_component_consuming_single_input(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("joiner", BranchJoiner(type_=int))
|
||||
state = _build_isolation_state(pp, {})
|
||||
# Two values are queued on the greedy variadic socket; greedy consumption keeps only the first.
|
||||
state["inputs"]["joiner"] = {"value": [{"sender": None, "value": 1}, {"sender": None, "value": 2}]}
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="joiner", **state)]
|
||||
|
||||
assert results == [{"joiner": {"value": 1}}]
|
||||
assert state["component_visits"]["joiner"] == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drains_in_flight_tasks_before_running(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("doubler", _Doubler())
|
||||
state = _build_isolation_state(pp, {"doubler": {"value": 3}})
|
||||
|
||||
async def _in_flight() -> dict:
|
||||
return {"value": 99}
|
||||
|
||||
task = asyncio.create_task(_in_flight())
|
||||
state["running_tasks"][task] = "other"
|
||||
state["scheduled_components"].add("other")
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="doubler", **state)]
|
||||
|
||||
# The in-flight task is drained (and its output yielded) before the isolated component runs.
|
||||
assert {"other": {"value": 99}} in results
|
||||
assert {"doubler": {"value": 6}} in results
|
||||
assert results.index({"other": {"value": 99}}) < results.index({"doubler": {"value": 6}})
|
||||
assert state["running_tasks"] == {}
|
||||
assert "other" not in state["scheduled_components"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_when_component_already_scheduled(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("doubler", _Doubler())
|
||||
state = _build_isolation_state(pp, {"doubler": {"value": 3}})
|
||||
state["scheduled_components"].add("doubler")
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="doubler", **state)]
|
||||
|
||||
# Already scheduled: the component is not run.
|
||||
assert results == []
|
||||
assert state["component_visits"]["doubler"] == 0
|
||||
assert state["pipeline_outputs"] == {}
|
||||
assert "doubler" in state["scheduled_components"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_distributes_outputs_downstream_and_prunes_consumed(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", _Doubler())
|
||||
pp.add_component("second", _Doubler())
|
||||
pp.connect("first.value", "second.value")
|
||||
state = _build_isolation_state(pp, {"first": {"value": 3}})
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="first", **state)]
|
||||
|
||||
# `first`'s output is consumed by `second`, so it is pruned: nothing is yielded or stored as a pipeline output.
|
||||
assert results == []
|
||||
assert state["pipeline_outputs"] == {}
|
||||
# `second` can now consume the distributed value.
|
||||
second = pp._get_component_with_graph_metadata_and_visits("second", 0)
|
||||
assert pp._consume_component_inputs("second", second, state["inputs"]) == {"value": 6}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_include_outputs_from_yields_even_when_consumed(self):
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", _Doubler())
|
||||
pp.add_component("second", _Doubler())
|
||||
pp.connect("first.value", "second.value")
|
||||
state = _build_isolation_state(pp, {"first": {"value": 3}})
|
||||
state["include_outputs_from"] = {"first"}
|
||||
|
||||
results = [out async for out in pp._run_component_in_isolation(component_name="first", **state)]
|
||||
|
||||
# Even though `first`'s output is consumed by `second`, include_outputs_from forces it to be surfaced.
|
||||
assert results == [{"first": {"value": 6}}]
|
||||
assert state["pipeline_outputs"] == {"first": {"value": 6}}
|
||||
|
||||
|
||||
class TestInFlightTaskCleanupOnError:
|
||||
@pytest.mark.asyncio
|
||||
async def test_sibling_tasks_cancelled_when_a_component_errors(self):
|
||||
"""When a component fails, the other in-flight tasks must be cancelled and not leak."""
|
||||
slow_started = asyncio.Event()
|
||||
slow_cancelled = False
|
||||
|
||||
@component
|
||||
class Slow:
|
||||
@component.output_types(value=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"value": text}
|
||||
|
||||
@component.output_types(value=str)
|
||||
async def run_async(self, text: str) -> dict[str, str]:
|
||||
nonlocal slow_cancelled
|
||||
slow_started.set()
|
||||
try:
|
||||
await asyncio.sleep(5)
|
||||
except asyncio.CancelledError:
|
||||
slow_cancelled = True
|
||||
raise
|
||||
return {"value": text}
|
||||
|
||||
@component
|
||||
class Failing:
|
||||
@component.output_types(value=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
@component.output_types(value=str)
|
||||
async def run_async(self, text: str) -> dict[str, str]:
|
||||
# Fail only once the sibling is actually running, so there is an in-flight task to clean up.
|
||||
await slow_started.wait()
|
||||
raise RuntimeError("boom")
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("slow", Slow())
|
||||
pp.add_component("failing", Failing())
|
||||
|
||||
with pytest.raises(PipelineRuntimeError):
|
||||
await pp.run_async({"slow": {"text": "x"}, "failing": {"text": "y"}}, concurrency_limit=2)
|
||||
|
||||
assert slow_cancelled is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_in_flight_tasks_cancelled_when_generator_iteration_is_abandoned(self):
|
||||
"""When the consumer stops iterating run_async_generator early, in-flight tasks must be cancelled."""
|
||||
slow_started = asyncio.Event()
|
||||
slow_cancelled = False
|
||||
|
||||
@component
|
||||
class Fast:
|
||||
@component.output_types(value=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"value": text}
|
||||
|
||||
@component.output_types(value=str)
|
||||
async def run_async(self, text: str) -> dict[str, str]:
|
||||
# Yield an output only once the sibling is actually running, so it is in flight when we abandon.
|
||||
await slow_started.wait()
|
||||
return {"value": text}
|
||||
|
||||
@component
|
||||
class Slow:
|
||||
@component.output_types(value=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"value": text}
|
||||
|
||||
@component.output_types(value=str)
|
||||
async def run_async(self, text: str) -> dict[str, str]:
|
||||
nonlocal slow_cancelled
|
||||
slow_started.set()
|
||||
try:
|
||||
await asyncio.sleep(5)
|
||||
except asyncio.CancelledError:
|
||||
slow_cancelled = True
|
||||
raise
|
||||
return {"value": text}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("fast", Fast())
|
||||
pp.add_component("slow", Slow())
|
||||
|
||||
generator = pp.run_async_generator({"fast": {"text": "x"}, "slow": {"text": "y"}}, concurrency_limit=2)
|
||||
async for _partial in generator:
|
||||
break # abandon iteration after the first partial output
|
||||
await generator.aclose()
|
||||
|
||||
assert slow_cancelled is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_component_run_in_thread_receives_contextvars():
|
||||
"""
|
||||
Regression test: contextvars set in the calling async context (e.g. the active tracing span) must propagate
|
||||
to sync-only components, which the async run path dispatches to a thread. `asyncio.to_thread` guarantees this by
|
||||
copying the current context; a plain `loop.run_in_executor` would not.
|
||||
"""
|
||||
|
||||
@component
|
||||
class SyncContextVarReader:
|
||||
@component.output_types(value=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
# Read inside the executor thread — only visible if the calling context was copied
|
||||
return {"value": _test_context_var.get()}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("reader", SyncContextVarReader())
|
||||
|
||||
_test_context_var.set("propagated")
|
||||
result = await pp.run_async({"reader": {"text": "irrelevant"}})
|
||||
|
||||
assert result["reader"]["value"] == "propagated"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_raises_when_multi_element_list_is_unwrapped_at_runtime():
|
||||
@component
|
||||
class MultiStrProducer:
|
||||
@component.output_types(texts=list[str])
|
||||
def run(self) -> dict[str, list[str]]:
|
||||
return {"texts": ["first", "second", "third"]}
|
||||
|
||||
@component
|
||||
class SingleStrConsumer:
|
||||
@component.output_types(out=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"out": text}
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("producer", MultiStrProducer())
|
||||
pipe.add_component("consumer", SingleStrConsumer())
|
||||
pipe.connect("producer.texts", "consumer.text")
|
||||
|
||||
with pytest.raises(PipelineRuntimeError, match="Cannot unwrap a list of 3 items"):
|
||||
await pipe.run_async({})
|
||||
@@ -0,0 +1,371 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Pipeline, component
|
||||
from haystack.core.errors import PipelineRuntimeError
|
||||
from haystack.dataclasses import AsyncStreamingCallbackT, StreamingChunk
|
||||
|
||||
|
||||
@component
|
||||
class StreamingEcho:
|
||||
"""Streaming component used by tests: emits `n_chunks` chunks then returns a final reply."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
prefix: str = "tok",
|
||||
n_chunks: int = 3,
|
||||
fail: bool = False,
|
||||
chunk_delay: float = 0.0,
|
||||
streaming_callback: AsyncStreamingCallbackT | None = None,
|
||||
) -> None:
|
||||
self.prefix = prefix
|
||||
self.n_chunks = n_chunks
|
||||
self.fail = fail
|
||||
self.chunk_delay = chunk_delay
|
||||
self.streaming_callback = streaming_callback
|
||||
|
||||
@component.output_types(reply=str)
|
||||
def run(self, prompt: str, streaming_callback: AsyncStreamingCallbackT | None = None) -> dict:
|
||||
return {"reply": f"{self.prefix}-final"}
|
||||
|
||||
@component.output_types(reply=str)
|
||||
async def run_async(self, prompt: str, streaming_callback: AsyncStreamingCallbackT | None = None) -> dict:
|
||||
for i in range(self.n_chunks):
|
||||
if self.chunk_delay:
|
||||
await asyncio.sleep(self.chunk_delay)
|
||||
chunk = StreamingChunk(content=f"{self.prefix}{i}")
|
||||
if streaming_callback is not None:
|
||||
await streaming_callback(chunk)
|
||||
if self.fail:
|
||||
raise RuntimeError("boom")
|
||||
return {"reply": f"{self.prefix}-final"}
|
||||
|
||||
|
||||
@component
|
||||
class DynamicStreamingEcho:
|
||||
"""Streaming component that declares `streaming_callback` via `set_input_type` instead of the run signature."""
|
||||
|
||||
def __init__(self, prefix: str = "tok", n_chunks: int = 2) -> None:
|
||||
self.prefix = prefix
|
||||
self.n_chunks = n_chunks
|
||||
self.streaming_callback = None
|
||||
component.set_input_type(self, "prompt", str)
|
||||
component.set_input_type(self, "streaming_callback", AsyncStreamingCallbackT | None, None)
|
||||
|
||||
@component.output_types(reply=str)
|
||||
def run(self, **kwargs: Any) -> dict:
|
||||
return {"reply": f"{self.prefix}-final"}
|
||||
|
||||
@component.output_types(reply=str)
|
||||
async def run_async(self, **kwargs: Any) -> dict:
|
||||
cb = kwargs.get("streaming_callback")
|
||||
for i in range(self.n_chunks):
|
||||
chunk = StreamingChunk(content=f"{self.prefix}{i}")
|
||||
if cb is not None:
|
||||
await cb(chunk)
|
||||
return {"reply": f"{self.prefix}-final"}
|
||||
|
||||
|
||||
@component
|
||||
class Passthrough:
|
||||
"""Non-streaming async component used by filter-validation tests."""
|
||||
|
||||
@component.output_types(prompt=str)
|
||||
def run(self, prompt: str) -> dict:
|
||||
return {"prompt": prompt}
|
||||
|
||||
@component.output_types(prompt=str)
|
||||
async def run_async(self, prompt: str) -> dict:
|
||||
return {"prompt": prompt}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_yields_chunks_and_returns_result():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=3))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi"}})
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert [c.content for c in chunks] == ["s0", "s1", "s2"]
|
||||
assert handle.result["streamer"] == {"reply": "s-final"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_yields_chunks_with_flat_input():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=3))
|
||||
|
||||
# flat input form (`{"prompt": ...}` instead of `{"streamer": {"prompt": ...}}`)
|
||||
handle = pipeline.stream(data={"prompt": "hi"})
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert [c.content for c in chunks] == ["s0", "s1", "s2"]
|
||||
assert handle.result["streamer"] == {"reply": "s-final"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_composes_with_init_streaming_callback():
|
||||
seen = []
|
||||
|
||||
async def init_callback(chunk: StreamingChunk) -> None:
|
||||
seen.append(chunk.content)
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=2, streaming_callback=init_callback))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi"}})
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert [c.content for c in chunks] == ["s0", "s1"]
|
||||
assert seen == ["s0", "s1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_runtime_callback_overrides_init():
|
||||
init_seen = []
|
||||
runtime_seen = []
|
||||
|
||||
async def init_callback(chunk: StreamingChunk) -> None:
|
||||
init_seen.append(chunk.content)
|
||||
|
||||
async def runtime_callback(chunk: StreamingChunk) -> None:
|
||||
runtime_seen.append(chunk.content)
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=2, streaming_callback=init_callback))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi", "streaming_callback": runtime_callback}})
|
||||
[c async for c in handle]
|
||||
|
||||
assert init_seen == []
|
||||
assert runtime_seen == ["s0", "s1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_warns_on_sync_runtime_callback(caplog):
|
||||
seen: list[str] = []
|
||||
|
||||
def sync_callback(chunk: StreamingChunk) -> None:
|
||||
seen.append(chunk.content)
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=2))
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi", "streaming_callback": sync_callback}})
|
||||
[c async for c in handle]
|
||||
|
||||
assert "sync streaming callback" in caplog.text
|
||||
assert seen == ["s0", "s1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_detects_streaming_callback_declared_via_set_input_type():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", DynamicStreamingEcho(prefix="d", n_chunks=3))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi"}})
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert [c.content for c in chunks] == ["d0", "d1", "d2"]
|
||||
assert handle.result["streamer"] == {"reply": "d-final"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_streams_all_components_by_default():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("a", StreamingEcho(prefix="a", n_chunks=2))
|
||||
pipeline.add_component("b", StreamingEcho(prefix="b", n_chunks=2))
|
||||
|
||||
handle = pipeline.stream(data={"a": {"prompt": "x"}, "b": {"prompt": "y"}})
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert {c.content for c in chunks} == {"a0", "a1", "b0", "b1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_filters_to_selected_components():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("a", StreamingEcho(prefix="a", n_chunks=2))
|
||||
pipeline.add_component("b", StreamingEcho(prefix="b", n_chunks=2))
|
||||
|
||||
handle = pipeline.stream(data={"a": {"prompt": "x"}, "b": {"prompt": "y"}}, streaming_components=["a"])
|
||||
chunks = [c async for c in handle]
|
||||
|
||||
assert {c.content for c in chunks} == {"a0", "a1"}
|
||||
assert handle.result["a"] == {"reply": "a-final"}
|
||||
assert handle.result["b"] == {"reply": "b-final"}
|
||||
|
||||
|
||||
def test_stream_raises_for_unknown_component_in_filter():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho())
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown components") as excinfo:
|
||||
pipeline.stream(data={"streamer": {"prompt": "x"}}, streaming_components=["typo"])
|
||||
assert "typo" in str(excinfo.value) # the message names the offending component
|
||||
|
||||
|
||||
def test_stream_raises_for_non_streaming_component_in_filter():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho())
|
||||
pipeline.add_component("passthrough", Passthrough())
|
||||
|
||||
with pytest.raises(ValueError, match="do not support streaming") as excinfo:
|
||||
pipeline.stream(
|
||||
data={"streamer": {"prompt": "x"}, "passthrough": {"prompt": "x"}}, streaming_components=["passthrough"]
|
||||
)
|
||||
assert "passthrough" in str(excinfo.value) # the message names the offending component
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_propagates_pipeline_exception_during_iteration():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=1, fail=True))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}})
|
||||
|
||||
with pytest.raises(PipelineRuntimeError) as excinfo:
|
||||
async for _ in handle:
|
||||
pass
|
||||
|
||||
# the error message indicates the failing component and the underlying error
|
||||
message = str(excinfo.value)
|
||||
assert "failed to run" in message
|
||||
assert "Component name: 'streamer'" in message
|
||||
assert "Component type: 'StreamingEcho'" in message
|
||||
assert "boom" in message
|
||||
|
||||
# the original exception is preserved as the direct cause
|
||||
cause = excinfo.value.__cause__
|
||||
assert isinstance(cause, RuntimeError)
|
||||
assert str(cause) == "boom"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failing_callback_does_not_drop_chunk():
|
||||
async def failing_callback(chunk: StreamingChunk) -> None:
|
||||
raise RuntimeError("callback boom")
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=3, streaming_callback=failing_callback))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "hi"}})
|
||||
|
||||
received = []
|
||||
with pytest.raises(PipelineRuntimeError) as excinfo:
|
||||
async for chunk in handle:
|
||||
received.append(chunk.content)
|
||||
|
||||
assert received == ["s0"] # queued before the callback raised
|
||||
|
||||
# the error message indicates the failing component and the underlying error
|
||||
message = str(excinfo.value)
|
||||
assert "failed to run" in message
|
||||
assert "Component name: 'streamer'" in message
|
||||
assert "callback boom" in message
|
||||
|
||||
# the original exception is preserved as the direct cause
|
||||
cause = excinfo.value.__cause__
|
||||
assert isinstance(cause, RuntimeError)
|
||||
assert str(cause) == "callback boom"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_result_reraises_original_failure():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=1, fail=True))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}})
|
||||
|
||||
with pytest.raises(PipelineRuntimeError) as iter_excinfo:
|
||||
async for _ in handle:
|
||||
pass
|
||||
|
||||
# `result` re-raises the exception raised during iteration
|
||||
with pytest.raises(PipelineRuntimeError) as result_excinfo:
|
||||
handle.result
|
||||
assert result_excinfo.value is iter_excinfo.value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_result_raises_when_pipeline_not_finished():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=5))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}})
|
||||
|
||||
# no await between `stream()` and reading `result`, so the task has not started yet
|
||||
with pytest.raises(RuntimeError, match="Pipeline has not finished"):
|
||||
handle.result
|
||||
|
||||
await handle.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_cancels_pipeline_and_result_reports_cancelled():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=100))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}})
|
||||
|
||||
await handle.aclose()
|
||||
|
||||
with pytest.raises(RuntimeError, match="Pipeline was cancelled"):
|
||||
handle.result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_cancellation_cancels_pipeline():
|
||||
pipeline = Pipeline()
|
||||
# chunk_delay keeps the producer running so the consumer can be cancelled mid-stream
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=100, chunk_delay=0.01))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}})
|
||||
|
||||
async def consumer() -> None:
|
||||
async for _ in handle:
|
||||
pass
|
||||
|
||||
task = asyncio.create_task(consumer())
|
||||
await asyncio.sleep(0.03) # let the producer emit a couple of chunks
|
||||
# send a cancel request and wait for the consumer to finish
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert handle._task.done()
|
||||
with pytest.raises(RuntimeError, match="cancelled"):
|
||||
handle.result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_on_abandon_false_lets_pipeline_finish():
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("streamer", StreamingEcho(prefix="s", n_chunks=3, chunk_delay=0.01))
|
||||
|
||||
handle = pipeline.stream(data={"streamer": {"prompt": "x"}}, cancel_on_abandon=False)
|
||||
|
||||
async def consumer() -> None:
|
||||
async for _ in handle:
|
||||
pass
|
||||
|
||||
task = asyncio.create_task(consumer())
|
||||
await asyncio.sleep(0.015) # consume ~1 chunk
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
# the consumer was cancelled, but cancel_on_abandon=False leaves the pipeline running
|
||||
await handle._task
|
||||
assert not handle._task.cancelled()
|
||||
assert handle.result["streamer"] == {"reply": "s-final"}
|
||||
@@ -0,0 +1,664 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import component
|
||||
from haystack.core.errors import BreakpointException, PipelineInvalidPipelineSnapshotError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.pipeline.breakpoint import (
|
||||
HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED,
|
||||
_create_pipeline_snapshot,
|
||||
_is_snapshot_save_enabled,
|
||||
_save_pipeline_snapshot,
|
||||
_transform_json_structure,
|
||||
load_pipeline_snapshot,
|
||||
)
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.dataclasses.breakpoints import Breakpoint, PipelineSnapshot, PipelineState
|
||||
from haystack.utils import _deserialize_value_with_schema
|
||||
|
||||
_EMPTY_OBJECT_PAYLOAD = {"serialization_schema": {"type": "object", "properties": {}}, "serialized_data": {}}
|
||||
|
||||
|
||||
def test_transform_json_structure_unwraps_sender_value():
|
||||
data = {
|
||||
"key1": [{"sender": None, "value": "some value"}],
|
||||
"key2": [{"sender": "comp1", "value": 42}],
|
||||
"key3": "direct value",
|
||||
}
|
||||
|
||||
result = _transform_json_structure(data)
|
||||
|
||||
assert result == {"key1": "some value", "key2": 42, "key3": "direct value"}
|
||||
|
||||
|
||||
def test_transform_json_structure_handles_nested_structures():
|
||||
data = {
|
||||
"key1": [{"sender": None, "value": "value1"}],
|
||||
"key2": {"nested": [{"sender": "comp1", "value": "value2"}], "direct": "value3"},
|
||||
"key3": [[{"sender": None, "value": "value4"}], [{"sender": "comp2", "value": "value5"}]],
|
||||
}
|
||||
|
||||
result = _transform_json_structure(data)
|
||||
|
||||
assert result == {"key1": "value1", "key2": {"nested": "value2", "direct": "value3"}, "key3": ["value4", "value5"]}
|
||||
|
||||
|
||||
def test_load_pipeline_snapshot_loads_valid_snapshot(tmp_path):
|
||||
pipeline_snapshot = {
|
||||
"break_point": {"component_name": "comp1", "visit_count": 0},
|
||||
"pipeline_state": {"inputs": {}, "component_visits": {"comp1": 0, "comp2": 0}, "pipeline_outputs": {}},
|
||||
"original_input_data": {},
|
||||
"ordered_component_names": ["comp1", "comp2"],
|
||||
"include_outputs_from": ["comp1", "comp2"],
|
||||
}
|
||||
pipeline_snapshot_file = tmp_path / "state.json"
|
||||
with open(pipeline_snapshot_file, "w") as f:
|
||||
json.dump(pipeline_snapshot, f)
|
||||
|
||||
loaded_snapshot = load_pipeline_snapshot(pipeline_snapshot_file)
|
||||
assert loaded_snapshot == PipelineSnapshot.from_dict(pipeline_snapshot)
|
||||
|
||||
|
||||
def test_load_state_handles_invalid_state(tmp_path):
|
||||
pipeline_snapshot = {
|
||||
"break_point": {"component_name": "comp1", "visit_count": 0},
|
||||
"pipeline_state": {"inputs": {}, "component_visits": {"comp1": 0, "comp2": 0}, "pipeline_outputs": {}},
|
||||
"original_input_data": {},
|
||||
"include_outputs_from": ["comp1", "comp2"],
|
||||
"ordered_component_names": ["comp1", "comp3"], # inconsistent with component_visits
|
||||
}
|
||||
|
||||
pipeline_snapshot_file = tmp_path / "invalid_pipeline_snapshot.json"
|
||||
with open(pipeline_snapshot_file, "w") as f:
|
||||
json.dump(pipeline_snapshot, f)
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid pipeline snapshot from"):
|
||||
load_pipeline_snapshot(pipeline_snapshot_file)
|
||||
|
||||
|
||||
def test_breakpoint_saves_intermediate_outputs(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, input_value: str) -> dict[str, str]:
|
||||
return {"result": f"processed_{input_value}"}
|
||||
|
||||
pipeline = Pipeline()
|
||||
comp1 = SimpleComponent()
|
||||
comp2 = SimpleComponent()
|
||||
pipeline.add_component("comp1", comp1)
|
||||
pipeline.add_component("comp2", comp2)
|
||||
pipeline.connect("comp1", "comp2")
|
||||
|
||||
# breakpoint on comp2
|
||||
break_point = Breakpoint(component_name="comp2", visit_count=0, snapshot_file_path=str(tmp_path))
|
||||
|
||||
try:
|
||||
# run with include_outputs_from to capture intermediate outputs
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, include_outputs_from={"comp1"}, break_point=break_point)
|
||||
except BreakpointException as e:
|
||||
# breakpoint should be triggered
|
||||
assert e.component == "comp2"
|
||||
|
||||
# verify snapshot file contains the intermediate outputs
|
||||
snapshot_files = list(tmp_path.glob("comp2_*.json"))
|
||||
assert len(snapshot_files) == 1, f"Expected exactly one snapshot file, found {len(snapshot_files)}"
|
||||
|
||||
snapshot_file = snapshot_files[0]
|
||||
loaded_snapshot = load_pipeline_snapshot(snapshot_file)
|
||||
|
||||
# verify the snapshot contains the intermediate outputs from comp1
|
||||
assert loaded_snapshot.pipeline_state.pipeline_outputs == (
|
||||
{
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {"comp1": {"type": "object", "properties": {"result": {"type": "string"}}}},
|
||||
},
|
||||
"serialized_data": {"comp1": {"result": "processed_test"}},
|
||||
}
|
||||
)
|
||||
|
||||
# verify the whole pipeline state contains the expected data
|
||||
assert loaded_snapshot.pipeline_state.component_visits["comp1"] == 1
|
||||
assert loaded_snapshot.pipeline_state.component_visits["comp2"] == 0
|
||||
assert "comp1" in loaded_snapshot.include_outputs_from
|
||||
assert isinstance(loaded_snapshot.break_point, Breakpoint)
|
||||
assert loaded_snapshot.break_point.component_name == "comp2"
|
||||
assert loaded_snapshot.break_point.visit_count == 0
|
||||
|
||||
|
||||
@component
|
||||
class _AppendingComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, input_value: str) -> dict[str, str]:
|
||||
return {"result": f"{input_value}_processed"}
|
||||
|
||||
|
||||
def _three_component_pipeline() -> Pipeline:
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("comp1", _AppendingComponent())
|
||||
pipeline.add_component("comp2", _AppendingComponent())
|
||||
pipeline.add_component("comp3", _AppendingComponent())
|
||||
pipeline.connect("comp1", "comp2")
|
||||
pipeline.connect("comp2", "comp3")
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_break_point_with_pipeline_snapshot_steps_through_pipeline():
|
||||
pipeline = _three_component_pipeline()
|
||||
|
||||
# run until the breakpoint on comp2
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=Breakpoint(component_name="comp2"))
|
||||
first_snapshot = exc_info.value.pipeline_snapshot
|
||||
assert first_snapshot is not None
|
||||
assert first_snapshot.pipeline_state.component_visits == {"comp1": 1, "comp2": 0, "comp3": 0}
|
||||
|
||||
# step: resume from the snapshot and pause again at comp3
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(data={}, pipeline_snapshot=first_snapshot, break_point=Breakpoint(component_name="comp3"))
|
||||
second_snapshot = exc_info.value.pipeline_snapshot
|
||||
assert second_snapshot is not None
|
||||
assert second_snapshot.pipeline_state.component_visits == {"comp1": 1, "comp2": 1, "comp3": 0}
|
||||
|
||||
# resume from the second snapshot and run to completion
|
||||
result = pipeline.run(data={}, pipeline_snapshot=second_snapshot)
|
||||
assert result["comp3"]["result"] == "test_processed_processed_processed"
|
||||
|
||||
|
||||
def test_break_point_on_earlier_component_than_pipeline_snapshot_never_triggers():
|
||||
pipeline = _three_component_pipeline()
|
||||
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=Breakpoint(component_name="comp2"))
|
||||
snapshot = exc_info.value.pipeline_snapshot
|
||||
|
||||
# comp1 already ran before the snapshot was taken, so a breakpoint on it never triggers
|
||||
# and the resumed run completes normally
|
||||
result = pipeline.run(data={}, pipeline_snapshot=snapshot, break_point=Breakpoint(component_name="comp1"))
|
||||
assert result["comp3"]["result"] == "test_processed_processed_processed"
|
||||
|
||||
|
||||
def test_break_point_matching_pipeline_snapshot_break_point_raises():
|
||||
pipeline = _three_component_pipeline()
|
||||
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=Breakpoint(component_name="comp2"))
|
||||
snapshot = exc_info.value.pipeline_snapshot
|
||||
|
||||
with pytest.raises(PipelineInvalidPipelineSnapshotError, match="different component or visit count"):
|
||||
pipeline.run(data={}, pipeline_snapshot=snapshot, break_point=Breakpoint(component_name="comp2", visit_count=0))
|
||||
|
||||
|
||||
class TestCreatePipelineSnapshot:
|
||||
def test_create_pipeline_snapshot_all_fields(self):
|
||||
break_point = Breakpoint(component_name="comp2")
|
||||
ordered_component_names = ["comp1", "comp2"]
|
||||
include_outputs_from = {"comp1"}
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={"comp1": {"input_value": [{"sender": None, "value": "test"}]}, "comp2": {}},
|
||||
component_inputs={"input_value": "processed_test"},
|
||||
break_point=break_point,
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={"comp1": {"input_value": "test"}},
|
||||
ordered_component_names=ordered_component_names,
|
||||
include_outputs_from=include_outputs_from,
|
||||
pipeline_outputs={"comp1": {"result": "processed_test"}},
|
||||
)
|
||||
|
||||
assert snapshot.original_input_data == {
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {"comp1": {"type": "object", "properties": {"input_value": {"type": "string"}}}},
|
||||
},
|
||||
"serialized_data": {"comp1": {"input_value": "test"}},
|
||||
}
|
||||
assert snapshot.ordered_component_names == ordered_component_names
|
||||
assert snapshot.break_point == break_point
|
||||
assert snapshot.include_outputs_from == include_outputs_from
|
||||
assert snapshot.pipeline_state == PipelineState(
|
||||
inputs={
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"comp1": {"type": "object", "properties": {"input_value": {"type": "string"}}},
|
||||
"comp2": {"type": "object", "properties": {"input_value": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
"serialized_data": {"comp1": {"input_value": "test"}, "comp2": {"input_value": "processed_test"}},
|
||||
},
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
pipeline_outputs={
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {"comp1": {"type": "object", "properties": {"result": {"type": "string"}}}},
|
||||
},
|
||||
"serialized_data": {"comp1": {"result": "processed_test"}},
|
||||
},
|
||||
)
|
||||
|
||||
def test_create_pipeline_snapshot_with_dataclasses_in_pipeline_outputs(self):
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2"),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from={"comp1"},
|
||||
pipeline_outputs={"comp1": {"result": ChatMessage.from_user("hello")}},
|
||||
)
|
||||
|
||||
assert snapshot.pipeline_state == PipelineState(
|
||||
inputs={
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {"comp2": {"type": "object", "properties": {}}},
|
||||
},
|
||||
"serialized_data": {"comp2": {}},
|
||||
},
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
pipeline_outputs={
|
||||
"serialization_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"comp1": {
|
||||
"type": "object",
|
||||
"properties": {"result": {"type": "haystack.dataclasses.chat_message.ChatMessage"}},
|
||||
}
|
||||
},
|
||||
},
|
||||
"serialized_data": {
|
||||
"comp1": {"result": {"role": "user", "meta": {}, "name": None, "content": [{"text": "hello"}]}}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def test_create_pipeline_snapshot_non_serializable_inputs(self, caplog):
|
||||
class NonSerializable:
|
||||
def to_dict(self):
|
||||
raise TypeError("Cannot serialize")
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
_create_pipeline_snapshot(
|
||||
inputs={"comp1": {"input_value": [{"sender": None, "value": NonSerializable()}]}, "comp2": {}},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2"),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={"comp1": {"input_value": NonSerializable()}},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from={"comp1"},
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
assert any("Failed to serialize the inputs of the current pipeline state" in msg for msg in caplog.messages)
|
||||
assert any("Failed to serialize original input data for `pipeline.run`." in msg for msg in caplog.messages)
|
||||
|
||||
def test_create_pipeline_snapshot_non_serializable_inputs_snapshot_is_resumable(self, caplog):
|
||||
"""
|
||||
Guards against the same non-resumable snapshot regression fixed at the agent level: when
|
||||
top-level pipeline inputs/outputs contain non-serializable values, the snapshot fields
|
||||
must still round-trip through ``_deserialize_value_with_schema`` instead of failing with
|
||||
``DeserializationError: ... Got: {}``. Serializable sibling components must stay intact.
|
||||
"""
|
||||
|
||||
class NonSerializable:
|
||||
def to_dict(self):
|
||||
raise TypeError("Cannot serialize")
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={
|
||||
"comp1": {"input_value": [{"sender": None, "value": NonSerializable()}]},
|
||||
"comp2": {"input_value": [{"sender": None, "value": "keep me"}]},
|
||||
},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp3"),
|
||||
component_visits={"comp1": 1, "comp2": 1, "comp3": 0},
|
||||
original_input_data={"comp1": {"input_value": NonSerializable()}},
|
||||
ordered_component_names=["comp1", "comp2", "comp3"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={"comp1": {"result": NonSerializable()}},
|
||||
)
|
||||
|
||||
# No DeserializationError on any of the three pipeline-level payloads.
|
||||
deserialized_inputs = _deserialize_value_with_schema(snapshot.pipeline_state.inputs)
|
||||
deserialized_original_input_data = _deserialize_value_with_schema(snapshot.original_input_data)
|
||||
deserialized_outputs = _deserialize_value_with_schema(snapshot.pipeline_state.pipeline_outputs)
|
||||
|
||||
# The non-serializable comp1 field is omitted while the serializable siblings are preserved.
|
||||
assert "comp1" not in deserialized_inputs
|
||||
assert deserialized_inputs["comp2"] == {"input_value": "keep me"}
|
||||
assert deserialized_inputs["comp3"] == {}
|
||||
# original_input_data and pipeline_outputs degrade to empty-but-valid payloads.
|
||||
assert deserialized_original_input_data == {}
|
||||
assert deserialized_outputs == {}
|
||||
assert any("Failed to serialize the inputs of the current pipeline state" in msg for msg in caplog.messages)
|
||||
assert any("Failed to serialize outputs of the current pipeline state" in msg for msg in caplog.messages)
|
||||
|
||||
|
||||
def test_save_pipeline_snapshot_raises_on_failure(tmp_path, caplog, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from={"comp1"},
|
||||
# We use a non-serializable type (bytes) directly in pipeline outputs to trigger the error
|
||||
pipeline_outputs={"comp1": {"result": b"test"}},
|
||||
)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
_save_pipeline_snapshot(snapshot)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
_save_pipeline_snapshot(snapshot, raise_on_failure=False)
|
||||
assert any("Failed to save pipeline snapshot to" in msg for msg in caplog.messages)
|
||||
|
||||
|
||||
class TestSnapshotCallback:
|
||||
def test_save_pipeline_snapshot_with_callback_no_file_created(self, tmp_path):
|
||||
captured_snapshots = []
|
||||
|
||||
def custom_callback(snapshot: PipelineSnapshot) -> str:
|
||||
captured_snapshots.append(snapshot)
|
||||
return "custom_path_or_id"
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot, snapshot_callback=custom_callback)
|
||||
|
||||
# Verify callback was invoked and returned expected value
|
||||
assert result == "custom_path_or_id"
|
||||
assert len(captured_snapshots) == 1
|
||||
assert captured_snapshots[0] == snapshot
|
||||
|
||||
# Verify NO file was created on disk (callback bypasses file saving)
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_save_pipeline_snapshot_callback_returns_none_no_file_created(self, tmp_path):
|
||||
captured_snapshots = []
|
||||
|
||||
def custom_callback(snapshot: PipelineSnapshot) -> None:
|
||||
captured_snapshots.append(snapshot)
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot, snapshot_callback=custom_callback)
|
||||
|
||||
assert result is None
|
||||
assert len(captured_snapshots) == 1
|
||||
|
||||
# Verify NO file was created on disk even when snapshot_file_path is set
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_save_pipeline_snapshot_without_callback_creates_file(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot)
|
||||
|
||||
# Verify file WAS created on disk
|
||||
snapshot_files = list(tmp_path.glob("comp2_*.json"))
|
||||
|
||||
# A file should be created when no callback is provided
|
||||
assert len(snapshot_files) == 1
|
||||
assert result == str(snapshot_files[0])
|
||||
|
||||
# Verify file contains valid snapshot data
|
||||
loaded = load_pipeline_snapshot(snapshot_files[0])
|
||||
assert isinstance(loaded.break_point, Breakpoint)
|
||||
assert loaded.break_point.component_name == "comp2"
|
||||
|
||||
def test_save_pipeline_snapshot_callback_raises_exception_no_file_created(self, tmp_path, caplog):
|
||||
def failing_callback(snapshot: PipelineSnapshot) -> str:
|
||||
raise RuntimeError("Database connection failed")
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
# Test with raise_on_failure=True (default)
|
||||
with pytest.raises(RuntimeError, match="Database connection failed"):
|
||||
_save_pipeline_snapshot(snapshot, snapshot_callback=failing_callback)
|
||||
|
||||
# Verify NO file was created even after exception
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
# Test with raise_on_failure=False
|
||||
with caplog.at_level(logging.ERROR):
|
||||
result = _save_pipeline_snapshot(snapshot, raise_on_failure=False, snapshot_callback=failing_callback)
|
||||
assert result is None
|
||||
assert any("Failed to handle pipeline snapshot with custom callback" in msg for msg in caplog.messages)
|
||||
|
||||
# Still no file should exist
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_pipeline_run_with_snapshot_callback(self, tmp_path):
|
||||
captured_snapshots = []
|
||||
|
||||
def custom_callback(snapshot: PipelineSnapshot) -> str:
|
||||
captured_snapshots.append(snapshot)
|
||||
return "custom_snapshot_id"
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, input_value: str) -> dict[str, str]:
|
||||
return {"result": f"processed_{input_value}"}
|
||||
|
||||
pipeline = Pipeline()
|
||||
comp1 = SimpleComponent()
|
||||
comp2 = SimpleComponent()
|
||||
pipeline.add_component("comp1", comp1)
|
||||
pipeline.add_component("comp2", comp2)
|
||||
pipeline.connect("comp1", "comp2")
|
||||
|
||||
# breakpoint on comp2
|
||||
break_point = Breakpoint(component_name="comp2", visit_count=0, snapshot_file_path=str(tmp_path))
|
||||
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(
|
||||
data={"comp1": {"input_value": "test"}}, break_point=break_point, snapshot_callback=custom_callback
|
||||
)
|
||||
|
||||
# Verify callback was called
|
||||
assert len(captured_snapshots) == 1
|
||||
assert isinstance(captured_snapshots[0].break_point, Breakpoint)
|
||||
assert captured_snapshots[0].break_point.component_name == "comp2"
|
||||
# Verify the file path in exception is from callback
|
||||
assert exc_info.value.pipeline_snapshot_file_path == "custom_snapshot_id"
|
||||
# Verify no file was saved to disk
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_pipeline_run_without_snapshot_callback_saves_file(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, input_value: str) -> dict[str, str]:
|
||||
return {"result": f"processed_{input_value}"}
|
||||
|
||||
pipeline = Pipeline()
|
||||
comp1 = SimpleComponent()
|
||||
comp2 = SimpleComponent()
|
||||
pipeline.add_component("comp1", comp1)
|
||||
pipeline.add_component("comp2", comp2)
|
||||
pipeline.connect("comp1", "comp2")
|
||||
|
||||
# breakpoint on comp2
|
||||
break_point = Breakpoint(component_name="comp2", visit_count=0, snapshot_file_path=str(tmp_path))
|
||||
|
||||
with pytest.raises(BreakpointException):
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=break_point)
|
||||
|
||||
# Verify file was saved to disk
|
||||
snapshot_files = list(tmp_path.glob("comp2_*.json"))
|
||||
assert len(snapshot_files) == 1
|
||||
|
||||
|
||||
class TestSnapshotSaveEnabled:
|
||||
def test_is_snapshot_save_enabled_default(self, monkeypatch):
|
||||
monkeypatch.delenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, raising=False)
|
||||
assert _is_snapshot_save_enabled() is False
|
||||
|
||||
@pytest.mark.parametrize("value", ["true", "TRUE", "True", "1"])
|
||||
def test_is_snapshot_save_enabled_truthy_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, value)
|
||||
assert _is_snapshot_save_enabled() is True
|
||||
|
||||
@pytest.mark.parametrize("value", ["false", "FALSE", "False", "0"])
|
||||
def test_is_snapshot_save_enabled_falsy_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, value)
|
||||
assert _is_snapshot_save_enabled() is False
|
||||
|
||||
def test_save_pipeline_snapshot_disabled_via_env_var(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "false")
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot)
|
||||
|
||||
# Verify no file was created
|
||||
assert result is None
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_save_pipeline_snapshot_enabled_via_env_var(self, tmp_path, monkeypatch):
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot)
|
||||
|
||||
# Verify file was created
|
||||
snapshot_files = list(tmp_path.glob("comp2_*.json"))
|
||||
assert len(snapshot_files) == 1
|
||||
assert result == str(snapshot_files[0])
|
||||
|
||||
def test_callback_still_invoked_when_env_var_disables_saving(self, tmp_path, monkeypatch):
|
||||
"""
|
||||
This is more a behaviour documentation test: we want to ensure that when the snapshot_callback is provided,
|
||||
the file-saving behaviour is always bypassed (the callback is invoked instead).
|
||||
"""
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "false")
|
||||
|
||||
captured_snapshots = []
|
||||
|
||||
def custom_callback(snapshot: PipelineSnapshot) -> str:
|
||||
captured_snapshots.append(snapshot)
|
||||
return "custom_result"
|
||||
|
||||
snapshot = _create_pipeline_snapshot(
|
||||
inputs={},
|
||||
component_inputs={},
|
||||
break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
|
||||
component_visits={"comp1": 1, "comp2": 0},
|
||||
original_input_data={},
|
||||
ordered_component_names=["comp1", "comp2"],
|
||||
include_outputs_from=set(),
|
||||
pipeline_outputs={},
|
||||
)
|
||||
|
||||
result = _save_pipeline_snapshot(snapshot, snapshot_callback=custom_callback)
|
||||
|
||||
# Callback should still be invoked
|
||||
assert result == "custom_result"
|
||||
assert len(captured_snapshots) == 1
|
||||
# No file should be created (callback handles it)
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
def test_pipeline_run_with_env_var_disabled(self, tmp_path, monkeypatch):
|
||||
"""Test that pipeline.run respects the env var when breakpoint is triggered."""
|
||||
monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "false")
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, input_value: str) -> dict[str, str]:
|
||||
return {"result": f"processed_{input_value}"}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("comp1", SimpleComponent())
|
||||
pipeline.add_component("comp2", SimpleComponent())
|
||||
pipeline.connect("comp1", "comp2")
|
||||
|
||||
break_point = Breakpoint(component_name="comp2", visit_count=0, snapshot_file_path=str(tmp_path))
|
||||
|
||||
with pytest.raises(BreakpointException) as exc_info:
|
||||
pipeline.run(data={"comp1": {"input_value": "test"}}, break_point=break_point)
|
||||
|
||||
# Verify no file was saved
|
||||
assert exc_info.value.pipeline_snapshot_file_path is None
|
||||
assert list(tmp_path.glob("*.json")) == []
|
||||
|
||||
# Verify snapshot object is still available for programmatic access
|
||||
assert exc_info.value.pipeline_snapshot is not None
|
||||
assert isinstance(exc_info.value.pipeline_snapshot.break_point, Breakpoint)
|
||||
assert exc_info.value.pipeline_snapshot.break_point.component_name == "comp2"
|
||||
@@ -0,0 +1,643 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pandas import DataFrame
|
||||
|
||||
from haystack.core.component.types import GreedyVariadic, InputSocket, OutputSocket, Variadic
|
||||
from haystack.core.pipeline.component_checks import (
|
||||
_NO_OUTPUT_PRODUCED,
|
||||
all_predecessors_executed,
|
||||
all_socket_predecessors_executed,
|
||||
any_predecessors_provided_input,
|
||||
any_socket_input_received,
|
||||
any_socket_value_from_predecessor_received,
|
||||
are_all_sockets_ready,
|
||||
can_component_run,
|
||||
can_not_receive_inputs_from_pipeline,
|
||||
has_any_trigger,
|
||||
has_lazy_variadic_socket_received_all_inputs,
|
||||
has_socket_received_all_inputs,
|
||||
has_user_input,
|
||||
is_any_greedy_socket_ready,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def basic_component():
|
||||
"""Basic component with one mandatory and one optional input."""
|
||||
return {
|
||||
"instance": "mock_instance",
|
||||
"visits": 0,
|
||||
"input_sockets": {
|
||||
"mandatory_input": InputSocket("mandatory_input", int, senders=["previous_component"]),
|
||||
"optional_input": InputSocket("optional_input", str, default_value="default"),
|
||||
},
|
||||
"output_sockets": {"output": OutputSocket("output", int)},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def variadic_component():
|
||||
"""Component with variadic input."""
|
||||
return {
|
||||
"instance": "mock_instance",
|
||||
"visits": 0,
|
||||
"input_sockets": {
|
||||
"variadic_input": InputSocket("variadic_input", Variadic[int], senders=["previous_component"]),
|
||||
"normal_input": InputSocket("normal_input", str, senders=["another_component"]),
|
||||
},
|
||||
"output_sockets": {"output": OutputSocket("output", int)},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def greedy_variadic_component():
|
||||
"""Component with greedy variadic input."""
|
||||
return {
|
||||
"instance": "mock_instance",
|
||||
"visits": 0,
|
||||
"input_sockets": {
|
||||
"greedy_input": InputSocket(
|
||||
"greedy_input", GreedyVariadic[int], senders=["previous_component", "other_component"]
|
||||
),
|
||||
"normal_input": InputSocket("normal_input", str),
|
||||
},
|
||||
"output_sockets": {"output": OutputSocket("output", int)},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def input_socket_with_sender():
|
||||
"""Regular input socket with a single sender."""
|
||||
socket = InputSocket("test_input", int)
|
||||
socket.senders = ["component1"]
|
||||
return socket
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def variadic_socket_with_senders():
|
||||
"""Variadic input socket with multiple senders."""
|
||||
socket = InputSocket("test_variadic", Variadic[int])
|
||||
socket.senders = ["component1", "component2"]
|
||||
return socket
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def component_with_multiple_sockets(input_socket_with_sender, variadic_socket_with_senders):
|
||||
"""Component with multiple input sockets including both regular and variadic."""
|
||||
return {
|
||||
"instance": "mock_instance",
|
||||
"input_sockets": {
|
||||
"socket1": input_socket_with_sender,
|
||||
"socket2": variadic_socket_with_senders,
|
||||
"socket3": InputSocket("socket3", str), # No senders
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def regular_socket():
|
||||
"""Regular input socket with one sender."""
|
||||
socket = InputSocket("regular", int)
|
||||
socket.senders = ["component1"]
|
||||
return socket
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lazy_variadic_socket():
|
||||
"""Lazy variadic input socket with multiple senders."""
|
||||
socket = InputSocket("lazy_variadic", Variadic[int])
|
||||
socket.senders = ["component1", "component2"]
|
||||
return socket
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def greedy_variadic_socket():
|
||||
"""Greedy variadic input socket with multiple senders."""
|
||||
socket = InputSocket("greedy_variadic", GreedyVariadic[int])
|
||||
socket.senders = ["component1", "component2", "component3"]
|
||||
return socket
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def complex_component(regular_socket, lazy_variadic_socket, greedy_variadic_socket):
|
||||
"""Component with all types of sockets."""
|
||||
return {
|
||||
"instance": "mock_instance",
|
||||
"input_sockets": {
|
||||
"regular": regular_socket,
|
||||
"lazy_var": lazy_variadic_socket,
|
||||
"greedy_var": greedy_variadic_socket,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestCanComponentRun:
|
||||
def test_component_with_all_mandatory_inputs_and_trigger(self, basic_component):
|
||||
"""Checks that the component runs if all mandatory inputs are received and triggered."""
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": 42}]}
|
||||
assert can_component_run(basic_component, inputs) is True
|
||||
|
||||
def test_component_missing_mandatory_input(self, basic_component):
|
||||
"""Checks that the component won't run if mandatory inputs are missing."""
|
||||
inputs = {"optional_input": [{"sender": "previous_component", "value": "test"}]}
|
||||
assert can_component_run(basic_component, inputs) is False
|
||||
|
||||
# We added these tests because a component that returned a pandas dataframe caused the pipeline to fail.
|
||||
# Previously, we compared the value of the socket using '!=' which leads to an error with dataframes.
|
||||
# Instead, we use 'is not' to compare with the sentinel value.
|
||||
def test_sockets_with_ambiguous_truth_value(self, basic_component, greedy_variadic_socket, regular_socket):
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": DataFrame.from_dict([{"value": 42}])}]}
|
||||
|
||||
assert are_all_sockets_ready(basic_component, inputs, only_check_mandatory=True) is True
|
||||
assert any_socket_value_from_predecessor_received(inputs["mandatory_input"]) is True
|
||||
assert any_socket_input_received(inputs["mandatory_input"]) is True
|
||||
assert (
|
||||
has_lazy_variadic_socket_received_all_inputs(
|
||||
basic_component["input_sockets"]["mandatory_input"], inputs["mandatory_input"]
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert has_socket_received_all_inputs(greedy_variadic_socket, inputs["mandatory_input"]) is True
|
||||
assert has_socket_received_all_inputs(regular_socket, inputs["mandatory_input"]) is True
|
||||
|
||||
def test_component_with_no_trigger_but_all_inputs(self, basic_component):
|
||||
"""
|
||||
Test case where all mandatory inputs are present with valid values,
|
||||
but there is no trigger (no new input from predecessor, not first visit).
|
||||
"""
|
||||
basic_component["visits"] = 1
|
||||
inputs = {"mandatory_input": [{"sender": None, "value": 42}]}
|
||||
assert can_component_run(basic_component, inputs) is False
|
||||
|
||||
def test_component_with_multiple_visits(self, basic_component):
|
||||
"""Checks that a component can still be triggered on subsequent visits by a predecessor."""
|
||||
basic_component["visits"] = 2
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": 42}]}
|
||||
assert can_component_run(basic_component, inputs) is True
|
||||
|
||||
def test_component_with_no_inputs_first_visit(self, basic_component):
|
||||
"""Checks that a component with no input sockets can be triggered on its first visit."""
|
||||
basic_component["input_sockets"] = {}
|
||||
assert can_component_run(basic_component, {}) is True
|
||||
|
||||
def test_component_triggered_on_second_visit_with_new_input(self, basic_component):
|
||||
"""
|
||||
Tests that a second visit is triggered if new predecessor input arrives
|
||||
(i.e. visits > 0, but a valid new input from a predecessor is provided).
|
||||
"""
|
||||
# First, simulate that the component has already run once.
|
||||
basic_component["visits"] = 1
|
||||
|
||||
# Now a predecessor provides a new input; this should re-trigger execution.
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": 99}]}
|
||||
assert can_component_run(basic_component, inputs) is True
|
||||
|
||||
|
||||
class TestHasAnyTrigger:
|
||||
def test_trigger_from_predecessor(self, basic_component):
|
||||
"""Ensures that new data from a predecessor can trigger a component."""
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": 42}]}
|
||||
assert has_any_trigger(basic_component, inputs) is True
|
||||
|
||||
def test_trigger_from_user_first_visit(self, basic_component):
|
||||
"""Checks that user input (sender=None) triggers the component on the first visit."""
|
||||
inputs = {"mandatory_input": [{"sender": None, "value": 42}]}
|
||||
assert has_any_trigger(basic_component, inputs) is True
|
||||
|
||||
def test_no_trigger_from_user_after_first_visit(self, basic_component):
|
||||
"""Checks that user input no longer triggers the component after the first visit."""
|
||||
basic_component["visits"] = 1
|
||||
inputs = {"mandatory_input": [{"sender": None, "value": 42}]}
|
||||
assert has_any_trigger(basic_component, inputs) is False
|
||||
|
||||
def test_trigger_without_inputs_first_visit(self, basic_component):
|
||||
"""Checks that a component with no inputs is triggered on the first visit."""
|
||||
basic_component["input_sockets"] = {}
|
||||
assert has_any_trigger(basic_component, {}) is True
|
||||
|
||||
def test_no_trigger_without_inputs_after_first_visit(self, basic_component):
|
||||
"""Checks that on subsequent visits, no inputs means no trigger."""
|
||||
basic_component["input_sockets"] = {}
|
||||
basic_component["visits"] = 1
|
||||
assert has_any_trigger(basic_component, {}) is False
|
||||
|
||||
|
||||
class TestAllMandatorySocketsReady:
|
||||
def test_all_mandatory_sockets_filled(self, basic_component):
|
||||
"""Checks that all mandatory sockets are ready when they have valid input."""
|
||||
inputs = {"mandatory_input": [{"sender": "previous_component", "value": 42}]}
|
||||
assert are_all_sockets_ready(basic_component, inputs) is True
|
||||
|
||||
def test_missing_mandatory_socket(self, basic_component):
|
||||
"""Ensures that if a mandatory socket is missing, the component is not ready."""
|
||||
inputs = {"optional_input": [{"sender": "previous_component", "value": "test"}]}
|
||||
assert are_all_sockets_ready(basic_component, inputs) is False
|
||||
|
||||
def test_variadic_socket_with_input(self, variadic_component):
|
||||
"""Verifies that a variadic socket is considered filled if it has at least one input."""
|
||||
inputs = {
|
||||
"variadic_input": [{"sender": "previous_component", "value": 42}],
|
||||
"normal_input": [{"sender": "previous_component", "value": "test"}],
|
||||
}
|
||||
assert are_all_sockets_ready(variadic_component, inputs) is True
|
||||
|
||||
def test_greedy_variadic_socket(self, greedy_variadic_component):
|
||||
"""Greedy variadic sockets are ready with at least one valid input."""
|
||||
inputs = {
|
||||
"greedy_input": [{"sender": "previous_component", "value": 42}],
|
||||
"normal_input": [{"sender": "previous_component", "value": "test"}],
|
||||
}
|
||||
assert are_all_sockets_ready(greedy_variadic_component, inputs) is True
|
||||
|
||||
def test_greedy_variadic_socket_and_missing_mandatory(self, greedy_variadic_component):
|
||||
"""All mandatory sockets need to be filled even with GreedyVariadic sockets."""
|
||||
inputs = {"greedy_input": [{"sender": "previous_component", "value": 42}]}
|
||||
assert are_all_sockets_ready(greedy_variadic_component, inputs, only_check_mandatory=True) is False
|
||||
|
||||
def test_variadic_socket_no_input(self, variadic_component):
|
||||
"""A variadic socket is not filled if it has zero valid inputs."""
|
||||
inputs = {"normal_input": [{"sender": "previous_component", "value": "test"}]}
|
||||
assert are_all_sockets_ready(variadic_component, inputs) is False
|
||||
|
||||
def test_mandatory_and_optional_sockets(self):
|
||||
input_sockets = {
|
||||
"mandatory": InputSocket(name="mandatory", type=str, senders=["previous_component"]),
|
||||
"optional": InputSocket(name="optional", type=str, senders=["previous_component"], default_value="test"),
|
||||
}
|
||||
|
||||
component = {"input_sockets": input_sockets}
|
||||
inputs = {"mandatory": [{"sender": "previous_component", "value": "hello"}]}
|
||||
assert are_all_sockets_ready(component, inputs) is False
|
||||
assert are_all_sockets_ready(component, inputs, only_check_mandatory=True) is True
|
||||
|
||||
def test_empty_inputs(self, basic_component):
|
||||
"""Checks that if there are no inputs at all, mandatory sockets are not ready."""
|
||||
assert are_all_sockets_ready(basic_component, {}) is False
|
||||
|
||||
def test_no_mandatory_sockets(self, basic_component):
|
||||
"""Ensures that if there are no mandatory sockets, the component is considered ready."""
|
||||
basic_component["input_sockets"] = {
|
||||
"optional_1": InputSocket("optional_1", str, default_value="default1"),
|
||||
"optional_2": InputSocket("optional_2", str, default_value="default2"),
|
||||
}
|
||||
assert are_all_sockets_ready(basic_component, {}) is True
|
||||
|
||||
def test_multiple_mandatory_sockets(self, basic_component):
|
||||
"""Checks readiness when multiple mandatory sockets are defined."""
|
||||
basic_component["input_sockets"] = {
|
||||
"mandatory_1": InputSocket("mandatory_1", int, senders=["previous_component"]),
|
||||
"mandatory_2": InputSocket("mandatory_2", str, senders=["some other component"]),
|
||||
"optional": InputSocket("optional", bool, default_value=False),
|
||||
}
|
||||
inputs = {
|
||||
"mandatory_1": [{"sender": "comp1", "value": 42}],
|
||||
"mandatory_2": [{"sender": "comp2", "value": "test"}],
|
||||
}
|
||||
assert are_all_sockets_ready(basic_component, inputs) is True
|
||||
|
||||
# Missing one mandatory input
|
||||
inputs = {"mandatory_1": [{"sender": "comp1", "value": 42}], "optional": [{"sender": "comp3", "value": True}]}
|
||||
assert are_all_sockets_ready(basic_component, inputs) is False
|
||||
|
||||
|
||||
class TestPredecessorInputDetection:
|
||||
def test_any_predecessors_provided_input_with_predecessor(self, component_with_multiple_sockets):
|
||||
"""
|
||||
Tests detection of predecessor input when a valid predecessor sends data.
|
||||
"""
|
||||
inputs = {"socket1": [{"sender": "component1", "value": 42}], "socket2": [{"sender": None, "value": "test"}]}
|
||||
assert any_predecessors_provided_input(component_with_multiple_sockets, inputs) is True
|
||||
|
||||
def test_any_predecessors_provided_input_no_predecessor(self, component_with_multiple_sockets):
|
||||
"""
|
||||
Checks that no predecessor inputs are detected if all senders are None (user inputs).
|
||||
"""
|
||||
inputs = {"socket1": [{"sender": None, "value": 42}], "socket2": [{"sender": None, "value": "test"}]}
|
||||
assert any_predecessors_provided_input(component_with_multiple_sockets, inputs) is False
|
||||
|
||||
def test_any_predecessors_provided_input_with_no_output(self, component_with_multiple_sockets):
|
||||
"""
|
||||
Ensures that _NO_OUTPUT_PRODUCED from a predecessor is ignored in the predecessor detection.
|
||||
"""
|
||||
inputs = {
|
||||
"socket1": [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}],
|
||||
"socket2": [{"sender": None, "value": "test"}],
|
||||
}
|
||||
assert any_predecessors_provided_input(component_with_multiple_sockets, inputs) is False
|
||||
|
||||
def test_any_predecessors_provided_input_empty_inputs(self, component_with_multiple_sockets):
|
||||
"""Ensures that empty inputs dictionary returns False."""
|
||||
assert any_predecessors_provided_input(component_with_multiple_sockets, {}) is False
|
||||
|
||||
|
||||
class TestSocketValueFromPredecessor:
|
||||
"""
|
||||
Tests for `any_socket_value_from_predecessor_received`, verifying whether
|
||||
any predecessor component provided valid output to a socket.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"socket_inputs, expected_result",
|
||||
[
|
||||
pytest.param([{"sender": "component1", "value": 42}], True, id="valid_input"),
|
||||
pytest.param([{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}], False, id="no_output"),
|
||||
pytest.param([{"sender": None, "value": 42}], False, id="user_input"),
|
||||
pytest.param(
|
||||
[
|
||||
{"sender": None, "value": 42},
|
||||
{"sender": "component1", "value": _NO_OUTPUT_PRODUCED},
|
||||
{"sender": "component2", "value": 100},
|
||||
],
|
||||
True,
|
||||
id="mixed_inputs",
|
||||
),
|
||||
pytest.param([], False, id="empty_list"),
|
||||
],
|
||||
)
|
||||
def test_any_socket_value_from_predecessor_received(self, socket_inputs, expected_result):
|
||||
"""
|
||||
Parametrized test to check whether any valid predecessor input
|
||||
exists in a list of socket inputs.
|
||||
"""
|
||||
assert any_socket_value_from_predecessor_received(socket_inputs) == expected_result
|
||||
|
||||
|
||||
class TestUserInputDetection:
|
||||
def test_has_user_input_with_user_input(self):
|
||||
"""Checks that having a sender=None input means user input is present."""
|
||||
inputs = {"socket1": [{"sender": None, "value": 42}], "socket2": [{"sender": "component1", "value": "test"}]}
|
||||
assert has_user_input(inputs) is True
|
||||
|
||||
def test_has_user_input_without_user_input(self):
|
||||
"""Ensures that if all senders are component-based, there's no user input."""
|
||||
inputs = {
|
||||
"socket1": [{"sender": "component1", "value": 42}],
|
||||
"socket2": [{"sender": "component2", "value": "test"}],
|
||||
}
|
||||
assert has_user_input(inputs) is False
|
||||
|
||||
def test_has_user_input_empty_inputs(self):
|
||||
"""Checks that an empty inputs dict has no user input."""
|
||||
assert has_user_input({}) is False
|
||||
|
||||
def test_has_user_input_with_no_output(self):
|
||||
"""
|
||||
Even if the input value is _NO_OUTPUT_PRODUCED, if sender=None
|
||||
it still counts as user input being provided.
|
||||
"""
|
||||
inputs = {"socket1": [{"sender": None, "value": _NO_OUTPUT_PRODUCED}]}
|
||||
assert has_user_input(inputs) is True
|
||||
|
||||
|
||||
class TestPipelineInputCapability:
|
||||
def test_cannot_receive_inputs_no_senders(self):
|
||||
"""Checks that a component with zero senders for each socket cannot receive pipeline inputs."""
|
||||
component = {"input_sockets": {"socket1": InputSocket("socket1", int), "socket2": InputSocket("socket2", str)}}
|
||||
assert can_not_receive_inputs_from_pipeline(component) is True
|
||||
|
||||
def test_cannot_receive_inputs_with_senders(self, component_with_multiple_sockets):
|
||||
"""If at least one socket has a sender, the component can receive pipeline inputs."""
|
||||
assert can_not_receive_inputs_from_pipeline(component_with_multiple_sockets) is False
|
||||
|
||||
def test_cannot_receive_inputs_mixed_senders(self, input_socket_with_sender):
|
||||
"""A single socket with a sender means the component can receive pipeline inputs."""
|
||||
component = {
|
||||
"input_sockets": {
|
||||
"socket1": input_socket_with_sender,
|
||||
"socket2": InputSocket("socket2", str), # No senders
|
||||
}
|
||||
}
|
||||
assert can_not_receive_inputs_from_pipeline(component) is False
|
||||
|
||||
|
||||
class TestSocketExecutionStatus:
|
||||
def test_regular_socket_predecessor_executed(self, input_socket_with_sender):
|
||||
"""Verifies that if the correct sender provides a value, the socket is marked as executed."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}]
|
||||
assert all_socket_predecessors_executed(input_socket_with_sender, socket_inputs) is True
|
||||
|
||||
def test_regular_socket_predecessor_not_executed(self, input_socket_with_sender):
|
||||
"""If there are no inputs, the predecessor is not considered executed."""
|
||||
assert all_socket_predecessors_executed(input_socket_with_sender, []) is False
|
||||
|
||||
def test_regular_socket_with_wrong_predecessor(self, input_socket_with_sender):
|
||||
"""Checks that a mismatch in sender means the socket is not yet executed."""
|
||||
socket_inputs = [{"sender": "component2", "value": 42}]
|
||||
assert all_socket_predecessors_executed(input_socket_with_sender, socket_inputs) is False
|
||||
|
||||
def test_variadic_socket_all_predecessors_executed(self, variadic_socket_with_senders):
|
||||
"""Variadic socket is executed only if all senders have produced at least one valid result."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}]
|
||||
assert all_socket_predecessors_executed(variadic_socket_with_senders, socket_inputs) is True
|
||||
|
||||
def test_variadic_socket_partial_execution(self, variadic_socket_with_senders):
|
||||
"""If only one of multiple senders produced an output, not all predecessors are executed."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}]
|
||||
assert all_socket_predecessors_executed(variadic_socket_with_senders, socket_inputs) is False
|
||||
|
||||
def test_variadic_socket_with_user_input(self, variadic_socket_with_senders):
|
||||
"""
|
||||
User input (sender=None) doesn't block the socket from being 'executed' if
|
||||
all named predecessors have also produced outputs.
|
||||
"""
|
||||
socket_inputs: list[dict[str, Any]] = [
|
||||
{"sender": "component1", "value": 42},
|
||||
{"sender": None, "value": 43},
|
||||
{"sender": "component2", "value": 44},
|
||||
]
|
||||
assert all_socket_predecessors_executed(variadic_socket_with_senders, socket_inputs) is True
|
||||
|
||||
def test_variadic_socket_no_execution(self, variadic_socket_with_senders):
|
||||
"""Empty inputs means no predecessor has executed."""
|
||||
assert all_socket_predecessors_executed(variadic_socket_with_senders, []) is False
|
||||
|
||||
|
||||
class TestSocketInputReceived:
|
||||
def test_any_socket_input_received_with_value(self):
|
||||
"""Checks that if there's a non-_NO_OUTPUT_PRODUCED value, the socket is marked as having input."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}]
|
||||
assert any_socket_input_received(socket_inputs) is True
|
||||
|
||||
def test_any_socket_input_received_with_no_output(self):
|
||||
"""If all inputs are _NO_OUTPUT_PRODUCED, the socket has no effective input."""
|
||||
socket_inputs = [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}]
|
||||
assert any_socket_input_received(socket_inputs) is False
|
||||
|
||||
def test_any_socket_input_received_mixed_inputs(self):
|
||||
"""A single valid input among many is enough to consider the socket as having input."""
|
||||
socket_inputs = [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}, {"sender": "component2", "value": 42}]
|
||||
assert any_socket_input_received(socket_inputs) is True
|
||||
|
||||
def test_any_socket_input_received_empty_list(self):
|
||||
"""Empty list: no input received."""
|
||||
assert any_socket_input_received([]) is False
|
||||
|
||||
|
||||
class TestLazyVariadicSocket:
|
||||
def test_lazy_variadic_all_inputs_received(self, variadic_socket_with_senders):
|
||||
"""Lazy variadic socket is ready only if all named senders provided outputs."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}]
|
||||
assert has_lazy_variadic_socket_received_all_inputs(variadic_socket_with_senders, socket_inputs) is True
|
||||
|
||||
def test_lazy_variadic_partial_inputs(self, variadic_socket_with_senders):
|
||||
"""Partial inputs from only some senders is insufficient for a lazy variadic socket."""
|
||||
socket_inputs = [{"sender": "component1", "value": 42}]
|
||||
assert has_lazy_variadic_socket_received_all_inputs(variadic_socket_with_senders, socket_inputs) is False
|
||||
|
||||
def test_lazy_variadic_with_no_output(self, variadic_socket_with_senders):
|
||||
"""_NO_OUTPUT_PRODUCED from a sender doesn't count as valid input, so it's not fully received."""
|
||||
socket_inputs = [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}, {"sender": "component2", "value": 42}]
|
||||
assert has_lazy_variadic_socket_received_all_inputs(variadic_socket_with_senders, socket_inputs) is False
|
||||
|
||||
def test_lazy_variadic_with_user_input(self, variadic_socket_with_senders):
|
||||
"""
|
||||
User input doesn't block a lazy variadic socket, as long as all named senders
|
||||
also provided outputs.
|
||||
"""
|
||||
socket_inputs: list[dict[str, Any]] = [
|
||||
{"sender": "component1", "value": 42},
|
||||
{"sender": None, "value": 43},
|
||||
{"sender": "component2", "value": 44},
|
||||
]
|
||||
assert has_lazy_variadic_socket_received_all_inputs(variadic_socket_with_senders, socket_inputs) is True
|
||||
|
||||
def test_lazy_variadic_empty_inputs(self, variadic_socket_with_senders):
|
||||
"""No inputs at all means the lazy variadic socket hasn't received everything yet."""
|
||||
assert has_lazy_variadic_socket_received_all_inputs(variadic_socket_with_senders, []) is False
|
||||
|
||||
|
||||
class TestSocketTypeDetection:
|
||||
def test_is_socket_lazy_variadic_with_lazy_socket(self, lazy_variadic_socket):
|
||||
"""Ensures that a non-greedy variadic socket is detected as lazy."""
|
||||
assert lazy_variadic_socket.is_lazy_variadic is True
|
||||
assert lazy_variadic_socket.is_greedy is False
|
||||
|
||||
def test_is_socket_lazy_variadic_with_greedy_socket(self, greedy_variadic_socket):
|
||||
"""Greedy variadic sockets should not be marked as lazy."""
|
||||
assert greedy_variadic_socket.is_lazy_variadic is False
|
||||
assert greedy_variadic_socket.is_greedy is True
|
||||
|
||||
def test_is_socket_lazy_variadic_with_regular_socket(self, regular_socket):
|
||||
"""Regular sockets are not variadic at all."""
|
||||
assert regular_socket.is_lazy_variadic is False
|
||||
assert regular_socket.is_greedy is False
|
||||
|
||||
|
||||
class TestSocketInputCompletion:
|
||||
def test_regular_socket_complete(self, regular_socket):
|
||||
"""A single valid input marks a regular socket as complete."""
|
||||
inputs = [{"sender": "component1", "value": 42}]
|
||||
assert has_socket_received_all_inputs(regular_socket, inputs) is True
|
||||
|
||||
def test_regular_socket_incomplete(self, regular_socket):
|
||||
"""_NO_OUTPUT_PRODUCED means the socket is not complete."""
|
||||
inputs = [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}]
|
||||
assert has_socket_received_all_inputs(regular_socket, inputs) is False
|
||||
|
||||
def test_regular_socket_no_inputs(self, regular_socket):
|
||||
"""No inputs at all means the socket is incomplete."""
|
||||
assert has_socket_received_all_inputs(regular_socket, []) is False
|
||||
|
||||
def test_lazy_variadic_socket_all_inputs(self, lazy_variadic_socket):
|
||||
"""Lazy variadic socket is complete only if all senders have produced valid outputs."""
|
||||
inputs = [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}]
|
||||
assert has_socket_received_all_inputs(lazy_variadic_socket, inputs) is True
|
||||
|
||||
def test_lazy_variadic_socket_partial_inputs(self, lazy_variadic_socket):
|
||||
"""Partial coverage of senders is insufficient for lazy variadic sockets."""
|
||||
inputs = [{"sender": "component1", "value": 42}]
|
||||
assert has_socket_received_all_inputs(lazy_variadic_socket, inputs) is False
|
||||
|
||||
def test_lazy_variadic_socket_with_no_output(self, lazy_variadic_socket):
|
||||
"""A sender that produces _NO_OUTPUT_PRODUCED does not fulfill the lazy socket requirement."""
|
||||
inputs = [{"sender": "component1", "value": 42}, {"sender": "component2", "value": _NO_OUTPUT_PRODUCED}]
|
||||
assert has_socket_received_all_inputs(lazy_variadic_socket, inputs) is False
|
||||
|
||||
def test_greedy_variadic_socket_one_input(self, greedy_variadic_socket):
|
||||
"""A greedy variadic socket is complete if it has at least one valid input."""
|
||||
inputs = [{"sender": "component1", "value": 42}]
|
||||
assert has_socket_received_all_inputs(greedy_variadic_socket, inputs) is True
|
||||
|
||||
def test_greedy_variadic_socket_multiple_inputs(self, greedy_variadic_socket):
|
||||
"""A greedy variadic socket with multiple inputs remains complete as soon as one is valid."""
|
||||
inputs = [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}]
|
||||
assert has_socket_received_all_inputs(greedy_variadic_socket, inputs) is True
|
||||
|
||||
def test_greedy_variadic_socket_no_valid_inputs(self, greedy_variadic_socket):
|
||||
"""All _NO_OUTPUT_PRODUCED means the greedy socket is not complete."""
|
||||
inputs = [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}]
|
||||
assert has_socket_received_all_inputs(greedy_variadic_socket, inputs) is False
|
||||
|
||||
|
||||
class TestPredecessorExecution:
|
||||
def test_all_predecessors_executed_complete(self, complex_component):
|
||||
"""
|
||||
Checks that if all named senders produce valid outputs for each socket,
|
||||
then all predecessors are considered executed.
|
||||
"""
|
||||
inputs = {
|
||||
"regular": [{"sender": "component1", "value": 42}],
|
||||
"lazy_var": [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}],
|
||||
"greedy_var": [
|
||||
{"sender": "component1", "value": 42},
|
||||
{"sender": "component2", "value": 43},
|
||||
{"sender": "component3", "value": 44},
|
||||
],
|
||||
}
|
||||
assert all_predecessors_executed(complex_component, inputs) is True
|
||||
|
||||
def test_all_predecessors_executed_partial(self, complex_component):
|
||||
"""If a lazy socket is missing one predecessor, not all predecessors are executed."""
|
||||
inputs = {
|
||||
"regular": [{"sender": "component1", "value": 42}],
|
||||
"lazy_var": [{"sender": "component1", "value": 42}], # Missing component2
|
||||
"greedy_var": [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}],
|
||||
}
|
||||
assert all_predecessors_executed(complex_component, inputs) is False
|
||||
|
||||
def test_all_predecessors_executed_with_user_input(self, complex_component):
|
||||
"""
|
||||
User input shouldn't affect predecessor execution for the lazy socket:
|
||||
we still need all named senders to produce output.
|
||||
"""
|
||||
inputs = {
|
||||
"regular": [{"sender": "component1", "value": 42}],
|
||||
"lazy_var": [{"sender": "component1", "value": 42}, {"sender": None, "value": 43}],
|
||||
"greedy_var": [
|
||||
{"sender": "component1", "value": 42},
|
||||
{"sender": "component2", "value": 43},
|
||||
{"sender": "component3", "value": 44},
|
||||
],
|
||||
}
|
||||
assert all_predecessors_executed(complex_component, inputs) is False
|
||||
|
||||
|
||||
class TestGreedySocketReadiness:
|
||||
def test_greedy_socket_ready(self, complex_component):
|
||||
"""A single valid input is enough for a greedy variadic socket to be considered ready."""
|
||||
inputs = {"greedy_var": [{"sender": "component1", "value": 42}]}
|
||||
assert is_any_greedy_socket_ready(complex_component, inputs) is True
|
||||
|
||||
def test_greedy_socket_multiple_inputs_ready(self, complex_component):
|
||||
"""Multiple valid inputs on a greedy socket is also fine—it's still ready."""
|
||||
inputs = {"greedy_var": [{"sender": "component1", "value": 42}, {"sender": "component2", "value": 43}]}
|
||||
assert is_any_greedy_socket_ready(complex_component, inputs) is True
|
||||
|
||||
def test_greedy_socket_not_ready(self, complex_component):
|
||||
"""If the only input is _NO_OUTPUT_PRODUCED, the greedy socket isn't ready."""
|
||||
inputs = {"greedy_var": [{"sender": "component1", "value": _NO_OUTPUT_PRODUCED}]}
|
||||
assert is_any_greedy_socket_ready(complex_component, inputs) is False
|
||||
|
||||
def test_greedy_socket_no_inputs(self, complex_component):
|
||||
"""No inputs at all: the greedy socket is not ready."""
|
||||
assert is_any_greedy_socket_ready(complex_component, {}) is False
|
||||
|
||||
def test_greedy_socket_with_user_input(self, complex_component):
|
||||
"""User input can also trigger readiness for a greedy variadic socket."""
|
||||
inputs = {"greedy_var": [{"sender": None, "value": 42}]}
|
||||
assert is_any_greedy_socket_ready(complex_component, inputs) is True
|
||||
@@ -0,0 +1,260 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from haystack.core.errors import PipelineDrawingError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.pipeline.draw import _to_mermaid_image, _to_mermaid_text, _validate_image_response
|
||||
from haystack.testing.sample_components import AddFixedValue, Double
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Temporarily skipped due to mermaid.ink issues")
|
||||
@pytest.mark.integration
|
||||
def test_to_mermaid_image():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
image_data = _to_mermaid_image(pipe.graph)
|
||||
# We just verify we received some data as testing the actual image is not reliable
|
||||
assert image_data
|
||||
|
||||
|
||||
@patch("haystack.core.pipeline.draw.httpx")
|
||||
def test_to_mermaid_image_does_not_edit_graph(mock_httpx):
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=3))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1.result", "comp2.value")
|
||||
pipe.connect("comp2.value", "comp1.value")
|
||||
|
||||
mock_httpx.get.return_value = MagicMock(
|
||||
status_code=200, content=b"\x89PNG\r\n\x1a\n", headers={"content-type": "image/png"}
|
||||
)
|
||||
expected_pipe = pipe.to_dict()
|
||||
_to_mermaid_image(pipe.graph)
|
||||
assert expected_pipe == pipe.to_dict()
|
||||
|
||||
|
||||
@patch("haystack.core.pipeline.draw.httpx")
|
||||
def test_to_mermaid_image_applies_timeout(mock_httpx):
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
mock_httpx.get.return_value = MagicMock(
|
||||
status_code=200, content=b"\x89PNG\r\n\x1a\n", headers={"content-type": "image/png"}
|
||||
)
|
||||
_to_mermaid_image(pipe.graph, timeout=1)
|
||||
|
||||
assert mock_httpx.get.call_args[1]["timeout"] == 1
|
||||
|
||||
|
||||
def test_to_mermaid_image_failing_request(tmp_path):
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
pipe.connect("comp2", "comp1")
|
||||
|
||||
with patch("haystack.core.pipeline.draw.httpx.get") as mock_get:
|
||||
|
||||
def raise_for_status(self):
|
||||
raise httpx.HTTPError("error")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 429
|
||||
mock_response.content = '{"error": "too many requests"}'
|
||||
mock_response.raise_for_status = raise_for_status
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
with pytest.raises(PipelineDrawingError, match="There was an issue with https://mermaid.ink"):
|
||||
_to_mermaid_image(pipe.graph)
|
||||
|
||||
|
||||
def test_to_mermaid_text():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=3))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1.result", "comp2.value")
|
||||
pipe.connect("comp2.value", "comp1.value")
|
||||
|
||||
init_params = {"theme": "neutral"}
|
||||
text = _to_mermaid_text(pipe.graph, init_params)
|
||||
assert (
|
||||
text
|
||||
== """
|
||||
%%{ init: {'theme': 'neutral'} }%%
|
||||
|
||||
graph TD;
|
||||
|
||||
comp1["<b>comp1</b><br><small><i>AddFixedValue<br><br>Optional inputs:<ul style='text-align:left;'><li>add (int | None)</li></ul></i></small>"]:::component -- "result -> value<br><small><i>int</i></small>" --> comp2["<b>comp2</b><br><small><i>Double</i></small>"]:::component
|
||||
comp2["<b>comp2</b><br><small><i>Double</i></small>"]:::component -- "value -> value<br><small><i>int</i></small>" --> comp1["<b>comp1</b><br><small><i>AddFixedValue<br><br>Optional inputs:<ul style='text-align:left;'><li>add (int | None)</li></ul></i></small>"]:::component
|
||||
|
||||
classDef component text-align:center;
|
||||
|
||||
""" # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def test_to_mermaid_text_does_not_edit_graph():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=3))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1.result", "comp2.value")
|
||||
pipe.connect("comp2.value", "comp1.value")
|
||||
|
||||
expected_pipe = pipe.to_dict()
|
||||
init_params = {"theme": "neutral"}
|
||||
_to_mermaid_text(pipe.graph, init_params)
|
||||
assert expected_pipe == pipe.to_dict()
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="This is a nice to have, but frequently fails due to mermaid.ink issues")
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.parametrize(
|
||||
"params",
|
||||
[
|
||||
{"format": "img", "type": "png", "theme": "dark"},
|
||||
{"format": "svg", "theme": "forest"},
|
||||
{"format": "pdf", "fit": True, "theme": "neutral"},
|
||||
],
|
||||
)
|
||||
def test_to_mermaid_image_valid_formats(params):
|
||||
# Test valid formats
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
image_data = _to_mermaid_image(pipe.graph, params=params)
|
||||
assert image_data # Ensure some data is returned
|
||||
|
||||
|
||||
def test_to_mermaid_image_invalid_format():
|
||||
# Test invalid format
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid image format:"):
|
||||
_to_mermaid_image(pipe.graph, params={"format": "invalid_format"})
|
||||
|
||||
|
||||
def test_to_mermaid_image_invalid_scale():
|
||||
# Test invalid scale
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
with pytest.raises(ValueError, match="Scale must be a number between 1 and 3."):
|
||||
_to_mermaid_image(pipe.graph, params={"format": "img", "scale": 5})
|
||||
|
||||
|
||||
def test_to_mermaid_image_scale_without_dimensions():
|
||||
# Test scale without width/height
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
with pytest.raises(ValueError, match="Scale is only allowed when width or height is set."):
|
||||
_to_mermaid_image(pipe.graph, params={"format": "img", "scale": 2})
|
||||
|
||||
|
||||
@patch("haystack.core.pipeline.draw.httpx.get")
|
||||
def test_to_mermaid_image_server_error(mock_get):
|
||||
# Test server failure
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
def raise_for_status(self):
|
||||
raise httpx.HTTPError("error")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.content = '{"error": "server error"}'
|
||||
mock_response.raise_for_status = raise_for_status
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
with pytest.raises(PipelineDrawingError, match="There was an issue with https://mermaid.ink"):
|
||||
_to_mermaid_image(pipe.graph)
|
||||
|
||||
|
||||
def test_to_mermaid_image_invalid_server_url():
|
||||
# Test invalid server URL
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=3))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1.result", "comp2.value")
|
||||
pipe.connect("comp2.value", "comp1.value")
|
||||
|
||||
server_url = "https://invalid.server"
|
||||
|
||||
with pytest.raises(PipelineDrawingError, match=f"There was an issue with {server_url}"):
|
||||
_to_mermaid_image(pipe.graph, server_url=server_url)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"params, content, content_type",
|
||||
[
|
||||
({"format": "img", "type": "png"}, b"\x89PNG\r\n\x1a\n" + b"rest", "image/png"),
|
||||
({"format": "img", "type": "jpeg"}, b"\xff\xd8\xff" + b"rest", "image/jpeg"),
|
||||
({"format": "img", "type": "webp"}, b"RIFF\x00\x00\x00\x00WEBPrest", "image/webp"),
|
||||
({"format": "svg"}, b'<?xml version="1.0"?><svg></svg>', "image/svg+xml"),
|
||||
({"format": "svg"}, b"<svg xmlns='...'></svg>", "image/svg+xml"),
|
||||
({"format": "pdf"}, b"%PDF-1.7\nrest", "application/pdf"),
|
||||
],
|
||||
)
|
||||
def test_validate_image_response_accepts_expected_formats(params, content, content_type):
|
||||
resp = MagicMock(content=content, headers={"content-type": content_type})
|
||||
# Should not raise
|
||||
_validate_image_response(resp, params)
|
||||
|
||||
|
||||
def test_validate_image_response_rejects_empty_body():
|
||||
resp = MagicMock(content=b"", headers={"content-type": "image/png"})
|
||||
with pytest.raises(PipelineDrawingError, match="empty response"):
|
||||
_validate_image_response(resp, {"format": "img", "type": "png"})
|
||||
|
||||
|
||||
def test_validate_image_response_rejects_mismatched_body():
|
||||
# A server returning an HTML error page (or attacker-controlled payload) while the caller
|
||||
# expects a PNG must be rejected so it never gets written to disk.
|
||||
resp = MagicMock(content=b"<html><body>error</body></html>", headers={"content-type": "image/png"})
|
||||
with pytest.raises(PipelineDrawingError, match="does not look like a valid PNG image"):
|
||||
_validate_image_response(resp, {"format": "img", "type": "png"})
|
||||
|
||||
|
||||
def test_validate_image_response_warns_on_spoofed_content_type_but_relies_on_body(caplog):
|
||||
# Content-Type is server-controlled, so a wrong header only warns as long as the body is valid.
|
||||
resp = MagicMock(content=b"\x89PNG\r\n\x1a\n" + b"rest", headers={"content-type": "text/html"})
|
||||
_validate_image_response(resp, {"format": "img", "type": "png"})
|
||||
assert "unexpected Content-Type" in caplog.text
|
||||
|
||||
|
||||
@patch("haystack.core.pipeline.draw.httpx.get")
|
||||
def test_to_mermaid_image_rejects_non_image_response(mock_get):
|
||||
# End-to-end: a 200 response with non-image content must not be returned for writing to disk.
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
mock_get.return_value = MagicMock(
|
||||
status_code=200, content=b"#!/bin/sh\nrm -rf /\n", headers={"content-type": "image/png"}
|
||||
)
|
||||
|
||||
with pytest.raises(PipelineDrawingError, match="does not look like a valid PNG image"):
|
||||
_to_mermaid_image(pipe.graph)
|
||||
@@ -0,0 +1,559 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.components.agents import Agent
|
||||
from haystack.components.joiners import BranchJoiner
|
||||
from haystack.core.component import component
|
||||
from haystack.core.errors import PipelineConnectError, PipelineRuntimeError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.dataclasses import ChatMessage, Document
|
||||
|
||||
|
||||
@component
|
||||
class MockChatGenerator:
|
||||
@component.output_types(replies=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||||
return {"replies": [ChatMessage.from_assistant("Hello, world!")]}
|
||||
|
||||
|
||||
@component
|
||||
class StringProducer:
|
||||
def __init__(self, value: str = "Hello"):
|
||||
self.value = value
|
||||
|
||||
@component.output_types(text=str)
|
||||
def run(self) -> dict[str, str]:
|
||||
return {"text": self.value}
|
||||
|
||||
|
||||
@component
|
||||
class ListStrProducer:
|
||||
def __init__(self, values: list[str] | None = None):
|
||||
self.values = values or ["Hello", "Hi"]
|
||||
|
||||
@component.output_types(texts=list[str])
|
||||
def run(self) -> dict[str, list[str]]:
|
||||
return {"texts": self.values}
|
||||
|
||||
|
||||
@component
|
||||
class ListStrAcceptor:
|
||||
@component.output_types(result=list[str])
|
||||
def run(self, texts: list[str]) -> dict[str, list[str]]:
|
||||
return {"result": texts}
|
||||
|
||||
|
||||
@component
|
||||
class ChatMessageProducer:
|
||||
def __init__(self, value: str = "Hello"):
|
||||
self.value = value
|
||||
|
||||
@component.output_types(message=ChatMessage)
|
||||
def run(self) -> dict[str, ChatMessage]:
|
||||
return {"message": ChatMessage.from_user(self.value)}
|
||||
|
||||
|
||||
@component
|
||||
class ListChatMessageProducer:
|
||||
def __init__(self, values: list[str] | None = None):
|
||||
self.values = values or ["Hello", "Hi"]
|
||||
|
||||
@component.output_types(messages=list[ChatMessage])
|
||||
def run(self) -> dict[str, list[ChatMessage]]:
|
||||
return {"messages": [ChatMessage.from_user(v) for v in self.values]}
|
||||
|
||||
|
||||
@component
|
||||
class ListChatMessageAcceptor:
|
||||
@component.output_types(result=list[ChatMessage])
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
|
||||
return {"result": messages}
|
||||
|
||||
|
||||
@component
|
||||
class WrongOutput:
|
||||
@component.output_types(output=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return "not_a_dict" # type: ignore[return-value]
|
||||
|
||||
|
||||
class TestPipeline:
|
||||
"""
|
||||
This class contains only unit tests for the Pipeline class.
|
||||
|
||||
It doesn't test Pipeline.run(), that is done separately in a different way.
|
||||
"""
|
||||
|
||||
def test_pipeline_thread_safety(self, waiting_component, spying_tracer):
|
||||
# Initialize pipeline with synchronous components
|
||||
pp = Pipeline()
|
||||
pp.add_component("wait", waiting_component())
|
||||
|
||||
run_data = [{"wait_for": 0.001}, {"wait_for": 0.002}]
|
||||
|
||||
# Use ThreadPoolExecutor to run pipeline calls in parallel
|
||||
with ThreadPoolExecutor(max_workers=len(run_data)) as executor:
|
||||
# Submit pipeline runs to the executor
|
||||
futures = [executor.submit(pp.run, data) for data in run_data]
|
||||
|
||||
# Wait for all futures to complete
|
||||
for future in futures:
|
||||
future.result()
|
||||
|
||||
# Verify component visits using tracer
|
||||
component_spans = [sp for sp in spying_tracer.spans if sp.operation_name == "haystack.component.run"]
|
||||
|
||||
for span in component_spans:
|
||||
assert span.tags["haystack.component.visits"] == 1
|
||||
|
||||
def test_prepare_component_inputs(self):
|
||||
pp = Pipeline()
|
||||
component_name = "joiner_1"
|
||||
pp.add_component(component_name, BranchJoiner(type_=str))
|
||||
pp.add_component("joiner_2", BranchJoiner(type_=str))
|
||||
pp.connect(component_name, "joiner_2")
|
||||
inputs = {"joiner_1": {"value": [{"sender": None, "value": "test_value"}]}}
|
||||
comp_dict = pp._get_component_with_graph_metadata_and_visits(component_name, 0)
|
||||
|
||||
_ = pp._consume_component_inputs(component_name=component_name, component=comp_dict, inputs=inputs)
|
||||
# We remove input in greedy variadic sockets, even if they are from the user
|
||||
assert inputs == {"joiner_1": {}}
|
||||
|
||||
def test__run_component_success(self):
|
||||
"""Test successful component execution"""
|
||||
pp = Pipeline()
|
||||
component_name = "joiner_1"
|
||||
pp.add_component(component_name, BranchJoiner(type_=str))
|
||||
pp.add_component("joiner_2", BranchJoiner(type_=str))
|
||||
pp.connect(component_name, "joiner_2")
|
||||
|
||||
outputs = pp._run_component(
|
||||
component_name=component_name,
|
||||
component=pp._get_component_with_graph_metadata_and_visits(component_name, 0),
|
||||
inputs={"value": ["test_value"]},
|
||||
component_visits={component_name: 0, "joiner_2": 0},
|
||||
)
|
||||
assert outputs == {"value": "test_value"}
|
||||
|
||||
def test__run_component_fail(self):
|
||||
"""Test error when component doesn't return a dictionary"""
|
||||
pp = Pipeline()
|
||||
pp.add_component("wrong", WrongOutput())
|
||||
|
||||
with pytest.raises(PipelineRuntimeError) as exc_info:
|
||||
pp._run_component(
|
||||
component_name="wrong",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("wrong", 0),
|
||||
inputs={"value": "test_value"},
|
||||
component_visits={"wrong": 0},
|
||||
)
|
||||
assert "Expected a dict" in str(exc_info.value)
|
||||
|
||||
def test_run_component_error(self):
|
||||
"""Test error when component fails to run"""
|
||||
|
||||
@component
|
||||
class ErroringComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self):
|
||||
raise ValueError("Test error")
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("erroring_component", ErroringComponent())
|
||||
|
||||
with pytest.raises(PipelineRuntimeError) as exc_info:
|
||||
pp._run_component(
|
||||
component_name="erroring_component",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("erroring_component", 0),
|
||||
inputs={"wrong": {"value": [{"sender": None, "value": "test_value"}]}},
|
||||
component_visits={"erroring_component": 0},
|
||||
)
|
||||
assert "Component name: 'erroring_component'" in str(exc_info.value)
|
||||
|
||||
def test_component_with_empty_dict_as_output_appears_in_results(self):
|
||||
"""Test that components that return an empty dict as output appear in results as an empty dict"""
|
||||
|
||||
@component
|
||||
class Producer:
|
||||
def __init__(self, prefix: str):
|
||||
self.prefix = prefix
|
||||
|
||||
@component.output_types(value=str | None)
|
||||
def run(self, text: str | None) -> dict[str, str | None]:
|
||||
return {"value": f"{self.prefix}: {text}"}
|
||||
|
||||
@component
|
||||
class EmptyProcessor:
|
||||
@component.output_types()
|
||||
def run(self, sources: list[str]) -> dict:
|
||||
# Returns empty dict when sources is empty
|
||||
return {}
|
||||
|
||||
@component
|
||||
class Combiner:
|
||||
@component.output_types(combined=str)
|
||||
def run(self, input_a: str | None, input_b: str | None) -> dict[str, str]:
|
||||
if input_a is None:
|
||||
input_a = ""
|
||||
if input_b is None:
|
||||
input_b = ""
|
||||
return {"combined": f"{input_a} | {input_b}"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("producer_a", Producer("A"))
|
||||
pp.add_component("producer_b", Producer("B"))
|
||||
pp.add_component("empty_processor", EmptyProcessor())
|
||||
pp.add_component("combiner", Combiner())
|
||||
|
||||
pp.connect("producer_a.value", "combiner.input_a")
|
||||
pp.connect("producer_b.value", "combiner.input_b")
|
||||
|
||||
result = pp.run(
|
||||
{"producer_a": {"text": "hello"}, "producer_b": {"text": "world"}, "empty_processor": {"sources": []}},
|
||||
include_outputs_from={"producer_a", "empty_processor", "combiner"},
|
||||
)
|
||||
|
||||
# Producer A should appear in results because it's in include_outputs_from
|
||||
assert "producer_a" in result
|
||||
assert result["producer_a"] == {"value": "A: hello"}
|
||||
# Producer B should NOT appear since it's not in include_outputs_from
|
||||
assert "producer_b" not in result
|
||||
# Combiner should appear in results
|
||||
assert "combiner" in result
|
||||
assert result["combiner"] == {"combined": "A: hello | B: world"}
|
||||
# Empty processor should appear in results even though it returns an empty dict
|
||||
# because it's in include_outputs_from
|
||||
assert "empty_processor" in result
|
||||
assert result["empty_processor"] == {}
|
||||
|
||||
def test__run_component_warns_on_extra_output_keys(self, caplog):
|
||||
"""Test that a warning is raised when a component returns undeclared output keys."""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class ExtraKeyComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"output": value, "extra_key": "unexpected"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("extra", ExtraKeyComponent())
|
||||
|
||||
pp._run_component(
|
||||
component_name="extra",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("extra", 0),
|
||||
inputs={"value": "test"},
|
||||
component_visits={"extra": 0},
|
||||
)
|
||||
assert "returned output keys" in caplog.text
|
||||
assert "extra_key" in caplog.text
|
||||
assert "not declared" in caplog.text
|
||||
|
||||
def test__run_component_no_warning_on_correct_output_keys(self, caplog):
|
||||
"""Test that no warning is raised when a component returns the correct output keys."""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class CorrectComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"output": value}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("correct", CorrectComponent())
|
||||
|
||||
pp._run_component(
|
||||
component_name="correct",
|
||||
component=pp._get_component_with_graph_metadata_and_visits("correct", 0),
|
||||
inputs={"value": "test"},
|
||||
component_visits={"correct": 0},
|
||||
)
|
||||
assert "returned output keys" not in caplog.text
|
||||
assert "did not produce output keys" not in caplog.text
|
||||
|
||||
def test_pipeline_is_possibly_blocked_warning_message(self, caplog):
|
||||
"""
|
||||
Test that the pipeline raises a warning when it is possibly blocked due to missing inputs.
|
||||
|
||||
The situation below looks a little contrived, but it has happened in practice that users create pipelines
|
||||
and accidentally made a mistake in their component code.
|
||||
"""
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
@component
|
||||
class MisconfiguredComponent:
|
||||
# Here we purposely declare other_output which is not actually returned by the run() method
|
||||
@component.output_types(output=str, other_output=str)
|
||||
def run(self, required_input: str) -> dict[str, str]:
|
||||
return {"output": "test"}
|
||||
|
||||
@component
|
||||
class SimpleComponentTwoInputs:
|
||||
@component.output_types(output=str)
|
||||
def run(self, required_input: str, second_required_input: str) -> dict[str, str]:
|
||||
return {"output": "test"}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", MisconfiguredComponent())
|
||||
pp.add_component("second", SimpleComponentTwoInputs())
|
||||
|
||||
# NOTE: We connect both outputs from the first component to the second component, but the first component
|
||||
# doesn't actually produce other_output, so the second component will be blocked due to missing input.
|
||||
pp.connect("first.output", "second.required_input")
|
||||
pp.connect("first.other_output", "second.second_required_input")
|
||||
|
||||
pp.run({"first": {"required_input": "test"}})
|
||||
assert "Cannot run pipeline - the pipeline appears to be blocked." in caplog.text
|
||||
assert " - 'second' (SimpleComponentTwoInputs)" in caplog.text
|
||||
|
||||
def test_pipeline_ensure_inputs_are_deep_copied(self):
|
||||
"""
|
||||
Test to ensure that pipeline deep copies the inputs before passing them to components.
|
||||
|
||||
This is important to prevent unintended side effects when components modify their inputs especially when
|
||||
the output from one component is passed to multiple other components.
|
||||
|
||||
Some other notes about how this situation can arise in practice:
|
||||
- When a component returns a mutable object (like a Document) and that output is passed to multiple other
|
||||
components.
|
||||
- This doesn't happen when using output types like strings or integers, because they are not shared by
|
||||
reference so we will only commonly see this for objects like our dataclasses.
|
||||
"""
|
||||
|
||||
@component
|
||||
class SimpleComponent:
|
||||
@component.output_types(output=Document)
|
||||
def run(self, document: Document) -> dict[str, Document]:
|
||||
# Creates a new document to avoid modifying in place
|
||||
new_document = Document(content=document.content)
|
||||
return {"output": new_document}
|
||||
|
||||
@component
|
||||
class ModifyingComponent:
|
||||
@component.output_types(output=Document)
|
||||
def run(self, document: Document) -> dict[str, Document]:
|
||||
# Modifies the incoming document inplace
|
||||
document.content = "modified"
|
||||
return {"output": document}
|
||||
|
||||
pp = Pipeline()
|
||||
pp.add_component("first", SimpleComponent())
|
||||
pp.add_component("modifier", ModifyingComponent())
|
||||
# It's important that the following component has a name lower down the alphabetical order than "modifier",
|
||||
# since the pipeline runs components in a first-in-first-out manner based on ordered_component_names which is
|
||||
# sorted alphabetically.
|
||||
pp.add_component("second", SimpleComponent())
|
||||
|
||||
pp.connect("first.output", "modifier.document")
|
||||
pp.connect("first.output", "second.document")
|
||||
|
||||
result = pp.run({"first": {"document": Document(content="original")}})
|
||||
|
||||
assert result["modifier"]["output"].content == "modified"
|
||||
# Without deep copying the inputs, the second component would also see the modified document and produce
|
||||
# "modified" instead of "original"
|
||||
assert result["second"]["output"].content == "original"
|
||||
|
||||
def test_pipeline_does_not_corrupt_outputs(self):
|
||||
"""
|
||||
Test that a component's output collected via include_outputs_from is not corrupted when a downstream
|
||||
component receives and mutates the same data in-place.
|
||||
"""
|
||||
|
||||
@component
|
||||
class Producer:
|
||||
@component.output_types(doc=Document)
|
||||
def run(self) -> dict:
|
||||
return {"doc": Document(content="original")}
|
||||
|
||||
@component
|
||||
class Mutator:
|
||||
@component.output_types(doc=Document)
|
||||
def run(self, doc: Document) -> dict:
|
||||
# Modifies the incoming document inplace
|
||||
doc.content = "mutated"
|
||||
return {"doc": doc}
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("producer", Producer())
|
||||
pipe.add_component("mutator", Mutator())
|
||||
pipe.connect("producer.doc", "mutator.doc")
|
||||
|
||||
result = pipe.run({}, include_outputs_from={"producer"})
|
||||
|
||||
assert result["producer"]["doc"].content == "original"
|
||||
assert result["mutator"]["doc"].content == "mutated"
|
||||
|
||||
def test_auto_variadic_connection_to_agent(self):
|
||||
@component
|
||||
class MessageProducer:
|
||||
@component.output_types(messages=list[ChatMessage])
|
||||
def run(self) -> dict[str, list[ChatMessage]]:
|
||||
return {"messages": [ChatMessage.from_user("Hello, world!")]}
|
||||
|
||||
p = Pipeline()
|
||||
p.add_component("message_producer", MessageProducer())
|
||||
p.add_component("message_producer2", MessageProducer())
|
||||
p.add_component("agent", Agent(chat_generator=MockChatGenerator()))
|
||||
p.connect("message_producer", "agent.messages")
|
||||
p.connect("message_producer2", "agent.messages")
|
||||
|
||||
result = p.run({})
|
||||
assert result["agent"]["messages"] == [
|
||||
ChatMessage.from_user("Hello, world!"),
|
||||
ChatMessage.from_user("Hello, world!"),
|
||||
ChatMessage.from_assistant("Hello, world!"),
|
||||
]
|
||||
|
||||
def test_run_auto_variadic_str_to_list_str(self):
|
||||
"""Two str producers connected to a list[str] input are auto-joined and flattened at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("producer1", StringProducer("hello"))
|
||||
p.add_component("producer2", StringProducer("world"))
|
||||
p.add_component("receiver", ListStrAcceptor())
|
||||
p.connect("producer1.text", "receiver.texts")
|
||||
p.connect("producer2.text", "receiver.texts")
|
||||
result = p.run({})
|
||||
assert result["receiver"]["result"] == ["hello", "world"]
|
||||
|
||||
def test_run_auto_variadic_str_and_list_str_to_list_str(self):
|
||||
"""A str producer and a list[str] producer connected to a list[str] input are auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("str_producer", StringProducer("hello"))
|
||||
p.add_component("list_producer", ListStrProducer(["world", "!"]))
|
||||
p.add_component("receiver", ListStrAcceptor())
|
||||
p.connect("str_producer.text", "receiver.texts")
|
||||
p.connect("list_producer.texts", "receiver.texts")
|
||||
result = p.run({})
|
||||
assert result["receiver"]["result"] == ["world", "!", "hello"]
|
||||
|
||||
def test_run_auto_variadic_chat_message_to_list_str(self):
|
||||
"""Two ChatMessage producers connected to a list[str] input are converted and auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("producer1", ChatMessageProducer("hello"))
|
||||
p.add_component("producer2", ChatMessageProducer("world"))
|
||||
p.add_component("receiver", ListStrAcceptor())
|
||||
p.connect("producer1.message", "receiver.texts")
|
||||
p.connect("producer2.message", "receiver.texts")
|
||||
result = p.run({})
|
||||
assert result["receiver"]["result"] == ["hello", "world"]
|
||||
|
||||
def test_run_auto_variadic_str_and_chat_message_to_list_str(self):
|
||||
"""A str producer and a ChatMessage producer connected to a list[str] input are auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("str_producer", StringProducer("hello"))
|
||||
p.add_component("chat_producer", ChatMessageProducer("world"))
|
||||
p.add_component("receiver", ListStrAcceptor())
|
||||
p.connect("str_producer.text", "receiver.texts")
|
||||
p.connect("chat_producer.message", "receiver.texts")
|
||||
result = p.run({})
|
||||
assert result["receiver"]["result"] == ["world", "hello"]
|
||||
|
||||
def test_run_auto_variadic_chat_message_to_list_chat_message(self):
|
||||
"""Two ChatMessage producers connected to a list[ChatMessage] input are auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("producer1", ChatMessageProducer("hello"))
|
||||
p.add_component("producer2", ChatMessageProducer("world"))
|
||||
p.add_component("receiver", ListChatMessageAcceptor())
|
||||
p.connect("producer1.message", "receiver.messages")
|
||||
p.connect("producer2.message", "receiver.messages")
|
||||
result = p.run({})
|
||||
assert [m.text for m in result["receiver"]["result"]] == ["hello", "world"]
|
||||
|
||||
def test_run_auto_variadic_str_to_list_chat_message(self):
|
||||
"""Two str producers connected to a list[ChatMessage] input are converted and auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("producer1", StringProducer("hello"))
|
||||
p.add_component("producer2", StringProducer("world"))
|
||||
p.add_component("receiver", ListChatMessageAcceptor())
|
||||
p.connect("producer1.text", "receiver.messages")
|
||||
p.connect("producer2.text", "receiver.messages")
|
||||
result = p.run({})
|
||||
assert [m.text for m in result["receiver"]["result"]] == ["hello", "world"]
|
||||
|
||||
def test_run_auto_variadic_str_and_chat_message_to_list_chat_message(self):
|
||||
"""A str and a ChatMessage producer connected to a list[ChatMessage] input are auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("str_producer", StringProducer("hello"))
|
||||
p.add_component("chat_producer", ChatMessageProducer("world"))
|
||||
p.add_component("receiver", ListChatMessageAcceptor())
|
||||
p.connect("str_producer.text", "receiver.messages")
|
||||
p.connect("chat_producer.message", "receiver.messages")
|
||||
result = p.run({})
|
||||
assert [m.text for m in result["receiver"]["result"]] == ["world", "hello"]
|
||||
|
||||
def test_run_auto_variadic_chat_message_and_list_chat_message_to_list_chat_message(self):
|
||||
"""A ChatMessage and a list[ChatMessage] producer connected to list[ChatMessage] are auto-joined at runtime."""
|
||||
p = Pipeline()
|
||||
p.add_component("chat_producer", ChatMessageProducer("hello"))
|
||||
p.add_component("list_producer", ListChatMessageProducer(["world", "!"]))
|
||||
p.add_component("receiver", ListChatMessageAcceptor())
|
||||
p.connect("chat_producer.message", "receiver.messages")
|
||||
p.connect("list_producer.messages", "receiver.messages")
|
||||
result = p.run({})
|
||||
assert [m.text for m in result["receiver"]["result"]] == ["hello", "world", "!"]
|
||||
|
||||
def test_connect_rejects_list_of_documents_to_single_document(self):
|
||||
@component
|
||||
class DocsProducer:
|
||||
@component.output_types(docs=list[Document])
|
||||
def run(self) -> dict[str, list[Document]]:
|
||||
return {"docs": [Document(content="a"), Document(content="b"), Document(content="c")]}
|
||||
|
||||
@component
|
||||
class DocConsumer:
|
||||
@component.output_types(out=Document)
|
||||
def run(self, doc: Document) -> dict[str, Document]:
|
||||
return {"out": doc}
|
||||
|
||||
p = Pipeline()
|
||||
p.add_component("producer", DocsProducer())
|
||||
p.add_component("consumer", DocConsumer())
|
||||
with pytest.raises(PipelineConnectError):
|
||||
p.connect("producer.docs", "consumer.doc")
|
||||
|
||||
def test_run_raises_when_multi_element_list_is_unwrapped_at_runtime(self):
|
||||
@component
|
||||
class MultiStrProducer:
|
||||
@component.output_types(texts=list[str])
|
||||
def run(self) -> dict[str, list[str]]:
|
||||
return {"texts": ["first", "second", "third"]}
|
||||
|
||||
@component
|
||||
class SingleStrConsumer:
|
||||
@component.output_types(out=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"out": text}
|
||||
|
||||
p = Pipeline()
|
||||
p.add_component("producer", MultiStrProducer())
|
||||
p.add_component("consumer", SingleStrConsumer())
|
||||
p.connect("producer.texts", "consumer.text")
|
||||
|
||||
with pytest.raises(PipelineRuntimeError, match="Cannot unwrap a list of 3 items"):
|
||||
p.run({})
|
||||
|
||||
def test_run_single_element_list_unwrap_still_works(self):
|
||||
@component
|
||||
class SingleStrProducer:
|
||||
@component.output_types(texts=list[str])
|
||||
def run(self) -> dict[str, list[str]]:
|
||||
return {"texts": ["only-one"]}
|
||||
|
||||
@component
|
||||
class SingleStrConsumer:
|
||||
@component.output_types(out=str)
|
||||
def run(self, text: str) -> dict[str, str]:
|
||||
return {"out": text}
|
||||
|
||||
p = Pipeline()
|
||||
p.add_component("producer", SingleStrProducer())
|
||||
p.add_component("consumer", SingleStrConsumer())
|
||||
p.connect("producer.texts", "consumer.text")
|
||||
assert p.run({}) == {"consumer": {"out": "only-one"}}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,107 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from haystack import Document, Pipeline
|
||||
from haystack.components.builders import ChatPromptBuilder
|
||||
from haystack.components.joiners import DocumentJoiner
|
||||
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
||||
from haystack.core.component import component
|
||||
from haystack.core.errors import PipelineRuntimeError
|
||||
from haystack.dataclasses import ChatMessage
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
|
||||
|
||||
def setup_document_store():
|
||||
documents = [Document(content="My name is Jean and I live in Paris.", embedding=[0.1, 0.3, 0.6])]
|
||||
document_store = InMemoryDocumentStore()
|
||||
document_store.write_documents(documents)
|
||||
return document_store
|
||||
|
||||
|
||||
@component
|
||||
class InvalidOutputEmbeddingRetriever:
|
||||
@component.output_types(documents=list[Document])
|
||||
def run(self, query_embedding: list[float]) -> dict[str, list[Document]]:
|
||||
# Return an int instead of the expected dict with 'documents' key
|
||||
# This will cause the pipeline to crash when trying to pass it to the next component
|
||||
return 42 # type: ignore[return-value]
|
||||
|
||||
|
||||
@component
|
||||
class MockTextEmbedder:
|
||||
@component.output_types(embedding=list[float])
|
||||
def run(self, text: str) -> dict[str, list[float]]:
|
||||
embedding = np.ones(384).tolist() # Mock embedding of size 384
|
||||
return {"embedding": embedding}
|
||||
|
||||
|
||||
class TestPipelineOutputsRaisedInException:
|
||||
def test_hybrid_rag_pipeline_crash_on_embedding_retriever(self):
|
||||
document_store = setup_document_store()
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("text_embedder", MockTextEmbedder())
|
||||
pipeline.add_component("embedding_retriever", InvalidOutputEmbeddingRetriever())
|
||||
pipeline.add_component("bm25_retriever", InMemoryBM25Retriever(document_store))
|
||||
pipeline.add_component("document_joiner", DocumentJoiner(join_mode="concatenate"))
|
||||
pipeline.add_component(
|
||||
"prompt_builder",
|
||||
ChatPromptBuilder(
|
||||
template=[ChatMessage.from_user("Context:\n{{ documents[0].content }}\n\nQuestion: {{question}}\n")],
|
||||
required_variables=["question", "documents"],
|
||||
),
|
||||
)
|
||||
|
||||
pipeline.connect("bm25_retriever", "document_joiner")
|
||||
pipeline.connect("text_embedder", "embedding_retriever")
|
||||
pipeline.connect("embedding_retriever", "document_joiner")
|
||||
pipeline.connect("document_joiner.documents", "prompt_builder.documents")
|
||||
|
||||
question = "Where does Mark live?"
|
||||
test_data = {
|
||||
"text_embedder": {"text": question},
|
||||
"bm25_retriever": {"query": question},
|
||||
"prompt_builder": {"question": question},
|
||||
}
|
||||
|
||||
# run pipeline and expect it to crash due to invalid output type
|
||||
with pytest.raises(PipelineRuntimeError) as exc_info:
|
||||
pipeline.run(
|
||||
data=test_data,
|
||||
include_outputs_from={
|
||||
"text_embedder",
|
||||
"embedding_retriever",
|
||||
"bm25_retriever",
|
||||
"document_joiner",
|
||||
"prompt_builder",
|
||||
},
|
||||
)
|
||||
|
||||
assert "Component name: 'embedding_retriever'" in str(exc_info.value)
|
||||
assert exc_info.value.component_name == "embedding_retriever"
|
||||
assert exc_info.value.component_type == InvalidOutputEmbeddingRetriever
|
||||
# File saving is disabled by default, so pipeline_snapshot_file_path is None
|
||||
assert exc_info.value.pipeline_snapshot_file_path is None
|
||||
|
||||
pipeline_snapshot = exc_info.value.pipeline_snapshot
|
||||
assert pipeline_snapshot is not None
|
||||
pipeline_outputs = pipeline_snapshot.pipeline_state.pipeline_outputs
|
||||
assert pipeline_outputs is not None, "Pipeline outputs should be captured in the exception"
|
||||
|
||||
# verify that bm25_retriever and text_embedder ran successfully before the crash
|
||||
assert "bm25_retriever" in pipeline_outputs["serialized_data"], "BM25 retriever output not captured"
|
||||
assert "documents" in (pipeline_outputs["serialized_data"]["bm25_retriever"]), (
|
||||
"BM25 retriever should have produced documents"
|
||||
)
|
||||
assert "text_embedder" in (pipeline_outputs["serialized_data"]), "Text embedder output not captured"
|
||||
assert "embedding" in (pipeline_outputs["serialized_data"]["text_embedder"]), (
|
||||
"Text embedder should have produced embeddings"
|
||||
)
|
||||
|
||||
# components after the crash point are not in the outputs
|
||||
assert "document_joiner" not in pipeline_outputs, "Document joiner should not have run due to crash"
|
||||
assert "prompt_builder" not in pipeline_outputs, "Prompt builder should not have run due to crash"
|
||||
@@ -0,0 +1,209 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import Pipeline, component
|
||||
|
||||
|
||||
@component
|
||||
class LifecycleRecorder:
|
||||
"""Records every lifecycle method called, so tests can assert which ones the pipeline picks."""
|
||||
|
||||
def __init__(self):
|
||||
self.events = []
|
||||
|
||||
def warm_up(self):
|
||||
self.events.append("warm_up")
|
||||
|
||||
async def warm_up_async(self):
|
||||
self.events.append("warm_up_async")
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self):
|
||||
self.events.append("run")
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self):
|
||||
self.events.append("run_async")
|
||||
return {"value": 1}
|
||||
|
||||
def close(self):
|
||||
self.events.append("close")
|
||||
|
||||
async def close_async(self):
|
||||
self.events.append("close_async")
|
||||
|
||||
|
||||
@component
|
||||
class SyncOnlyRecorder:
|
||||
"""Implements only the synchronous warm_up and close, to exercise the async fallbacks."""
|
||||
|
||||
def __init__(self):
|
||||
self.events = []
|
||||
|
||||
def warm_up(self):
|
||||
self.events.append("warm_up")
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self):
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self):
|
||||
return {"value": 1}
|
||||
|
||||
def close(self):
|
||||
self.events.append("close")
|
||||
|
||||
|
||||
@component
|
||||
class BareComponent:
|
||||
"""Implements no lifecycle method at all."""
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self):
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self):
|
||||
return {"value": 1}
|
||||
|
||||
|
||||
class LoopBoundAsyncClient:
|
||||
"""Mimics a real async client (aiohttp): binds to the loop it is created on and refuses any other."""
|
||||
|
||||
def __init__(self):
|
||||
self._loop = asyncio.get_running_loop()
|
||||
|
||||
async def use(self):
|
||||
if asyncio.get_running_loop() is not self._loop:
|
||||
raise RuntimeError("async client used on a different event loop than the one it was created on")
|
||||
|
||||
|
||||
@component
|
||||
class AsyncClientComponent:
|
||||
"""Creates a loop-bound async client in warm_up_async and uses it in run_async."""
|
||||
|
||||
def __init__(self):
|
||||
self.client: LoopBoundAsyncClient | None = None
|
||||
|
||||
async def warm_up_async(self):
|
||||
if self.client is None:
|
||||
self.client = LoopBoundAsyncClient()
|
||||
|
||||
@component.output_types(value=int)
|
||||
async def run_async(self):
|
||||
assert self.client is not None
|
||||
await self.client.use()
|
||||
return {"value": 1}
|
||||
|
||||
@component.output_types(value=int)
|
||||
def run(self):
|
||||
return {"value": 1}
|
||||
|
||||
|
||||
async def test_run_async_uses_warm_up_async():
|
||||
"""When a component implements warm_up_async, run_async uses it and does not also call its sync warm_up."""
|
||||
rec = LifecycleRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
await pipe.run_async({"rec": {}})
|
||||
assert "warm_up_async" in rec.events
|
||||
assert "warm_up" not in rec.events
|
||||
|
||||
|
||||
async def test_warm_up_async_falls_back_to_sync_warm_up():
|
||||
"""A component with only the sync warm_up is still warmed by run_async through that method."""
|
||||
rec = SyncOnlyRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
await pipe.run_async({"rec": {}})
|
||||
assert rec.events == ["warm_up"]
|
||||
|
||||
|
||||
def test_sync_run_uses_sync_warm_up():
|
||||
"""The sync run path warms components via the sync warm_up, never warm_up_async."""
|
||||
rec = LifecycleRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
pipe.run({"rec": {}})
|
||||
assert "warm_up" in rec.events
|
||||
assert "warm_up_async" not in rec.events
|
||||
|
||||
|
||||
def test_pipeline_close_calls_sync_close_only():
|
||||
"""Pipeline.close() calls each component's sync close, never close_async."""
|
||||
rec = LifecycleRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
pipe.close()
|
||||
assert "close" in rec.events
|
||||
assert "close_async" not in rec.events
|
||||
|
||||
|
||||
async def test_pipeline_close_async_calls_async_close_only():
|
||||
"""When a component implements close_async, Pipeline.close_async() uses it and does not also call its sync close."""
|
||||
rec = LifecycleRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
await pipe.close_async()
|
||||
assert "close_async" in rec.events
|
||||
assert "close" not in rec.events
|
||||
|
||||
|
||||
async def test_close_async_falls_back_to_sync_close():
|
||||
"""A component with only the sync close is still released by close_async through that method."""
|
||||
rec = SyncOnlyRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
await pipe.close_async()
|
||||
assert rec.events == ["close"]
|
||||
|
||||
|
||||
async def test_run_does_not_auto_close():
|
||||
"""Running a pipeline (sync or async) never closes components; closing is always explicit."""
|
||||
rec = LifecycleRecorder()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("rec", rec)
|
||||
pipe.run({"rec": {}})
|
||||
await pipe.run_async({"rec": {}})
|
||||
assert "close" not in rec.events
|
||||
assert "close_async" not in rec.events
|
||||
|
||||
|
||||
async def test_lifecycle_methods_are_optional():
|
||||
"""A component without lifecycle methods works: every call is hasattr-guarded and skipped."""
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("bare", BareComponent())
|
||||
await pipe.warm_up_async()
|
||||
pipe.close()
|
||||
await pipe.close_async()
|
||||
await pipe.run_async({"bare": {}})
|
||||
|
||||
|
||||
def test_loop_bound_client_rejects_other_loop():
|
||||
"""The fake client raises when used from a different loop.
|
||||
|
||||
This ensures the affinity test below enforces loop binding.
|
||||
"""
|
||||
|
||||
async def _make_loop_bound_client():
|
||||
return LoopBoundAsyncClient()
|
||||
|
||||
client = asyncio.run(_make_loop_bound_client())
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(client.use())
|
||||
|
||||
|
||||
async def test_async_client_bound_to_run_loop():
|
||||
"""warm_up_async creates the async client on the loop run_async uses, so it stays usable there."""
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("client_component", AsyncClientComponent())
|
||||
await pipe.warm_up_async()
|
||||
# Would raise if warm_up_async had bound the client to a different loop than run_async
|
||||
await pipe.run_async({"client_component": {}})
|
||||
@@ -0,0 +1,188 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import asyncio
|
||||
from dataclasses import replace
|
||||
from unittest.mock import ANY
|
||||
|
||||
import pytest
|
||||
from _pytest.monkeypatch import MonkeyPatch
|
||||
|
||||
from haystack import Pipeline, component
|
||||
from haystack.dataclasses import Document
|
||||
from haystack.tracing.tracer import tracer
|
||||
from test.tracing.utils import SpyingSpan, SpyingTracer
|
||||
|
||||
|
||||
@component
|
||||
class Hello:
|
||||
@component.output_types(output=str)
|
||||
def run(self, word: str | None) -> dict[str, str]: # use optional to spice up the typing tags
|
||||
"""
|
||||
Takes a string in input and returns "Hello, <string>!" in output.
|
||||
"""
|
||||
return {"output": f"Hello, {word}!"}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def pipeline() -> Pipeline:
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("hello", Hello())
|
||||
pipeline.add_component("hello2", Hello())
|
||||
pipeline.connect("hello.output", "hello2.word")
|
||||
return pipeline
|
||||
|
||||
|
||||
class TestTracing:
|
||||
def test_with_enabled_tracing(self, pipeline: Pipeline, spying_tracer: SpyingTracer) -> None:
|
||||
pipeline.run(data={"word": "world"})
|
||||
|
||||
assert len(spying_tracer.spans) == 3
|
||||
|
||||
assert spying_tracer.spans == [
|
||||
pipeline_span := SpyingSpan(
|
||||
operation_name="haystack.pipeline.run",
|
||||
tags={
|
||||
"haystack.pipeline.input_data": {"hello": {"word": "world"}},
|
||||
"haystack.pipeline.output_data": {"hello2": {"output": "Hello, Hello, world!!"}},
|
||||
"haystack.pipeline.metadata": {},
|
||||
"haystack.pipeline.max_runs_per_component": 100,
|
||||
"haystack.pipeline.execution_mode": "sync",
|
||||
},
|
||||
parent_span=None,
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
SpyingSpan(
|
||||
operation_name="haystack.component.run",
|
||||
tags={
|
||||
"haystack.component.name": "hello",
|
||||
"haystack.component.type": "Hello",
|
||||
"haystack.component.fully_qualified_type": "test.core.pipeline.test_tracing.Hello",
|
||||
"haystack.component.input_types": {"word": "str"},
|
||||
"haystack.component.input_spec": {"word": {"type": ANY, "senders": []}},
|
||||
"haystack.component.input": {"word": "world"},
|
||||
"haystack.component.output_spec": {"output": {"type": "str", "receivers": ["hello2"]}},
|
||||
"haystack.component.output": {"output": "Hello, world!"},
|
||||
"haystack.component.visits": 1,
|
||||
},
|
||||
parent_span=pipeline_span,
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
SpyingSpan(
|
||||
operation_name="haystack.component.run",
|
||||
tags={
|
||||
"haystack.component.name": "hello2",
|
||||
"haystack.component.type": "Hello",
|
||||
"haystack.component.fully_qualified_type": "test.core.pipeline.test_tracing.Hello",
|
||||
"haystack.component.input_types": {"word": "str"},
|
||||
"haystack.component.input_spec": {"word": {"type": ANY, "senders": ["hello"]}},
|
||||
"haystack.component.input": {"word": "Hello, world!"},
|
||||
"haystack.component.output_spec": {"output": {"type": "str", "receivers": []}},
|
||||
"haystack.component.output": {"output": "Hello, Hello, world!!"},
|
||||
"haystack.component.visits": 1,
|
||||
},
|
||||
parent_span=pipeline_span,
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
]
|
||||
|
||||
# We need to check the type of the input_spec because it can be rendered differently
|
||||
# depending on the Python version 🫠
|
||||
assert spying_tracer.spans[1].tags["haystack.component.input_spec"]["word"]["type"] in [
|
||||
"typing.Union[str, NoneType]",
|
||||
"typing.Optional[str]",
|
||||
"str | None",
|
||||
]
|
||||
|
||||
def test_with_enabled_content_tracing(
|
||||
self, spying_tracer: SpyingTracer, monkeypatch: MonkeyPatch, pipeline: Pipeline
|
||||
) -> None:
|
||||
# Monkeypatch to avoid impact on other tests
|
||||
monkeypatch.setattr(tracer, "is_content_tracing_enabled", True)
|
||||
|
||||
pipeline.run(data={"word": "world"})
|
||||
|
||||
assert len(spying_tracer.spans) == 3
|
||||
assert spying_tracer.spans == [
|
||||
pipeline_span := SpyingSpan(
|
||||
operation_name="haystack.pipeline.run",
|
||||
tags={
|
||||
"haystack.pipeline.metadata": {},
|
||||
"haystack.pipeline.max_runs_per_component": 100,
|
||||
"haystack.pipeline.input_data": {"hello": {"word": "world"}},
|
||||
"haystack.pipeline.output_data": {"hello2": {"output": "Hello, Hello, world!!"}},
|
||||
"haystack.pipeline.execution_mode": "sync",
|
||||
},
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
SpyingSpan(
|
||||
operation_name="haystack.component.run",
|
||||
tags={
|
||||
"haystack.component.name": "hello",
|
||||
"haystack.component.type": "Hello",
|
||||
"haystack.component.fully_qualified_type": "test.core.pipeline.test_tracing.Hello",
|
||||
"haystack.component.input_types": {"word": "str"},
|
||||
"haystack.component.input_spec": {"word": {"type": "str | None", "senders": []}},
|
||||
"haystack.component.output_spec": {"output": {"type": "str", "receivers": ["hello2"]}},
|
||||
"haystack.component.input": {"word": "world"},
|
||||
"haystack.component.visits": 1,
|
||||
"haystack.component.output": {"output": "Hello, world!"},
|
||||
},
|
||||
parent_span=pipeline_span,
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
SpyingSpan(
|
||||
operation_name="haystack.component.run",
|
||||
tags={
|
||||
"haystack.component.name": "hello2",
|
||||
"haystack.component.type": "Hello",
|
||||
"haystack.component.fully_qualified_type": "test.core.pipeline.test_tracing.Hello",
|
||||
"haystack.component.input_types": {"word": "str"},
|
||||
"haystack.component.input_spec": {"word": {"type": "str | None", "senders": ["hello"]}},
|
||||
"haystack.component.output_spec": {"output": {"type": "str", "receivers": []}},
|
||||
"haystack.component.input": {"word": "Hello, world!"},
|
||||
"haystack.component.visits": 1,
|
||||
"haystack.component.output": {"output": "Hello, Hello, world!!"},
|
||||
},
|
||||
parent_span=pipeline_span,
|
||||
trace_id=ANY,
|
||||
span_id=ANY,
|
||||
),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("run_async", [False, True])
|
||||
def test_span_input_not_affected_by_component_mutation(self, run_async, spying_tracer, monkeypatch):
|
||||
"""
|
||||
Verify that the haystack.component.input span tag retains the pre-execution value.
|
||||
|
||||
Parametrized to cover both the synchronous (`run`) and asynchronous (`run_async`) execution paths.
|
||||
"""
|
||||
monkeypatch.setattr(tracer, "is_content_tracing_enabled", True)
|
||||
|
||||
@component
|
||||
class MutatingComponent:
|
||||
@component.output_types(doc=Document)
|
||||
def run(self, doc: Document) -> dict:
|
||||
return {"doc": replace(doc, content="mutated")}
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("mutator", MutatingComponent())
|
||||
|
||||
if run_async:
|
||||
result = asyncio.run(pipe.run_async({"mutator": {"doc": Document(content="original")}}))
|
||||
else:
|
||||
result = pipe.run({"mutator": {"doc": Document(content="original")}})
|
||||
|
||||
pipeline_span = spying_tracer.spans[0]
|
||||
component_span = spying_tracer.spans[1]
|
||||
|
||||
assert pipeline_span.operation_name == "haystack.pipeline.run"
|
||||
assert pipeline_span.tags["haystack.pipeline.execution_mode"] == ("async" if run_async else "sync")
|
||||
assert component_span.tags["haystack.component.input"]["doc"].content == "original"
|
||||
assert result["mutator"]["doc"].content == "mutated"
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component import component
|
||||
from haystack.core.pipeline import Pipeline
|
||||
|
||||
|
||||
@component
|
||||
class NonPep604Producer:
|
||||
@component.output_types(value=Optional[str], items=Union[str, list[str]])
|
||||
def run(self, text: str) -> dict[str, Union[Optional[str], Union[str, list[str]]]]:
|
||||
return {"value": text, "items": ["a", "b", "c"]}
|
||||
|
||||
|
||||
@component
|
||||
class NonPep604Consumer:
|
||||
@component.output_types(result=str)
|
||||
def run(self, value: Optional[str], items: Union[str, list[str]]) -> dict[str, str]:
|
||||
return {"result": f"{value}: {items}"}
|
||||
|
||||
|
||||
@component
|
||||
class Pep604Producer:
|
||||
@component.output_types(value=str | None, items=str | list[str])
|
||||
def run(self, text: str) -> dict[str, str | None | str | list[str]]:
|
||||
return {"value": text, "items": ["a", "b", "c"]}
|
||||
|
||||
|
||||
@component
|
||||
class Pep604Consumer:
|
||||
@component.output_types(result=str)
|
||||
def run(self, value: str | None, items: str | list[str]) -> dict[str, str]:
|
||||
return {"result": f"{value}: {items}"}
|
||||
|
||||
|
||||
class TestTypeSyntaxCompatibility:
|
||||
"""Tests for type syntax compatibility between non-PEP 604 and PEP 604."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"producer, consumer", [(NonPep604Producer(), Pep604Consumer()), (Pep604Producer(), NonPep604Consumer())]
|
||||
)
|
||||
def test_type_syntax_compatibility(self, producer, consumer):
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("producer", producer)
|
||||
pipe.add_component("consumer", consumer)
|
||||
pipe.connect("producer.value", "consumer.value")
|
||||
pipe.connect("producer.items", "consumer.items")
|
||||
result = pipe.run({"producer": {"text": "hello"}})
|
||||
assert result["consumer"]["result"] == "hello: ['a', 'b', 'c']"
|
||||
|
||||
def test_non_pep604_pipeline(self):
|
||||
yaml_content = """
|
||||
components:
|
||||
producer:
|
||||
init_parameters: {}
|
||||
type: test.core.pipeline.test_type_syntax_compatibility.NonPep604Producer
|
||||
consumer:
|
||||
init_parameters: {}
|
||||
type: test.core.pipeline.test_type_syntax_compatibility.NonPep604Consumer
|
||||
connection_type_validation: true
|
||||
connections:
|
||||
- receiver: consumer.value
|
||||
sender: producer.value
|
||||
- receiver: consumer.items
|
||||
sender: producer.items
|
||||
max_runs_per_component: 100
|
||||
metadata: {}
|
||||
"""
|
||||
|
||||
pipe = Pipeline.loads(yaml_content)
|
||||
|
||||
result = pipe.run({"producer": {"text": "test"}})
|
||||
assert result["consumer"]["result"] == "test: ['a', 'b', 'c']"
|
||||
@@ -0,0 +1,304 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.components.builders.prompt_builder import PromptBuilder
|
||||
from haystack.components.generators.chat.openai import OpenAIChatGenerator
|
||||
from haystack.core.pipeline.utils import (
|
||||
FIFOPriorityQueue,
|
||||
_deepcopy_with_exceptions,
|
||||
args_deprecated,
|
||||
parse_connect_string,
|
||||
)
|
||||
from haystack.tools import ComponentTool
|
||||
|
||||
|
||||
def get_weather_report(city: str) -> str:
|
||||
return f"Weather report for {city}: 20°C, sunny"
|
||||
|
||||
|
||||
def test_parse_connection():
|
||||
assert parse_connect_string("foobar") == ("foobar", None)
|
||||
assert parse_connect_string("foo.bar") == ("foo", "bar")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def empty_queue():
|
||||
"""Fixture providing a fresh empty queue for each test."""
|
||||
return FIFOPriorityQueue()
|
||||
|
||||
|
||||
def test_empty_queue_initialization(empty_queue):
|
||||
"""Test that a new queue is empty."""
|
||||
assert len(empty_queue) == 0
|
||||
assert not bool(empty_queue)
|
||||
|
||||
|
||||
def test_push_single_item(empty_queue):
|
||||
"""Test pushing a single item."""
|
||||
empty_queue.push("item1", 1)
|
||||
assert len(empty_queue) == 1
|
||||
assert bool(empty_queue)
|
||||
assert empty_queue.peek() == (1, "item1")
|
||||
|
||||
|
||||
def test_push_multiple_items_different_priorities(empty_queue):
|
||||
"""Test pushing multiple items with different priorities."""
|
||||
items = [("item3", 3), ("item1", 1), ("item2", 2)]
|
||||
for item, priority in items:
|
||||
empty_queue.push(item, priority)
|
||||
|
||||
# Items should come out in priority order
|
||||
assert empty_queue.pop() == (1, "item1")
|
||||
assert empty_queue.pop() == (2, "item2")
|
||||
assert empty_queue.pop() == (3, "item3")
|
||||
|
||||
|
||||
def test_push_multiple_items_same_priority(empty_queue):
|
||||
"""Test FIFO behavior for items with equal priority."""
|
||||
items = [("first", 1), ("second", 1), ("third", 1)]
|
||||
for item, priority in items:
|
||||
empty_queue.push(item, priority)
|
||||
|
||||
# Items should come out in insertion order
|
||||
assert empty_queue.pop() == (1, "first")
|
||||
assert empty_queue.pop() == (1, "second")
|
||||
assert empty_queue.pop() == (1, "third")
|
||||
|
||||
|
||||
def test_mixed_priority_and_fifo(empty_queue):
|
||||
"""Test mixed priority levels with some equal priorities."""
|
||||
empty_queue.push("medium1", 2)
|
||||
empty_queue.push("high", 1)
|
||||
empty_queue.push("medium2", 2)
|
||||
empty_queue.push("low", 3)
|
||||
|
||||
# Check extraction order
|
||||
assert empty_queue.pop() == (1, "high")
|
||||
assert empty_queue.pop() == (2, "medium1")
|
||||
assert empty_queue.pop() == (2, "medium2")
|
||||
assert empty_queue.pop() == (3, "low")
|
||||
|
||||
|
||||
def test_peek_behavior(empty_queue):
|
||||
"""Test that peek returns items without removing them."""
|
||||
empty_queue.push("item1", 1)
|
||||
empty_queue.push("item2", 2)
|
||||
|
||||
# Peek multiple times
|
||||
for _ in range(3):
|
||||
assert empty_queue.peek() == (1, "item1")
|
||||
assert len(empty_queue) == 2
|
||||
|
||||
|
||||
def test_get_behavior(empty_queue):
|
||||
"""Test the get method with both empty and non-empty queues."""
|
||||
# Test on empty queue
|
||||
assert empty_queue.get() is None
|
||||
|
||||
# Test with items
|
||||
empty_queue.push("item1", 1)
|
||||
assert empty_queue.get() == (1, "item1")
|
||||
assert empty_queue.get() is None # Queue should be empty again
|
||||
|
||||
|
||||
def test_pop_empty_queue(empty_queue):
|
||||
"""Test that pop raises IndexError on empty queue."""
|
||||
with pytest.raises(IndexError, match="pop from empty queue"):
|
||||
empty_queue.pop()
|
||||
|
||||
|
||||
def test_peek_empty_queue(empty_queue):
|
||||
"""Test that peek raises IndexError on empty queue."""
|
||||
with pytest.raises(IndexError, match="peek at empty queue"):
|
||||
empty_queue.peek()
|
||||
|
||||
|
||||
def test_length_updates(empty_queue):
|
||||
"""Test that length updates correctly with pushes and pops."""
|
||||
initial_len = len(empty_queue)
|
||||
assert initial_len == 0
|
||||
|
||||
# Test length increases
|
||||
empty_queue.push("item1", 1)
|
||||
assert len(empty_queue) == 1
|
||||
empty_queue.push("item2", 2)
|
||||
assert len(empty_queue) == 2
|
||||
|
||||
# Test length decreases
|
||||
empty_queue.pop()
|
||||
assert len(empty_queue) == 1
|
||||
empty_queue.pop()
|
||||
assert len(empty_queue) == 0
|
||||
|
||||
|
||||
def test_bool_conversion(empty_queue):
|
||||
"""Test boolean conversion in various states."""
|
||||
# Empty queue should be False
|
||||
assert not bool(empty_queue)
|
||||
|
||||
# Queue with items should be True
|
||||
empty_queue.push("item", 1)
|
||||
assert bool(empty_queue)
|
||||
|
||||
# Queue should be False again after removing item
|
||||
empty_queue.pop()
|
||||
assert not bool(empty_queue)
|
||||
|
||||
|
||||
def test_large_number_of_items(empty_queue):
|
||||
"""Test handling of a large number of items with mixed priorities."""
|
||||
# Add 1000 items with 10 different priority levels
|
||||
for i in range(1000):
|
||||
priority = i % 10
|
||||
empty_queue.push(f"item{i}", priority)
|
||||
|
||||
# Verify FIFO order within same priority
|
||||
last_priority = -1
|
||||
last_index = -1
|
||||
for _ in range(1000):
|
||||
priority, item = empty_queue.pop()
|
||||
current_index = int(item[4:]) # Extract index from "itemX"
|
||||
|
||||
if priority == last_priority:
|
||||
assert current_index > last_index, "FIFO order violated within same priority"
|
||||
else:
|
||||
assert priority > last_priority, "Priority order violated"
|
||||
|
||||
last_priority = priority
|
||||
last_index = current_index
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"items",
|
||||
[
|
||||
[(1, "A"), (1, "B"), (1, "C")], # Same priority
|
||||
[(3, "A"), (2, "B"), (1, "C")], # Different priorities
|
||||
[(2, "A"), (1, "B"), (2, "C")], # Mixed priorities
|
||||
],
|
||||
)
|
||||
def test_queue_ordering_parametrized(empty_queue, items):
|
||||
"""Parametrized test for different ordering scenarios."""
|
||||
for priority, item in items:
|
||||
empty_queue.push(item, priority)
|
||||
|
||||
sorted_items = sorted(items, key=lambda x: (x[0], items.index(x)))
|
||||
for priority, item in sorted_items:
|
||||
assert empty_queue.pop() == (priority, item)
|
||||
|
||||
|
||||
class Copyable:
|
||||
def __init__(self, name="copyable"):
|
||||
self.name = name
|
||||
|
||||
|
||||
class NotCopyable:
|
||||
def __init__(self, name="not_copyable"):
|
||||
self.name = name
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
raise TypeError("This object cannot be deepcopied.")
|
||||
|
||||
|
||||
class TestDeepcopyWithFallback:
|
||||
def test_deepcopy_with_fallback_copyable(self, caplog):
|
||||
original = {"class": Copyable()}
|
||||
with caplog.at_level(logging.INFO):
|
||||
copy = _deepcopy_with_exceptions(original)
|
||||
assert "Deepcopy failed for object of type" not in caplog.text
|
||||
assert copy["class"] is not original["class"]
|
||||
|
||||
def test_deepcopy_with_fallback_not_copyable_error(self, caplog):
|
||||
original = {"class": NotCopyable()}
|
||||
with caplog.at_level(logging.INFO):
|
||||
copy = _deepcopy_with_exceptions(original)
|
||||
assert "Deepcopy failed for object of type 'NotCopyable'" in caplog.text
|
||||
assert copy["class"] is original["class"]
|
||||
|
||||
def test_deepcopy_with_fallback_mixed_copyable_list(self, caplog):
|
||||
obj1 = Copyable()
|
||||
obj2 = NotCopyable()
|
||||
original = {"objects": [obj1, obj2]}
|
||||
with caplog.at_level(logging.INFO):
|
||||
copy = _deepcopy_with_exceptions(original)
|
||||
assert "Deepcopy failed for object of type 'NotCopyable'" in caplog.text
|
||||
assert copy["objects"][0] is not original["objects"][0]
|
||||
assert copy["objects"][1] is original["objects"][1]
|
||||
|
||||
def test_deepcopy_with_fallback_mixed_copyable_tuple(self, caplog):
|
||||
obj1 = Copyable()
|
||||
obj2 = NotCopyable()
|
||||
original = {"objects": (obj1, obj2)}
|
||||
with caplog.at_level(logging.INFO):
|
||||
copy = _deepcopy_with_exceptions(original)
|
||||
assert "Deepcopy failed for object of type 'NotCopyable'" in caplog.text
|
||||
assert copy["objects"][0] is not original["objects"][0]
|
||||
assert copy["objects"][1] is original["objects"][1]
|
||||
|
||||
def test_deepcopy_with_fallback_tool(self):
|
||||
tool = ComponentTool(
|
||||
name="problematic_tool", description="This is a problematic tool.", component=PromptBuilder("{{query}}")
|
||||
)
|
||||
original = {"tools": tool}
|
||||
copy = _deepcopy_with_exceptions(original)
|
||||
assert copy["tools"] is original["tools"]
|
||||
|
||||
def test_deepcopy_with_fallback_component(self, monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test")
|
||||
comp = OpenAIChatGenerator()
|
||||
original = {"component": comp}
|
||||
res = _deepcopy_with_exceptions(original)
|
||||
assert res["component"] is original["component"]
|
||||
|
||||
|
||||
class TestArgsDeprecated:
|
||||
@pytest.fixture
|
||||
def sample_function(self):
|
||||
@args_deprecated
|
||||
def sample_func(param1: str = "default1", param2: int = 42) -> str:
|
||||
return f"{param1}-{param2}"
|
||||
|
||||
return sample_func
|
||||
|
||||
def test_warning_with_positional_args(self, sample_function):
|
||||
# using positional arguments only
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
result = sample_function("test", 123)
|
||||
assert result == "test-123"
|
||||
assert len(w) == 1
|
||||
assert issubclass(w[0].category, DeprecationWarning)
|
||||
assert (
|
||||
"Warning: In an upcoming release, this method will require keyword arguments for all parameters"
|
||||
in str(w[0].message)
|
||||
)
|
||||
|
||||
def test_warning_with_mixed_args(self, sample_function):
|
||||
# mixing positional and keyword arguments
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
result = sample_function("test", 123)
|
||||
assert result == "test-123"
|
||||
assert len(w) == 1
|
||||
assert issubclass(w[0].category, DeprecationWarning)
|
||||
assert (
|
||||
"Warning: In an upcoming release, this method will require keyword arguments for all parameters"
|
||||
in str(w[0].message)
|
||||
)
|
||||
|
||||
def test_no_warning_with_default_args(self, sample_function):
|
||||
# using default arguments
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
result = sample_function()
|
||||
assert result == "default1-42"
|
||||
assert len(w) == 0
|
||||
|
||||
def test_no_warning_with_keyword_args(self, sample_function):
|
||||
# using keyword arguments
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
result = sample_function(param1="test", param2=123)
|
||||
assert result == "test-123"
|
||||
assert len(w) == 0
|
||||
@@ -0,0 +1,210 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component.types import InputSocket, OutputSocket, Variadic
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.pipeline.descriptions import find_pipeline_inputs, find_pipeline_outputs
|
||||
from haystack.testing.factory import component_class
|
||||
from haystack.testing.sample_components import AddFixedValue, Double, Parity, Sum
|
||||
|
||||
|
||||
def test_find_pipeline_input_no_input():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
pipe.connect("comp2", "comp1")
|
||||
|
||||
assert find_pipeline_inputs(pipe.graph) == {"comp1": [], "comp2": []}
|
||||
|
||||
|
||||
def test_find_pipeline_input_one_input():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
assert find_pipeline_inputs(pipe.graph) == {"comp1": [InputSocket(name="value", type=int)], "comp2": []}
|
||||
|
||||
|
||||
def test_find_pipeline_input_two_inputs_same_component():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
assert find_pipeline_inputs(pipe.graph) == {
|
||||
"comp1": [InputSocket(name="value", type=int), InputSocket(name="add", type=int | None, default_value=None)],
|
||||
"comp2": [],
|
||||
}
|
||||
|
||||
|
||||
def test_find_pipeline_input_some_inputs_different_components():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.add_component("comp3", AddFixedValue())
|
||||
pipe.connect("comp1.result", "comp3.value")
|
||||
pipe.connect("comp2.value", "comp3.add")
|
||||
|
||||
assert find_pipeline_inputs(pipe.graph) == {
|
||||
"comp1": [InputSocket(name="value", type=int), InputSocket(name="add", type=int | None, default_value=None)],
|
||||
"comp2": [InputSocket(name="value", type=int)],
|
||||
"comp3": [],
|
||||
}
|
||||
|
||||
|
||||
def test_find_pipeline_variable_input_nodes_in_the_pipeline():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.add_component("comp3", Sum())
|
||||
|
||||
assert find_pipeline_inputs(pipe.graph) == {
|
||||
"comp1": [InputSocket(name="value", type=int), InputSocket(name="add", type=int | None, default_value=None)],
|
||||
"comp2": [InputSocket(name="value", type=int)],
|
||||
"comp3": [InputSocket(name="values", type=Variadic[int])],
|
||||
}
|
||||
|
||||
|
||||
def test_find_pipeline_output_no_output():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
pipe.connect("comp2", "comp1")
|
||||
|
||||
assert find_pipeline_outputs(pipe.graph) == {"comp1": [], "comp2": []}
|
||||
|
||||
|
||||
def test_find_pipeline_output_one_output():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
assert find_pipeline_outputs(pipe.graph) == {"comp1": [], "comp2": [OutputSocket(name="value", type=int)]}
|
||||
|
||||
|
||||
def test_find_pipeline_some_outputs_same_component():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Parity())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
assert find_pipeline_outputs(pipe.graph) == {
|
||||
"comp1": [],
|
||||
"comp2": [OutputSocket(name="even", type=int), OutputSocket(name="odd", type=int)],
|
||||
}
|
||||
|
||||
|
||||
def test_find_pipeline_some_outputs_different_components():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Parity())
|
||||
pipe.add_component("comp3", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
pipe.connect("comp1", "comp3")
|
||||
|
||||
assert find_pipeline_outputs(pipe.graph) == {
|
||||
"comp1": [],
|
||||
"comp2": [OutputSocket(name="even", type=int), OutputSocket(name="odd", type=int)],
|
||||
"comp3": [OutputSocket(name="value", type=int)],
|
||||
}
|
||||
|
||||
|
||||
def test_validate_pipeline_input_pipeline_with_no_inputs():
|
||||
pipe = Pipeline()
|
||||
NoInputs = component_class("NoInputs", input_types={}, output={"value": 10})
|
||||
pipe.add_component("no_inputs", NoInputs())
|
||||
res = pipe.run({})
|
||||
assert res == {"no_inputs": {"value": 10}}
|
||||
|
||||
|
||||
def test_validate_pipeline_input_pipeline_with_no_inputs_no_outputs():
|
||||
pipe = Pipeline()
|
||||
NoIO = component_class("NoIO", input_types={}, output={})
|
||||
pipe.add_component("no_inputs", NoIO())
|
||||
res = pipe.run({})
|
||||
assert res == {}
|
||||
|
||||
|
||||
def test_validate_pipeline_input_unknown_component():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"test_component": {"value": 1}})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_all_necessary_input_is_present():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_all_necessary_input_is_present_considering_defaults():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
pipe.run({"comp1": {"value": 1}})
|
||||
pipe.run({"comp1": {"value": 1, "add": 2}})
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"comp1": {"add": 3}})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_only_expected_input_is_present():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"comp1": {"value": 1}, "comp2": {"value": 2}})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_only_expected_input_is_present_falsy():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"comp1": {"value": 1}, "comp2": {"value": 0}})
|
||||
|
||||
|
||||
def test_validate_pipeline_falsy_input_present():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp", Double())
|
||||
assert pipe.run({"comp": {"value": 0}}) == {"comp": {"value": 0}}
|
||||
|
||||
|
||||
def test_validate_pipeline_falsy_input_missing():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp", Double())
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"comp": {}})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_only_expected_input_is_present_including_unknown_names():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", Double())
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
pipe.run({"comp1": {"value": 1, "add": 2}})
|
||||
|
||||
|
||||
def test_validate_pipeline_input_only_expected_input_is_present_and_defaults_dont_interfere():
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=10))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1", "comp2")
|
||||
assert pipe.run({"comp1": {"value": 1, "add": 5}}) == {"comp2": {"value": 12}}
|
||||
@@ -0,0 +1,74 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components.accumulate import Accumulate, _default_function
|
||||
|
||||
|
||||
def my_subtract(first, second):
|
||||
return first - second
|
||||
|
||||
|
||||
def test_to_dict():
|
||||
accumulate = Accumulate()
|
||||
res = accumulate.to_dict()
|
||||
assert res == {
|
||||
"type": "haystack.testing.sample_components.accumulate.Accumulate",
|
||||
"init_parameters": {"function": "haystack.testing.sample_components.accumulate._default_function"},
|
||||
}
|
||||
|
||||
|
||||
def test_to_dict_with_custom_function():
|
||||
accumulate = Accumulate(function=my_subtract)
|
||||
res = accumulate.to_dict()
|
||||
assert res == {
|
||||
"type": "haystack.testing.sample_components.accumulate.Accumulate",
|
||||
"init_parameters": {"function": "test_accumulate.my_subtract"},
|
||||
}
|
||||
|
||||
|
||||
def test_from_dict():
|
||||
data = {"type": "haystack.testing.sample_components.accumulate.Accumulate", "init_parameters": {}}
|
||||
accumulate = Accumulate.from_dict(data)
|
||||
assert accumulate.function == _default_function
|
||||
|
||||
|
||||
def test_from_dict_with_default_function():
|
||||
data = {
|
||||
"type": "haystack.testing.sample_components.accumulate.Accumulate",
|
||||
"init_parameters": {"function": "haystack.testing.sample_components.accumulate._default_function"},
|
||||
}
|
||||
accumulate = Accumulate.from_dict(data)
|
||||
assert accumulate.function == _default_function
|
||||
|
||||
|
||||
def test_from_dict_with_custom_function():
|
||||
data = {
|
||||
"type": "haystack.testing.sample_components.accumulate.Accumulate",
|
||||
"init_parameters": {"function": "test_accumulate.my_subtract"},
|
||||
}
|
||||
accumulate = Accumulate.from_dict(data)
|
||||
assert accumulate.function == my_subtract
|
||||
|
||||
|
||||
def test_accumulate_default():
|
||||
component = Accumulate()
|
||||
results = component.run(value=10)
|
||||
assert results == {"value": 10}
|
||||
assert component.state == 10
|
||||
|
||||
results = component.run(value=1)
|
||||
assert results == {"value": 11}
|
||||
assert component.state == 11
|
||||
|
||||
|
||||
def test_accumulate_callable():
|
||||
component = Accumulate(function=my_subtract)
|
||||
|
||||
results = component.run(value=10)
|
||||
assert results == {"value": -10}
|
||||
assert component.state == -10
|
||||
|
||||
results = component.run(value=1)
|
||||
assert results == {"value": -11}
|
||||
assert component.state == -11
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import AddFixedValue
|
||||
|
||||
|
||||
def test_run():
|
||||
component = AddFixedValue()
|
||||
results = component.run(value=50, add=10)
|
||||
assert results == {"result": 60}
|
||||
@@ -0,0 +1,29 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Concatenate
|
||||
|
||||
|
||||
def test_input_lists():
|
||||
component = Concatenate()
|
||||
res = component.run(first=["This"], second=["That"])
|
||||
assert res == {"value": ["This", "That"]}
|
||||
|
||||
|
||||
def test_input_strings():
|
||||
component = Concatenate()
|
||||
res = component.run(first="This", second="That")
|
||||
assert res == {"value": ["This", "That"]}
|
||||
|
||||
|
||||
def test_input_first_list_second_string():
|
||||
component = Concatenate()
|
||||
res = component.run(first=["This"], second="That")
|
||||
assert res == {"value": ["This", "That"]}
|
||||
|
||||
|
||||
def test_input_first_string_second_list():
|
||||
component = Concatenate()
|
||||
res = component.run(first="This", second=["That"])
|
||||
assert res == {"value": ["This", "That"]}
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Double
|
||||
|
||||
|
||||
def test_double_default():
|
||||
component = Double()
|
||||
results = component.run(value=10)
|
||||
assert results == {"value": 20}
|
||||
@@ -0,0 +1,43 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.testing.sample_components import FString
|
||||
|
||||
|
||||
def test_fstring_with_one_var():
|
||||
fstring = FString(template="Hello, {name}!", variables=["name"])
|
||||
output = fstring.run(name="Alice")
|
||||
assert output == {"string": "Hello, Alice!"}
|
||||
|
||||
|
||||
def test_fstring_with_no_vars():
|
||||
fstring = FString(template="No variables in this template.", variables=[])
|
||||
output = fstring.run()
|
||||
assert output == {"string": "No variables in this template."}
|
||||
|
||||
|
||||
def test_fstring_with_template_at_runtime():
|
||||
fstring = FString(template="Hello {name}", variables=["name"])
|
||||
output = fstring.run(template="Goodbye {name}!", name="Alice")
|
||||
assert output == {"string": "Goodbye Alice!"}
|
||||
|
||||
|
||||
def test_fstring_with_vars_mismatch():
|
||||
fstring = FString(template="Hello {name}", variables=["name"])
|
||||
with pytest.raises(KeyError):
|
||||
fstring.run(template="Goodbye {person}!", name="Alice")
|
||||
|
||||
|
||||
def test_fstring_with_vars_in_excess():
|
||||
fstring = FString(template="Hello {name}", variables=["name"])
|
||||
output = fstring.run(template="Goodbye!", name="Alice")
|
||||
assert output == {"string": "Goodbye!"}
|
||||
|
||||
|
||||
def test_fstring_with_vars_missing():
|
||||
fstring = FString(template="{greeting}, {name}!", variables=["name"])
|
||||
with pytest.raises(KeyError):
|
||||
fstring.run(greeting="Hello")
|
||||
@@ -0,0 +1,15 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import logging
|
||||
|
||||
from haystack.testing.sample_components import Greet
|
||||
|
||||
|
||||
def test_greet_message(caplog):
|
||||
caplog.set_level(logging.WARNING)
|
||||
component = Greet()
|
||||
results = component.run(value=10, message="Hello, that's {value}", log_level="WARNING")
|
||||
assert results == {"value": 10}
|
||||
assert "Hello, that's 10" in caplog.text
|
||||
@@ -0,0 +1,13 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Parity
|
||||
|
||||
|
||||
def test_parity():
|
||||
component = Parity()
|
||||
results = component.run(value=1)
|
||||
assert results == {"odd": 1}
|
||||
results = component.run(value=2)
|
||||
assert results == {"even": 2}
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.testing.sample_components import Remainder
|
||||
|
||||
|
||||
def test_remainder_default():
|
||||
component = Remainder()
|
||||
results = component.run(value=4)
|
||||
assert results == {"remainder_is_1": 4}
|
||||
|
||||
|
||||
def test_remainder_with_divisor():
|
||||
component = Remainder(divisor=4)
|
||||
results = component.run(value=4)
|
||||
assert results == {"remainder_is_0": 4}
|
||||
|
||||
|
||||
def test_remainder_zero():
|
||||
with pytest.raises(ValueError):
|
||||
Remainder(divisor=0)
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Repeat
|
||||
|
||||
|
||||
def test_repeat_default():
|
||||
component = Repeat(outputs=["one", "two"])
|
||||
results = component.run(value=10)
|
||||
assert results == {"one": 10, "two": 10}
|
||||
@@ -0,0 +1,11 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Subtract
|
||||
|
||||
|
||||
def test_subtract():
|
||||
component = Subtract()
|
||||
results = component.run(first_value=10, second_value=7)
|
||||
assert results == {"difference": 3}
|
||||
@@ -0,0 +1,21 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Sum
|
||||
|
||||
|
||||
def test_sum_receives_no_values():
|
||||
component = Sum()
|
||||
results = component.run(values=[])
|
||||
assert results == {"total": 0}
|
||||
|
||||
|
||||
def test_sum_receives_one_value():
|
||||
component = Sum()
|
||||
assert component.run(values=[10]) == {"total": 10}
|
||||
|
||||
|
||||
def test_sum_receives_few_values():
|
||||
component = Sum()
|
||||
assert component.run(values=[10, 2]) == {"total": 12}
|
||||
@@ -0,0 +1,15 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from haystack.testing.sample_components import Threshold
|
||||
|
||||
|
||||
def test_threshold():
|
||||
component = Threshold()
|
||||
|
||||
results = component.run(value=5, threshold=10)
|
||||
assert results == {"below": 5}
|
||||
|
||||
results = component.run(value=15, threshold=10)
|
||||
assert results == {"above": 15}
|
||||
@@ -0,0 +1,551 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from haystack import Document, Pipeline, SuperComponent, component, super_component
|
||||
from haystack.components.builders import AnswerBuilder, PromptBuilder
|
||||
from haystack.components.joiners import DocumentJoiner
|
||||
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
|
||||
from haystack.core.pipeline.base import component_from_dict, component_to_dict
|
||||
from haystack.core.serialization import default_from_dict, default_to_dict
|
||||
from haystack.core.super_component.super_component import InvalidMappingTypeError, InvalidMappingValueError
|
||||
from haystack.dataclasses import GeneratedAnswer
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
from haystack.document_stores.types import DuplicatePolicy
|
||||
from haystack.testing.sample_components import AddFixedValue, Double
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def documents():
|
||||
"""Create test documents for the document store."""
|
||||
return [
|
||||
Document(content="Paris is the capital of France."),
|
||||
Document(content="Berlin is the capital of Germany."),
|
||||
Document(content="Rome is the capital of Italy."),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def document_store(documents):
|
||||
"""Create and populate a test document store."""
|
||||
store = InMemoryDocumentStore()
|
||||
store.write_documents(documents, policy=DuplicatePolicy.OVERWRITE)
|
||||
yield store
|
||||
store.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rag_pipeline(document_store):
|
||||
"""Create a simple RAG pipeline."""
|
||||
|
||||
@component
|
||||
class FakeGenerator:
|
||||
@component.output_types(replies=list[str])
|
||||
def run(self, prompt: str, **kwargs: Any) -> dict[str, list[str]]:
|
||||
return {"replies": ["This is a test response about capitals."]}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
|
||||
pipeline.add_component(
|
||||
"prompt_builder",
|
||||
PromptBuilder(
|
||||
template="Given these documents: {{documents|join(', ',attribute='content')}} Answer: {{query}}",
|
||||
required_variables="*",
|
||||
),
|
||||
)
|
||||
pipeline.add_component("llm", FakeGenerator())
|
||||
pipeline.add_component("answer_builder", AnswerBuilder())
|
||||
pipeline.add_component("joiner", DocumentJoiner())
|
||||
|
||||
pipeline.connect("retriever", "prompt_builder.documents")
|
||||
pipeline.connect("prompt_builder", "llm")
|
||||
pipeline.connect("llm.replies", "answer_builder.replies")
|
||||
pipeline.connect("retriever.documents", "joiner.documents")
|
||||
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_rag_pipeline(document_store):
|
||||
"""Create a simple asyncRAG pipeline."""
|
||||
|
||||
@component
|
||||
class FakeGenerator:
|
||||
@component.output_types(replies=list[str])
|
||||
def run(self, prompt: str, **kwargs: Any) -> dict[str, list[str]]:
|
||||
return {"replies": ["This is a test response about capitals."]}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
|
||||
pipeline.add_component(
|
||||
"prompt_builder",
|
||||
PromptBuilder(
|
||||
template="Given these documents: {{documents|join(', ',attribute='content')}} Answer: {{query}}",
|
||||
required_variables="*",
|
||||
),
|
||||
)
|
||||
pipeline.add_component("llm", FakeGenerator())
|
||||
pipeline.add_component("answer_builder", AnswerBuilder())
|
||||
pipeline.add_component("joiner", DocumentJoiner())
|
||||
|
||||
pipeline.connect("retriever", "prompt_builder.documents")
|
||||
pipeline.connect("prompt_builder", "llm")
|
||||
pipeline.connect("llm.replies", "answer_builder.replies")
|
||||
pipeline.connect("retriever.documents", "joiner.documents")
|
||||
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_super_component():
|
||||
"""Creates a sample SuperComponent for testing visualization methods"""
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("comp1", AddFixedValue(add=3))
|
||||
pipe.add_component("comp2", Double())
|
||||
pipe.connect("comp1.result", "comp2.value")
|
||||
|
||||
return SuperComponent(pipeline=pipe)
|
||||
|
||||
|
||||
@super_component
|
||||
class CustomSuperComponent:
|
||||
def __init__(self, var1: int, var2: str = "test"):
|
||||
self.var1 = var1
|
||||
self.var2 = var2
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("joiner", DocumentJoiner())
|
||||
self.pipeline = pipeline
|
||||
|
||||
|
||||
class TestSuperComponent:
|
||||
def test_split_component_path(self):
|
||||
path = "router.chat_query"
|
||||
components = SuperComponent._split_component_path(path)
|
||||
assert components == ("router", "chat_query")
|
||||
|
||||
def test_split_component_path_error(self):
|
||||
path = "router"
|
||||
with pytest.raises(InvalidMappingValueError):
|
||||
SuperComponent._split_component_path(path)
|
||||
|
||||
def test_invalid_input_mapping_type(self, rag_pipeline):
|
||||
input_mapping = {"search_query": "not_a_list"} # Should be a list
|
||||
with pytest.raises(InvalidMappingTypeError):
|
||||
SuperComponent(pipeline=rag_pipeline, input_mapping=input_mapping) # type: ignore[arg-type]
|
||||
|
||||
def test_invalid_input_mapping_value(self, rag_pipeline):
|
||||
input_mapping = {"search_query": ["nonexistent_component.query"]}
|
||||
with pytest.raises(InvalidMappingValueError):
|
||||
SuperComponent(pipeline=rag_pipeline, input_mapping=input_mapping)
|
||||
|
||||
def test_invalid_output_mapping_type(self, rag_pipeline):
|
||||
output_mapping = {"answer_builder.answers": 123} # Should be a string
|
||||
with pytest.raises(InvalidMappingTypeError):
|
||||
SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping) # type: ignore[arg-type]
|
||||
|
||||
def test_invalid_output_mapping_value(self, rag_pipeline):
|
||||
output_mapping = {"nonexistent_component.answers": "final_answers"}
|
||||
with pytest.raises(InvalidMappingValueError):
|
||||
SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping)
|
||||
|
||||
def test_duplicate_output_names(self, rag_pipeline):
|
||||
output_mapping = {
|
||||
"answer_builder.answers": "final_answers",
|
||||
"llm.replies": "final_answers", # Different path but same output name
|
||||
}
|
||||
with pytest.raises(InvalidMappingValueError):
|
||||
SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping)
|
||||
|
||||
def test_explicit_input_mapping(self, rag_pipeline):
|
||||
input_mapping = {"search_query": ["retriever.query", "prompt_builder.query", "answer_builder.query"]}
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline, input_mapping=input_mapping)
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(input_sockets.keys()) == {"search_query"}
|
||||
assert input_sockets["search_query"].type == str
|
||||
|
||||
def test_explicit_output_mapping(self, rag_pipeline):
|
||||
output_mapping = {"answer_builder.answers": "final_answers"}
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping)
|
||||
output_sockets = wrapper.__haystack_output__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(output_sockets.keys()) == {"final_answers"}
|
||||
assert output_sockets["final_answers"].type == list[GeneratedAnswer]
|
||||
|
||||
def test_auto_input_mapping(self, rag_pipeline):
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline)
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(input_sockets.keys()) == {
|
||||
"documents",
|
||||
"expand_reference_ranges",
|
||||
"filters",
|
||||
"meta",
|
||||
"pattern",
|
||||
"query",
|
||||
"reference_pattern",
|
||||
"scale_score",
|
||||
"template",
|
||||
"template_variables",
|
||||
"top_k",
|
||||
}
|
||||
|
||||
def test_auto_output_mapping(self, rag_pipeline):
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline)
|
||||
output_sockets = wrapper.__haystack_output__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(output_sockets.keys()) == {"answers", "documents"}
|
||||
|
||||
def test_auto_mapping_sockets(self, rag_pipeline):
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline)
|
||||
|
||||
output_sockets = wrapper.__haystack_output__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(output_sockets.keys()) == {"answers", "documents"}
|
||||
assert output_sockets["answers"].type == list[GeneratedAnswer]
|
||||
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert set(input_sockets.keys()) == {
|
||||
"documents",
|
||||
"expand_reference_ranges",
|
||||
"filters",
|
||||
"meta",
|
||||
"pattern",
|
||||
"query",
|
||||
"reference_pattern",
|
||||
"scale_score",
|
||||
"template",
|
||||
"template_variables",
|
||||
"top_k",
|
||||
}
|
||||
assert input_sockets["query"].type == str
|
||||
|
||||
def test_super_component_run(self, rag_pipeline):
|
||||
input_mapping = {"search_query": ["retriever.query", "prompt_builder.query", "answer_builder.query"]}
|
||||
output_mapping = {"answer_builder.answers": "final_answers"}
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline, input_mapping=input_mapping, output_mapping=output_mapping)
|
||||
wrapper.warm_up()
|
||||
result = wrapper.run(search_query="What is the capital of France?")
|
||||
assert "final_answers" in result
|
||||
assert isinstance(result["final_answers"][0], GeneratedAnswer)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_super_component_run_async(self, async_rag_pipeline):
|
||||
input_mapping = {"search_query": ["retriever.query", "prompt_builder.query", "answer_builder.query"]}
|
||||
output_mapping = {"answer_builder.answers": "final_answers"}
|
||||
wrapper = SuperComponent(
|
||||
pipeline=async_rag_pipeline, input_mapping=input_mapping, output_mapping=output_mapping
|
||||
)
|
||||
wrapper.warm_up()
|
||||
result = await wrapper.run_async(search_query="What is the capital of France?")
|
||||
assert "final_answers" in result
|
||||
assert isinstance(result["final_answers"][0], GeneratedAnswer)
|
||||
|
||||
def test_wrapper_serialization(self, document_store):
|
||||
"""Test serialization and deserialization of pipeline wrapper."""
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
|
||||
|
||||
wrapper = SuperComponent(
|
||||
pipeline=pipeline,
|
||||
input_mapping={"query": ["retriever.query"]},
|
||||
output_mapping={"retriever.documents": "documents"},
|
||||
)
|
||||
|
||||
# Test serialization
|
||||
serialized = wrapper.to_dict()
|
||||
assert "type" in serialized
|
||||
assert "init_parameters" in serialized
|
||||
assert "pipeline" in serialized["init_parameters"]
|
||||
|
||||
# Test deserialization
|
||||
deserialized = SuperComponent.from_dict(serialized)
|
||||
assert isinstance(deserialized, SuperComponent)
|
||||
assert deserialized.input_mapping == wrapper.input_mapping
|
||||
assert deserialized.output_mapping == wrapper.output_mapping
|
||||
|
||||
deserialized.warm_up()
|
||||
result = deserialized.run(query="What is the capital of France?")
|
||||
assert "documents" in result
|
||||
assert result["documents"][0].content == "Paris is the capital of France."
|
||||
|
||||
def test_from_dict_ignores_legacy_is_pipeline_async(self, document_store):
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
|
||||
wrapper = SuperComponent(pipeline=pipeline, output_mapping={"retriever.documents": "documents"})
|
||||
|
||||
serialized = wrapper.to_dict()
|
||||
serialized["init_parameters"]["is_pipeline_async"] = True
|
||||
|
||||
deserialized = SuperComponent.from_dict(serialized)
|
||||
assert isinstance(deserialized.pipeline, Pipeline)
|
||||
|
||||
def test_subclass_serialization(self, rag_pipeline):
|
||||
super_comp = SuperComponent(rag_pipeline)
|
||||
serialized = super_comp.to_dict()
|
||||
|
||||
@component
|
||||
class CustomSuperComponent(SuperComponent):
|
||||
def __init__(self, pipeline, instance_attribute="test"):
|
||||
self.instance_attribute = instance_attribute
|
||||
super(CustomSuperComponent, self).__init__(pipeline) # noqa: UP008
|
||||
|
||||
def to_dict(self):
|
||||
return default_to_dict(
|
||||
self, instance_attribute=self.instance_attribute, pipeline=self.pipeline.to_dict()
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
data["init_parameters"]["pipeline"] = Pipeline.from_dict(data["init_parameters"]["pipeline"])
|
||||
return default_from_dict(cls, data)
|
||||
|
||||
custom_super_component = CustomSuperComponent(rag_pipeline)
|
||||
custom_serialized = custom_super_component.to_dict()
|
||||
|
||||
assert custom_serialized["type"] == "test_super_component.CustomSuperComponent"
|
||||
assert custom_super_component._to_super_component_dict() == serialized
|
||||
|
||||
def test_super_component_non_leaf_output(self, rag_pipeline):
|
||||
# 'retriever' is not a leaf, but should now be allowed
|
||||
output_mapping = {"retriever.documents": "retrieved_docs", "answer_builder.answers": "final_answers"}
|
||||
wrapper = SuperComponent(pipeline=rag_pipeline, output_mapping=output_mapping)
|
||||
wrapper.warm_up()
|
||||
result = wrapper.run(query="What is the capital of France?")
|
||||
assert "final_answers" in result # leaf output
|
||||
assert "retrieved_docs" in result # non-leaf output
|
||||
assert isinstance(result["retrieved_docs"][0], Document)
|
||||
|
||||
def test_custom_super_component_to_dict(self, rag_pipeline):
|
||||
custom_super_component = CustomSuperComponent(1)
|
||||
data = component_to_dict(custom_super_component, "custom_super_component")
|
||||
assert data == {
|
||||
"type": "test_super_component.CustomSuperComponent",
|
||||
"init_parameters": {"var1": 1, "var2": "test"},
|
||||
}
|
||||
|
||||
def test_custom_super_component_from_dict(self):
|
||||
data = {"type": "test_super_component.CustomSuperComponent", "init_parameters": {"var1": 1, "var2": "test"}}
|
||||
custom_super_component = component_from_dict(CustomSuperComponent, data, "custom_super_component")
|
||||
assert isinstance(custom_super_component, CustomSuperComponent)
|
||||
assert custom_super_component.var1 == 1
|
||||
assert custom_super_component.var2 == "test"
|
||||
|
||||
@patch("haystack.core.pipeline.Pipeline.show")
|
||||
def test_show_delegates_to_pipeline(self, mock_show, sample_super_component):
|
||||
"""Test that SuperComponent.show() correctly delegates to Pipeline.show() with all parameters"""
|
||||
|
||||
server_url = "https://custom.mermaid.server"
|
||||
params = {"theme": "dark", "format": "svg"}
|
||||
timeout = 60
|
||||
|
||||
sample_super_component.show(server_url=server_url, params=params, timeout=timeout)
|
||||
mock_show.assert_called_once_with(server_url=server_url, params=params, timeout=timeout)
|
||||
|
||||
@patch("haystack.core.pipeline.Pipeline.draw")
|
||||
def test_draw_delegates_to_pipeline(self, mock_draw, sample_super_component, tmp_path):
|
||||
"""Test that SuperComponent.draw() correctly delegates to Pipeline.draw() with all parameters"""
|
||||
|
||||
path = tmp_path / "test_pipeline.png"
|
||||
server_url = "https://custom.mermaid.server"
|
||||
params = {"theme": "dark", "format": "png"}
|
||||
timeout = 60
|
||||
|
||||
sample_super_component.draw(path=path, server_url=server_url, params=params, timeout=timeout)
|
||||
mock_draw.assert_called_once_with(path=path, server_url=server_url, params=params, timeout=timeout)
|
||||
|
||||
@patch("haystack.core.pipeline.Pipeline.show")
|
||||
def test_show_with_default_parameters(self, mock_show, sample_super_component):
|
||||
"""Test that SuperComponent.show() works with default parameters"""
|
||||
|
||||
sample_super_component.show()
|
||||
mock_show.assert_called_once_with(server_url="https://mermaid.ink", params=None, timeout=30)
|
||||
|
||||
@patch("haystack.core.pipeline.Pipeline.draw")
|
||||
def test_draw_with_default_parameters(self, mock_draw, sample_super_component, tmp_path):
|
||||
"""Test that SuperComponent.draw() works with default parameters except path"""
|
||||
|
||||
path = tmp_path / "test_pipeline.png"
|
||||
|
||||
sample_super_component.draw(path=path)
|
||||
mock_draw.assert_called_once_with(path=path, server_url="https://mermaid.ink", params=None, timeout=30)
|
||||
|
||||
def test_input_types_reconciliation(self):
|
||||
"""Test that input types are properly reconciled when they are compatible but not identical."""
|
||||
|
||||
@component
|
||||
class TypeTestComponent:
|
||||
@component.output_types(result_int=int, result_any=Any)
|
||||
def run(self, input_int: int, input_any: Any) -> dict[str, Any]:
|
||||
return {"result_int": input_int, "result_any": input_any}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("test1", TypeTestComponent())
|
||||
pipeline.add_component("test2", TypeTestComponent())
|
||||
|
||||
input_mapping = {"number": ["test1.input_int", "test2.input_any"]}
|
||||
output_mapping = {"test2.result_int": "result_int"}
|
||||
wrapper = SuperComponent(pipeline=pipeline, input_mapping=input_mapping, output_mapping=output_mapping)
|
||||
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert "number" in input_sockets
|
||||
assert input_sockets["number"].type == int
|
||||
|
||||
def test_union_type_reconciliation(self):
|
||||
"""Test that Union types are properly reconciled when creating a SuperComponent."""
|
||||
|
||||
@component
|
||||
class UnionTypeComponent1:
|
||||
@component.output_types(result=int | str)
|
||||
def run(self, inp: int | str) -> dict[str, int | str]:
|
||||
return {"result": inp}
|
||||
|
||||
@component
|
||||
class UnionTypeComponent2:
|
||||
@component.output_types(result=float | str)
|
||||
def run(self, inp: float | str) -> dict[str, float | str]:
|
||||
return {"result": inp}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("test1", UnionTypeComponent1())
|
||||
pipeline.add_component("test2", UnionTypeComponent2())
|
||||
|
||||
input_mapping = {"data": ["test1.inp", "test2.inp"]}
|
||||
output_mapping = {"test2.result": "result"}
|
||||
wrapper = SuperComponent(pipeline=pipeline, input_mapping=input_mapping, output_mapping=output_mapping)
|
||||
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert "data" in input_sockets
|
||||
assert input_sockets["data"].type == str
|
||||
|
||||
def test_input_types_with_any(self):
|
||||
"""Test that Any type is properly handled when reconciling types."""
|
||||
|
||||
@component
|
||||
class AnyTypeComponent:
|
||||
@component.output_types(result=str)
|
||||
def run(self, specific: str, generic: Any) -> dict[str, str]:
|
||||
return {"result": specific}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("test", AnyTypeComponent())
|
||||
|
||||
input_mapping = {"text": ["test.specific", "test.generic"]}
|
||||
wrapper = SuperComponent(pipeline=pipeline, input_mapping=input_mapping)
|
||||
|
||||
input_sockets = wrapper.__haystack_input__._sockets_dict # type: ignore[attr-defined]
|
||||
assert "text" in input_sockets
|
||||
assert input_sockets["text"].type == str
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_super_component_async_serialization_deserialization(self):
|
||||
"""
|
||||
Test for async SuperComponent serialization and deserialization.
|
||||
|
||||
Test that when using the SuperComponent class, a SuperComponent based on an async pipeline can be serialized and
|
||||
deserialized correctly.
|
||||
"""
|
||||
|
||||
@component
|
||||
class AsyncComponent:
|
||||
@component.output_types(output=str)
|
||||
def run(self):
|
||||
return {"output": "irrelevant"}
|
||||
|
||||
@component.output_types(output=str)
|
||||
async def run_async(self):
|
||||
return {"output": "Hello world"}
|
||||
|
||||
pipeline = Pipeline()
|
||||
pipeline.add_component("hello", AsyncComponent())
|
||||
|
||||
async_super_component = SuperComponent(pipeline=pipeline)
|
||||
serialized_super_component = async_super_component.to_dict()
|
||||
|
||||
deserialized_super_component = SuperComponent.from_dict(serialized_super_component)
|
||||
assert isinstance(deserialized_super_component.pipeline, Pipeline)
|
||||
|
||||
result = await deserialized_super_component.run_async()
|
||||
assert result == {"output": "Hello world"}
|
||||
|
||||
|
||||
class TestSuperComponentLifecycle:
|
||||
def test_warm_up_delegates_to_pipeline(self, sample_super_component):
|
||||
with patch.object(sample_super_component.pipeline, "warm_up") as mock_warm_up:
|
||||
sample_super_component.warm_up()
|
||||
mock_warm_up.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_and_async_warm_up_are_independent(self, sample_super_component):
|
||||
with (
|
||||
patch.object(sample_super_component.pipeline, "warm_up") as mock_warm_up,
|
||||
patch.object(sample_super_component.pipeline, "warm_up_async", new=AsyncMock()) as mock_warm_up_async,
|
||||
):
|
||||
sample_super_component.warm_up()
|
||||
mock_warm_up.assert_called_once()
|
||||
await sample_super_component.warm_up_async()
|
||||
mock_warm_up_async.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_warm_up_async_delegates_to_pipeline(self, sample_super_component):
|
||||
with patch.object(sample_super_component.pipeline, "warm_up_async", new=AsyncMock()) as mock_warm_up_async:
|
||||
await sample_super_component.warm_up_async()
|
||||
mock_warm_up_async.assert_awaited_once()
|
||||
|
||||
def test_close_delegates_to_pipeline(self, sample_super_component):
|
||||
with patch.object(sample_super_component.pipeline, "close") as mock_close:
|
||||
sample_super_component.close()
|
||||
mock_close.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_async_delegates_to_pipeline(self, sample_super_component):
|
||||
with patch.object(sample_super_component.pipeline, "close_async", new=AsyncMock()) as mock_close_async:
|
||||
await sample_super_component.close_async()
|
||||
mock_close_async.assert_awaited_once()
|
||||
|
||||
|
||||
@component
|
||||
class _CaptureValue:
|
||||
"""Captures whatever value reaches the inner pipeline, for both sync and async runs."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.captured: dict[str, Any] = {}
|
||||
|
||||
@component.output_types(seen=bool)
|
||||
def run(self, value: Any = None) -> dict[str, bool]:
|
||||
self.captured["value"] = value
|
||||
return {"seen": True}
|
||||
|
||||
@component.output_types(seen=bool)
|
||||
async def run_async(self, value: Any = None) -> dict[str, bool]:
|
||||
self.captured["value"] = value
|
||||
return {"seen": True}
|
||||
|
||||
|
||||
class TestSuperComponentDelegateDefaultFiltering:
|
||||
"""
|
||||
Regression for the `_delegate_default` filter: must be `is not`, not `!=`, otherwise
|
||||
numpy / pandas / torch inputs raise `ValueError: The truth value of ... is ambiguous`.
|
||||
"""
|
||||
|
||||
def test_run_accepts_numpy_ndarray_input(self):
|
||||
inner = _CaptureValue()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("capture", inner)
|
||||
|
||||
result = SuperComponent(pipe).run(value=np.array([1, 2, 3]))
|
||||
|
||||
assert result == {"seen": True}
|
||||
assert np.array_equal(inner.captured["value"], np.array([1, 2, 3]))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_async_accepts_numpy_ndarray_input(self):
|
||||
inner = _CaptureValue()
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("capture", inner)
|
||||
|
||||
result = await SuperComponent(pipe).run_async(value=np.array([4, 5, 6]))
|
||||
|
||||
assert result == {"seen": True}
|
||||
assert np.array_equal(inner.captured["value"], np.array([4, 5, 6]))
|
||||
@@ -0,0 +1,264 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component.types import GreedyVariadic, Variadic
|
||||
from haystack.core.super_component.utils import _is_compatible
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"left,right,expected_common",
|
||||
[
|
||||
(str, str, str),
|
||||
(int, int, int),
|
||||
(float, float, float),
|
||||
(bool, bool, bool),
|
||||
(list, list, list),
|
||||
(dict, dict, dict),
|
||||
(set, set, set),
|
||||
(tuple, tuple, tuple),
|
||||
],
|
||||
)
|
||||
def test_basic_types_compatible(left, right, expected_common):
|
||||
"""Test compatible basic Python types."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert is_compat
|
||||
assert common == expected_common
|
||||
|
||||
|
||||
@pytest.mark.parametrize("left,right", [(str, int), (float, int), (bool, str), (list, dict), (set, tuple)])
|
||||
def test_basic_types_incompatible(left, right):
|
||||
"""Test incompatible basic Python types."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert not is_compat
|
||||
assert common is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"left,right,expected_common",
|
||||
[
|
||||
(Any, Any, Any),
|
||||
# int
|
||||
(int, Any, int),
|
||||
(Any, int, int),
|
||||
# str
|
||||
(Any, str, str),
|
||||
(str, Any, str),
|
||||
# float
|
||||
(float, Any, float),
|
||||
(Any, float, float),
|
||||
# list
|
||||
(list, Any, list),
|
||||
(Any, list, list),
|
||||
# dict
|
||||
(dict, Any, dict),
|
||||
(Any, dict, dict),
|
||||
],
|
||||
)
|
||||
def test_any_type(left, right, expected_common):
|
||||
"""Test Any type compatibility."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert is_compat
|
||||
assert common == expected_common
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"left,right,expected_common",
|
||||
[
|
||||
(int, Union[int, str], int),
|
||||
(Union[int, str], int, int),
|
||||
(Union[int, str], Union[str, int], Union[int, str]),
|
||||
(str, Union[int, str], str),
|
||||
(str, Union[str, None], str),
|
||||
(Union[str, None], str, str),
|
||||
(Union[str, None], Optional[str], Optional[str]),
|
||||
(Optional[str], Union[str, None], Optional[str]),
|
||||
(Union[str, None], Union[str, None], Union[str, None]),
|
||||
],
|
||||
)
|
||||
def test_union_types_compatible(left, right, expected_common):
|
||||
"""Test compatible Union types."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert is_compat
|
||||
assert common == expected_common
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"left,right,expected_common",
|
||||
[
|
||||
(int, int | str, int),
|
||||
(int | str, int, int),
|
||||
(int | str, str | int, int | str),
|
||||
(str, str | None, str),
|
||||
(str | None, str, str),
|
||||
(str | None, str | None, str | None),
|
||||
# Mixed PEP 604 and typing.Union and Optional
|
||||
(int | str, Union[int, str], int | str),
|
||||
(Union[int, str], int | str, Union[int, str]),
|
||||
(str | None, Optional[str], str | None),
|
||||
(Optional[str], str | None, Optional[str]),
|
||||
],
|
||||
)
|
||||
def test_union_types_pep604(left, right, expected_common):
|
||||
"""Test PEP 604 union type compatibility."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert is_compat
|
||||
assert common == expected_common
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"left,right",
|
||||
[
|
||||
(bool, Union[int, str]),
|
||||
(float, Union[int, str]),
|
||||
(Union[int, str], Union[float, bool]),
|
||||
# Pep 604
|
||||
(bool, int | str),
|
||||
(float, int | str),
|
||||
(int | str, float | bool),
|
||||
],
|
||||
)
|
||||
def test_union_types_incompatible(left, right):
|
||||
"""Test incompatible Union types."""
|
||||
is_compat, common = _is_compatible(left, right)
|
||||
assert not is_compat
|
||||
assert common is None
|
||||
|
||||
|
||||
def test_variadic_type_compatibility():
|
||||
"""Test compatibility with Variadic and GreedyVariadic types."""
|
||||
# Basic type compatibility
|
||||
variadic_int = Variadic[int]
|
||||
greedy_int = GreedyVariadic[int]
|
||||
|
||||
is_compat, common = _is_compatible(variadic_int, int)
|
||||
assert is_compat and common == int
|
||||
|
||||
is_compat, common = _is_compatible(int, variadic_int)
|
||||
assert is_compat and common == int
|
||||
|
||||
is_compat, common = _is_compatible(greedy_int, int)
|
||||
assert is_compat and common == int
|
||||
|
||||
is_compat, common = _is_compatible(int, greedy_int)
|
||||
assert is_compat and common == int
|
||||
|
||||
# List type compatibility
|
||||
variadic_list = Variadic[list[int]]
|
||||
greedy_list = GreedyVariadic[list[int]]
|
||||
|
||||
is_compat, common = _is_compatible(variadic_list, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(list[int], variadic_list)
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(greedy_list, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(list[int], greedy_list)
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
# PEP 604 with Variadic
|
||||
variadic_union = Variadic[int | str]
|
||||
|
||||
is_compat, common = _is_compatible(variadic_union, int | str)
|
||||
assert is_compat and common == int | str
|
||||
|
||||
is_compat, common = _is_compatible(int | str, variadic_union)
|
||||
assert is_compat and common == int | str
|
||||
|
||||
# PEP 604 optional with GreedyVariadic
|
||||
greedy_opt = GreedyVariadic[int | None]
|
||||
|
||||
is_compat, common = _is_compatible(greedy_opt, int)
|
||||
assert is_compat and common == int
|
||||
|
||||
is_compat, common = _is_compatible(int, greedy_opt)
|
||||
assert is_compat and common == int
|
||||
|
||||
|
||||
def test_nested_type_unwrapping():
|
||||
"""Test nested type unwrapping behavior with unwrap_nested parameter."""
|
||||
# Test with unwrap_nested=True (default)
|
||||
nested_optional = Variadic[list[int | None]]
|
||||
|
||||
is_compat, common = _is_compatible(nested_optional, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(list[int], nested_optional)
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
nested_union = Variadic[list[Union[int, None]]]
|
||||
|
||||
is_compat, common = _is_compatible(nested_union, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(list[int], nested_union)
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
# PEP 604 (X | Y and X | None syntax)
|
||||
nested_pep604_optional = Variadic[list[int | None]]
|
||||
|
||||
is_compat, common = _is_compatible(nested_pep604_optional, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
is_compat, common = _is_compatible(list[int], nested_pep604_optional)
|
||||
assert is_compat and common == list[int]
|
||||
|
||||
nested_pep604_union = Variadic[list[str | int]]
|
||||
|
||||
is_compat, common = _is_compatible(nested_pep604_union, list[str | int])
|
||||
assert is_compat and common == list[str | int]
|
||||
|
||||
|
||||
def test_complex_nested_types():
|
||||
"""Test complex nested type scenarios."""
|
||||
# Multiple levels of nesting
|
||||
complex_type = Variadic[list[list[Variadic[int]]]]
|
||||
target_type = list[list[int]]
|
||||
|
||||
# With unwrap_nested=True
|
||||
is_compat, common = _is_compatible(complex_type, target_type)
|
||||
assert is_compat and common == list[list[int]]
|
||||
|
||||
is_compat, common = _is_compatible(target_type, complex_type)
|
||||
assert is_compat and common == list[list[int]]
|
||||
|
||||
# With unwrap_nested=False
|
||||
is_compat, common = _is_compatible(complex_type, target_type, unwrap_nested=False)
|
||||
assert not is_compat and common is None
|
||||
|
||||
is_compat, common = _is_compatible(target_type, complex_type, unwrap_nested=False)
|
||||
assert not is_compat and common is None
|
||||
|
||||
|
||||
def test_mixed_variadic_types():
|
||||
"""Test mixing Variadic and GreedyVariadic with other type constructs."""
|
||||
# Variadic with Union
|
||||
var_union = Variadic[Union[int, str]]
|
||||
|
||||
is_compat, common = _is_compatible(var_union, Union[int, str]) # type: ignore[arg-type]
|
||||
assert is_compat and common == Union[int, str]
|
||||
|
||||
is_compat, common = _is_compatible(Union[int, str], var_union) # type: ignore[arg-type]
|
||||
assert is_compat and common == Union[int, str]
|
||||
|
||||
# GreedyVariadic with Optional
|
||||
greedy_opt = GreedyVariadic[Optional[int]]
|
||||
|
||||
is_compat, common = _is_compatible(greedy_opt, int)
|
||||
assert is_compat and common == int
|
||||
|
||||
is_compat, common = _is_compatible(int, greedy_opt)
|
||||
assert is_compat and common == int
|
||||
|
||||
# Nested Variadic and GreedyVariadic
|
||||
nested_var = Variadic[list[GreedyVariadic[int]]]
|
||||
|
||||
is_compat, common = _is_compatible(nested_var, list[int])
|
||||
assert is_compat and common == list[int]
|
||||
@@ -0,0 +1,646 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import sys
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack.core.component import component
|
||||
from haystack.core.errors import DeserializationError, SerializationError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.serialization import (
|
||||
component_from_dict,
|
||||
component_to_dict,
|
||||
default_from_dict,
|
||||
default_to_dict,
|
||||
generate_qualified_class_name,
|
||||
import_class_by_name,
|
||||
)
|
||||
from haystack.document_stores.in_memory import InMemoryDocumentStore
|
||||
from haystack.testing import factory
|
||||
from haystack.utils import ComponentDevice, Secret
|
||||
from haystack.utils.auth import EnvVarSecret
|
||||
from haystack.utils.device import Device, DeviceMap
|
||||
|
||||
|
||||
def test_default_component_to_dict():
|
||||
MyComponent = factory.component_class("MyComponent")
|
||||
comp = MyComponent()
|
||||
res = default_to_dict(comp)
|
||||
assert res == {"type": "haystack.testing.factory.MyComponent", "init_parameters": {}}
|
||||
|
||||
|
||||
def test_default_component_to_dict_with_init_parameters():
|
||||
MyComponent = factory.component_class("MyComponent")
|
||||
comp = MyComponent()
|
||||
res = default_to_dict(comp, some_key="some_value")
|
||||
assert res == {"type": "haystack.testing.factory.MyComponent", "init_parameters": {"some_key": "some_value"}}
|
||||
|
||||
|
||||
def test_default_component_from_dict():
|
||||
def custom_init(self, some_param):
|
||||
self.some_param = some_param
|
||||
|
||||
extra_fields = {"__init__": custom_init}
|
||||
MyComponent = factory.component_class("MyComponent", extra_fields=extra_fields)
|
||||
comp = default_from_dict(
|
||||
MyComponent, {"type": "haystack.testing.factory.MyComponent", "init_parameters": {"some_param": 10}}
|
||||
)
|
||||
assert isinstance(comp, MyComponent)
|
||||
assert comp.some_param == 10 # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def test_default_component_from_dict_without_type():
|
||||
with pytest.raises(DeserializationError, match="Missing 'type' in serialization data"):
|
||||
default_from_dict(Mock, {})
|
||||
|
||||
|
||||
def test_default_component_from_dict_unregistered_component(request):
|
||||
# We use the test function name as component name to make sure it's not registered.
|
||||
# Since the registry is global we risk to have a component with the same name registered in another test.
|
||||
component_name = request.node.name
|
||||
|
||||
with pytest.raises(DeserializationError, match=f"Class '{component_name}' can't be deserialized as 'Mock'"):
|
||||
default_from_dict(Mock, {"type": component_name})
|
||||
|
||||
|
||||
def test_from_dict_import_type():
|
||||
pipeline_dict = {
|
||||
"metadata": {},
|
||||
"components": {
|
||||
"greeter": {
|
||||
"type": "haystack.testing.sample_components.greet.Greet",
|
||||
"init_parameters": {
|
||||
"message": "\nGreeting component says: Hi! The value is {value}\n",
|
||||
"log_level": "INFO",
|
||||
},
|
||||
}
|
||||
},
|
||||
"connections": [],
|
||||
}
|
||||
|
||||
# remove the target component from the registry if already there
|
||||
component.registry.pop("haystack.testing.sample_components.greet.Greet", None)
|
||||
# remove the module from sys.modules if already there
|
||||
sys.modules.pop("haystack.testing.sample_components.greet", None)
|
||||
|
||||
p = Pipeline.from_dict(pipeline_dict)
|
||||
|
||||
from haystack.testing.sample_components.greet import Greet
|
||||
|
||||
assert type(p.get_component("greeter")) == Greet
|
||||
|
||||
|
||||
def test_get_qualified_class_name():
|
||||
MyComponent = factory.component_class("MyComponent")
|
||||
comp = MyComponent()
|
||||
res = generate_qualified_class_name(type(comp))
|
||||
assert res == "haystack.testing.factory.MyComponent"
|
||||
|
||||
|
||||
def test_import_class_by_name():
|
||||
data = "haystack.core.pipeline.Pipeline"
|
||||
class_object = import_class_by_name(data)
|
||||
class_instance = class_object()
|
||||
assert isinstance(class_instance, Pipeline)
|
||||
|
||||
|
||||
def test_import_class_by_name_no_valid_class():
|
||||
# A name that passes the deserialization allowlist but cannot be resolved should raise ImportError.
|
||||
data = "haystack.does.not.exist.Class"
|
||||
with pytest.raises(ImportError):
|
||||
import_class_by_name(data)
|
||||
|
||||
|
||||
def test_import_class_by_name_rejects_untrusted_module():
|
||||
# A module outside the default allowlist is rejected before the import is attempted.
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
import_class_by_name("some.invalid.class")
|
||||
|
||||
|
||||
class CustomData:
|
||||
def __init__(self, a: int, b: str) -> None:
|
||||
self.a = a
|
||||
self.b = b
|
||||
|
||||
|
||||
@component()
|
||||
class UnserializableClass:
|
||||
def __init__(self, a: int, b: str, c: CustomData) -> None:
|
||||
self.a = a
|
||||
self.b = b
|
||||
self.c = c
|
||||
|
||||
def run(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_component_to_dict_invalid_type():
|
||||
with pytest.raises(SerializationError, match="unsupported value of type 'CustomData'"):
|
||||
component_to_dict(UnserializableClass(1, "s", CustomData(99, "aa")), "invalid_component")
|
||||
|
||||
|
||||
@component()
|
||||
class NestedNonStringKeyClass:
|
||||
def __init__(self, mapping: dict[Any, Any]) -> None:
|
||||
self.mapping = mapping
|
||||
|
||||
def run(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_component_to_dict_nested_non_string_key():
|
||||
# A non-string key nested inside init_parameters must be rejected too, not just at the top level.
|
||||
with pytest.raises(SerializationError, match="non-string key"):
|
||||
component_to_dict(NestedNonStringKeyClass(mapping={1: "a"}), "invalid_component")
|
||||
|
||||
|
||||
@component
|
||||
class CustomComponentWithSecrets:
|
||||
def __init__(self, api_key: Secret | None = None, token: Secret | None = None, regular_param: str | None = None):
|
||||
self.api_key = api_key
|
||||
self.token = token
|
||||
self.regular_param = regular_param
|
||||
|
||||
@component.output_types(value=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"value": value}
|
||||
|
||||
|
||||
def test_component_to_dict_with_secret():
|
||||
"""Test that Secret instances are automatically serialized in component_to_dict."""
|
||||
# Test with EnvVarSecret (serializable)
|
||||
env_secret = Secret.from_env_var("TEST_API_KEY")
|
||||
comp = CustomComponentWithSecrets(api_key=env_secret)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["api_key"] == env_secret.to_dict()
|
||||
assert res["init_parameters"]["api_key"]["type"] == "env_var"
|
||||
|
||||
# Test with None
|
||||
comp = CustomComponentWithSecrets(api_key=None)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["api_key"] is None
|
||||
|
||||
# Test with regular value (not a Secret)
|
||||
comp = CustomComponentWithSecrets(regular_param="regular_string")
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["regular_param"] == "regular_string"
|
||||
|
||||
# Test with multiple secrets
|
||||
env_secret1 = Secret.from_env_var("TEST_API_KEY1")
|
||||
env_secret2 = Secret.from_env_var("TEST_API_KEY2")
|
||||
comp = CustomComponentWithSecrets(api_key=env_secret1, token=env_secret2, regular_param="test")
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["api_key"] == env_secret1.to_dict()
|
||||
assert res["init_parameters"]["api_key"]["type"] == "env_var"
|
||||
assert res["init_parameters"]["token"] == env_secret2.to_dict()
|
||||
assert res["init_parameters"]["token"]["type"] == "env_var"
|
||||
assert res["init_parameters"]["regular_param"] == "test"
|
||||
|
||||
|
||||
def test_component_from_dict_with_secret():
|
||||
"""Test that serialized Secret dictionaries are automatically deserialized in component_from_dict."""
|
||||
# Test with EnvVarSecret
|
||||
env_secret = Secret.from_env_var("TEST_API_KEY")
|
||||
serialized_secret = env_secret.to_dict()
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithSecrets),
|
||||
"init_parameters": {"api_key": serialized_secret, "regular_param": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithSecrets, data, "test_component")
|
||||
assert isinstance(comp, CustomComponentWithSecrets)
|
||||
assert isinstance(comp.api_key, Secret)
|
||||
assert comp.api_key.type.value == "env_var"
|
||||
assert comp.regular_param == "test"
|
||||
|
||||
# Test with None
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithSecrets),
|
||||
"init_parameters": {"api_key": None, "regular_param": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithSecrets, data, "test_component")
|
||||
assert comp.api_key is None
|
||||
assert comp.regular_param == "test"
|
||||
|
||||
# Test with regular dict (not a Secret)
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithSecrets),
|
||||
"init_parameters": {"api_key": {"some": "dict"}, "regular_param": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithSecrets, data, "test_component")
|
||||
assert comp.api_key == {"some": "dict"}
|
||||
assert comp.regular_param == "test"
|
||||
|
||||
# Test with multiple secrets
|
||||
env_secret1 = Secret.from_env_var("TEST_API_KEY1")
|
||||
env_secret2 = Secret.from_env_var("TEST_API_KEY2")
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithSecrets),
|
||||
"init_parameters": {"api_key": env_secret1.to_dict(), "token": env_secret2.to_dict(), "regular_param": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithSecrets, data, "test_component")
|
||||
assert isinstance(comp.api_key, Secret)
|
||||
assert isinstance(comp.token, Secret)
|
||||
assert comp.regular_param == "test"
|
||||
|
||||
|
||||
def test_component_from_dict_does_not_mutate_input():
|
||||
"""default_from_dict must not mutate the caller's data dict when deserializing nested objects."""
|
||||
serialized_secret = Secret.from_env_var("TEST_API_KEY").to_dict()
|
||||
data: dict[str, Any] = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithSecrets),
|
||||
"init_parameters": {"api_key": serialized_secret, "regular_param": "test"},
|
||||
}
|
||||
expected = deepcopy(data)
|
||||
|
||||
comp = component_from_dict(CustomComponentWithSecrets, data, "test_component")
|
||||
|
||||
# the returned component is correctly deserialized ...
|
||||
assert isinstance(comp.api_key, Secret)
|
||||
# ... but the caller's dict is left untouched (the Secret sub-dict is not replaced in place)
|
||||
assert data == expected
|
||||
assert isinstance(data["init_parameters"]["api_key"], dict)
|
||||
|
||||
|
||||
def test_component_to_dict_and_from_dict_roundtrip_with_secret():
|
||||
"""Test that serialization and deserialization work together for Secrets."""
|
||||
# Test roundtrip with EnvVarSecret
|
||||
original_secret = Secret.from_env_var("TEST_API_KEY")
|
||||
comp = CustomComponentWithSecrets(api_key=original_secret)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert serialized["init_parameters"]["api_key"]["type"] == "env_var"
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithSecrets, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.api_key, EnvVarSecret)
|
||||
assert deserialized_comp.api_key.type.value == "env_var"
|
||||
assert isinstance(original_secret, EnvVarSecret)
|
||||
assert deserialized_comp.api_key._env_vars == original_secret._env_vars
|
||||
|
||||
# Test roundtrip with multiple secrets
|
||||
env_secret1 = Secret.from_env_var("TEST_API_KEY1")
|
||||
env_secret2 = Secret.from_env_var("TEST_API_KEY2")
|
||||
comp = CustomComponentWithSecrets(api_key=env_secret1, token=env_secret2, regular_param="test")
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert serialized["init_parameters"]["api_key"]["type"] == "env_var"
|
||||
assert serialized["init_parameters"]["token"]["type"] == "env_var"
|
||||
assert serialized["init_parameters"]["regular_param"] == "test"
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithSecrets, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.api_key, EnvVarSecret)
|
||||
assert isinstance(deserialized_comp.token, EnvVarSecret)
|
||||
assert deserialized_comp.api_key.type.value == "env_var"
|
||||
assert deserialized_comp.token.type.value == "env_var"
|
||||
assert deserialized_comp.regular_param == "test"
|
||||
assert isinstance(env_secret1, EnvVarSecret)
|
||||
assert isinstance(env_secret2, EnvVarSecret)
|
||||
assert deserialized_comp.api_key._env_vars == env_secret1._env_vars
|
||||
assert deserialized_comp.token._env_vars == env_secret2._env_vars
|
||||
|
||||
|
||||
@component
|
||||
class CustomComponentWithDevice:
|
||||
def __init__(
|
||||
self,
|
||||
device: ComponentDevice | None = None,
|
||||
other_device: ComponentDevice | None = None,
|
||||
name: str | None = None,
|
||||
):
|
||||
self.device = device
|
||||
self.other_device = other_device
|
||||
self.name = name
|
||||
|
||||
@component.output_types(value=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"value": value}
|
||||
|
||||
|
||||
def test_component_to_dict_with_component_device():
|
||||
"""Test that ComponentDevice instances are automatically serialized in component_to_dict."""
|
||||
# Test with single device (CPU)
|
||||
device = ComponentDevice.from_single(Device.cpu())
|
||||
comp = CustomComponentWithDevice(device=device)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["device"] == {"type": "single", "device": "cpu"}
|
||||
|
||||
# Test with single device (GPU with id)
|
||||
device = ComponentDevice.from_single(Device.gpu(1))
|
||||
comp = CustomComponentWithDevice(device=device)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["device"] == {"type": "single", "device": "cuda:1"}
|
||||
|
||||
# Test with None
|
||||
comp = CustomComponentWithDevice(device=None)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["device"] is None
|
||||
|
||||
# Test with multiple devices (device map)
|
||||
device_map = DeviceMap({"layer1": Device.gpu(0), "layer2": Device.gpu(1)})
|
||||
device = ComponentDevice.from_multiple(device_map)
|
||||
comp = CustomComponentWithDevice(device=device)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["device"] == {
|
||||
"type": "multiple",
|
||||
"device_map": {"layer1": "cuda:0", "layer2": "cuda:1"},
|
||||
}
|
||||
|
||||
# Test with multiple ComponentDevice params
|
||||
device1 = ComponentDevice.from_single(Device.cpu())
|
||||
device2 = ComponentDevice.from_single(Device.gpu(0))
|
||||
comp = CustomComponentWithDevice(device=device1, other_device=device2, name="test")
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["device"] == {"type": "single", "device": "cpu"}
|
||||
assert res["init_parameters"]["other_device"] == {"type": "single", "device": "cuda:0"}
|
||||
assert res["init_parameters"]["name"] == "test"
|
||||
|
||||
|
||||
def test_component_from_dict_with_component_device():
|
||||
"""Test that serialized ComponentDevice dictionaries are automatically deserialized in component_from_dict."""
|
||||
# Test with single device (CPU)
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {"device": {"type": "single", "device": "cpu"}, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert isinstance(comp, CustomComponentWithDevice)
|
||||
assert isinstance(comp.device, ComponentDevice)
|
||||
assert comp.device.to_torch_str() == "cpu"
|
||||
assert comp.name == "test"
|
||||
|
||||
# Test with single device (GPU with id)
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {"device": {"type": "single", "device": "cuda:1"}, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert isinstance(comp.device, ComponentDevice)
|
||||
assert comp.device.to_torch_str() == "cuda:1"
|
||||
|
||||
# Test with None
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {"device": None, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert comp.device is None
|
||||
assert comp.name == "test"
|
||||
|
||||
# Test with multiple devices (device map)
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {"device": {"type": "multiple", "device_map": {"layer1": "cuda:0", "layer2": "cuda:1"}}},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert isinstance(comp.device, ComponentDevice)
|
||||
assert comp.device.has_multiple_devices
|
||||
|
||||
# Test with regular dict (not a ComponentDevice - different structure)
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {"device": {"some": "dict"}, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert comp.device == {"some": "dict"}
|
||||
assert comp.name == "test"
|
||||
|
||||
# Test with multiple ComponentDevice params
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDevice),
|
||||
"init_parameters": {
|
||||
"device": {"type": "single", "device": "cpu"},
|
||||
"other_device": {"type": "single", "device": "cuda:0"},
|
||||
"name": "test",
|
||||
},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDevice, data, "test_component")
|
||||
assert isinstance(comp.device, ComponentDevice)
|
||||
assert isinstance(comp.other_device, ComponentDevice)
|
||||
assert comp.device.to_torch_str() == "cpu"
|
||||
assert comp.other_device.to_torch_str() == "cuda:0"
|
||||
assert comp.name == "test"
|
||||
|
||||
|
||||
def test_component_to_dict_and_from_dict_roundtrip_with_component_device():
|
||||
"""Test that serialization and deserialization work together for ComponentDevice."""
|
||||
# Test roundtrip with single device
|
||||
original_device = ComponentDevice.from_single(Device.cpu())
|
||||
comp = CustomComponentWithDevice(device=original_device)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert serialized["init_parameters"]["device"]["type"] == "single"
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.device, ComponentDevice)
|
||||
assert deserialized_comp.device.to_torch_str() == original_device.to_torch_str()
|
||||
|
||||
# Test roundtrip with GPU device
|
||||
original_device = ComponentDevice.from_single(Device.gpu(2))
|
||||
comp = CustomComponentWithDevice(device=original_device)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component")
|
||||
assert deserialized_comp.device.to_torch_str() == "cuda:2"
|
||||
|
||||
# Test roundtrip with device map
|
||||
device_map = DeviceMap({"layer1": Device.gpu(0), "layer2": Device.cpu()})
|
||||
original_device = ComponentDevice.from_multiple(device_map)
|
||||
comp = CustomComponentWithDevice(device=original_device)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert serialized["init_parameters"]["device"]["type"] == "multiple"
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.device, ComponentDevice)
|
||||
assert deserialized_comp.device.has_multiple_devices
|
||||
|
||||
# Test roundtrip with multiple ComponentDevice params
|
||||
device1 = ComponentDevice.from_single(Device.cpu())
|
||||
device2 = ComponentDevice.from_single(Device.gpu(0))
|
||||
comp = CustomComponentWithDevice(device=device1, other_device=device2, name="test")
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert serialized["init_parameters"]["device"]["type"] == "single"
|
||||
assert serialized["init_parameters"]["other_device"]["type"] == "single"
|
||||
assert serialized["init_parameters"]["name"] == "test"
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDevice, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.device, ComponentDevice)
|
||||
assert isinstance(deserialized_comp.other_device, ComponentDevice)
|
||||
assert deserialized_comp.device.to_torch_str() == "cpu"
|
||||
assert deserialized_comp.other_device.to_torch_str() == "cuda:0"
|
||||
assert deserialized_comp.name == "test"
|
||||
|
||||
|
||||
@component
|
||||
class CustomComponentWithDocumentStore:
|
||||
def __init__(self, document_store: InMemoryDocumentStore | None = None, name: str | None = None):
|
||||
self.document_store = document_store
|
||||
self.name = name
|
||||
|
||||
@component.output_types(value=str)
|
||||
def run(self, value: str) -> dict[str, str]:
|
||||
return {"value": value}
|
||||
|
||||
|
||||
def test_component_to_dict_with_document_store(in_memory_doc_store):
|
||||
"""Test that DocumentStore instances are automatically serialized in component_to_dict."""
|
||||
# Test with InMemoryDocumentStore
|
||||
comp = CustomComponentWithDocumentStore(document_store=in_memory_doc_store)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert "type" in res["init_parameters"]["document_store"]
|
||||
assert "init_parameters" in res["init_parameters"]["document_store"]
|
||||
assert (
|
||||
res["init_parameters"]["document_store"]["type"]
|
||||
== "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore"
|
||||
)
|
||||
|
||||
# Test with None
|
||||
comp = CustomComponentWithDocumentStore(document_store=None)
|
||||
res = component_to_dict(comp, "test_component")
|
||||
assert res["init_parameters"]["document_store"] is None
|
||||
|
||||
|
||||
def test_component_from_dict_with_document_store(in_memory_doc_store):
|
||||
"""Test that serialized DocumentStore dictionaries are automatically deserialized in component_from_dict."""
|
||||
# Test with InMemoryDocumentStore
|
||||
serialized_doc_store = in_memory_doc_store.to_dict()
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {"document_store": serialized_doc_store, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDocumentStore, data, "test_component")
|
||||
assert isinstance(comp, CustomComponentWithDocumentStore)
|
||||
assert isinstance(comp.document_store, InMemoryDocumentStore)
|
||||
assert comp.name == "test"
|
||||
|
||||
# Test with None
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {"document_store": None, "name": "test"},
|
||||
}
|
||||
comp = component_from_dict(CustomComponentWithDocumentStore, data, "test_component")
|
||||
assert comp.document_store is None
|
||||
assert comp.name == "test"
|
||||
|
||||
|
||||
def test_component_to_dict_and_from_dict_roundtrip_with_document_store(in_memory_doc_store):
|
||||
"""Test that serialization and deserialization work together for DocumentStore."""
|
||||
# Test roundtrip with InMemoryDocumentStore
|
||||
comp = CustomComponentWithDocumentStore(document_store=in_memory_doc_store)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
assert "type" in serialized["init_parameters"]["document_store"]
|
||||
assert (
|
||||
serialized["init_parameters"]["document_store"]["type"]
|
||||
== "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore"
|
||||
)
|
||||
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDocumentStore, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.document_store, InMemoryDocumentStore)
|
||||
assert deserialized_comp.document_store.bm25_algorithm == in_memory_doc_store.bm25_algorithm
|
||||
assert (
|
||||
deserialized_comp.document_store.embedding_similarity_function
|
||||
== in_memory_doc_store.embedding_similarity_function
|
||||
)
|
||||
|
||||
# Test roundtrip with custom parameters
|
||||
in_memory_doc_store = InMemoryDocumentStore(
|
||||
bm25_algorithm="BM25Okapi", embedding_similarity_function="cosine", return_embedding=False
|
||||
)
|
||||
comp = CustomComponentWithDocumentStore(document_store=in_memory_doc_store)
|
||||
|
||||
serialized = component_to_dict(comp, "test_component")
|
||||
deserialized_comp = component_from_dict(CustomComponentWithDocumentStore, serialized, "test_component")
|
||||
assert isinstance(deserialized_comp.document_store, InMemoryDocumentStore)
|
||||
assert deserialized_comp.document_store.bm25_algorithm == "BM25Okapi"
|
||||
assert deserialized_comp.document_store.embedding_similarity_function == "cosine"
|
||||
assert deserialized_comp.document_store.return_embedding is False
|
||||
|
||||
|
||||
def test_default_to_dict_with_document_store(in_memory_doc_store):
|
||||
"""Test that DocumentStore instances are automatically serialized in default_to_dict."""
|
||||
res = default_to_dict(in_memory_doc_store)
|
||||
assert res["type"] == "haystack.document_stores.in_memory.document_store.InMemoryDocumentStore"
|
||||
assert "init_parameters" in res
|
||||
|
||||
# Test that DocumentStore is serialized when passed as a parameter
|
||||
comp = CustomComponentWithDocumentStore(document_store=in_memory_doc_store)
|
||||
res = default_to_dict(comp, document_store=in_memory_doc_store, name="test")
|
||||
assert "type" in res["init_parameters"]["document_store"]
|
||||
assert res["init_parameters"]["name"] == "test"
|
||||
|
||||
|
||||
def test_default_from_dict_with_document_store(in_memory_doc_store):
|
||||
"""Test that serialized DocumentStore dictionaries are automatically deserialized in default_from_dict."""
|
||||
serialized = in_memory_doc_store.to_dict()
|
||||
|
||||
# Test direct deserialization
|
||||
deserialized = default_from_dict(InMemoryDocumentStore, serialized)
|
||||
assert isinstance(deserialized, InMemoryDocumentStore)
|
||||
assert deserialized.bm25_algorithm == in_memory_doc_store.bm25_algorithm
|
||||
|
||||
# Test deserialization when DocumentStore is in init_parameters
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {"document_store": serialized, "name": "test"},
|
||||
}
|
||||
comp = default_from_dict(CustomComponentWithDocumentStore, data)
|
||||
assert isinstance(comp.document_store, InMemoryDocumentStore)
|
||||
assert comp.name == "test"
|
||||
|
||||
|
||||
def test_default_from_dict_with_invalid_class_name():
|
||||
"""Test that deserialization raises ImportError with improved message when class cannot be imported."""
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {
|
||||
# Use a class name that passes the allowlist (haystack.*) but cannot be resolved.
|
||||
"document_store": {"type": "haystack.does.not.exist.Class", "init_parameters": {}},
|
||||
"name": "test",
|
||||
},
|
||||
}
|
||||
# Verify the error message includes the parameter key and original error
|
||||
with pytest.raises(
|
||||
ImportError, match=r"Failed to deserialize 'document_store':.*haystack\.does\.not\.exist\.Class"
|
||||
):
|
||||
default_from_dict(CustomComponentWithDocumentStore, data)
|
||||
|
||||
|
||||
def test_default_from_dict_rejects_untrusted_nested_class():
|
||||
"""A nested class with a module outside the allowlist should be rejected."""
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {
|
||||
"document_store": {"type": "nonexistent.module.Class", "init_parameters": {}},
|
||||
"name": "test",
|
||||
},
|
||||
}
|
||||
with pytest.raises(
|
||||
DeserializationError, match=r"Failed to deserialize 'document_store':.*not on the trusted-module allowlist"
|
||||
):
|
||||
default_from_dict(CustomComponentWithDocumentStore, data)
|
||||
|
||||
|
||||
def test_default_from_dict_rejects_unknown_nested_parameter():
|
||||
"""A nested ``{type: ...}`` dict on a parameter that the class does not accept must be rejected
|
||||
before the smuggled type is imported (Option 3: type-aware deserialization)."""
|
||||
data = {
|
||||
"type": generate_qualified_class_name(CustomComponentWithDocumentStore),
|
||||
"init_parameters": {
|
||||
# `payload` is not an init parameter of CustomComponentWithDocumentStore.
|
||||
"payload": {"type": "haystack.testing.factory.MyComponent", "init_parameters": {}},
|
||||
"name": "test",
|
||||
},
|
||||
}
|
||||
with pytest.raises(DeserializationError) as exc_info:
|
||||
default_from_dict(CustomComponentWithDocumentStore, data)
|
||||
message = str(exc_info.value)
|
||||
assert "Refusing to deserialize unknown parameter 'payload'" in message
|
||||
# The message lists the accepted parameters (sorted) and tells the user how to fix it.
|
||||
assert "Valid parameters are: 'document_store', 'name'." in message
|
||||
assert "Correct the parameter name or remove it from the serialized data." in message
|
||||
@@ -0,0 +1,423 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import io
|
||||
import json
|
||||
import subprocess
|
||||
from collections.abc import Callable
|
||||
|
||||
import pytest
|
||||
|
||||
from haystack import component as component_module
|
||||
from haystack.core.errors import DeserializationError
|
||||
from haystack.core.pipeline import Pipeline
|
||||
from haystack.core.serialization import (
|
||||
allow_deserialization_module,
|
||||
default_from_dict,
|
||||
generate_qualified_class_name,
|
||||
import_class_by_name,
|
||||
)
|
||||
from haystack.core.serialization_security import (
|
||||
_DENIED_BUILTIN_NAMES,
|
||||
DESERIALIZATION_ALLOWLIST_ENV_VAR,
|
||||
_check_module_allowed,
|
||||
_current_context,
|
||||
_deserialization_context,
|
||||
_DeserializationContext,
|
||||
_extra_allowed_modules,
|
||||
_is_module_allowed,
|
||||
_module_matches,
|
||||
)
|
||||
from haystack.marshal import YamlMarshaller
|
||||
from haystack.utils import deserialize_callable
|
||||
from haystack.utils.type_serialization import deserialize_type
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_allowlist_state(monkeypatch):
|
||||
"""
|
||||
Force a clean (safe-default, no extra patterns) state for every test in this module so we are
|
||||
testing the actual security model. The top-level test conftest extends the process-wide
|
||||
allowlist with test-only patterns (`test_*`, `pydantic`, ...); we must clear those here so
|
||||
"untrusted" really means untrusted.
|
||||
"""
|
||||
monkeypatch.delenv(DESERIALIZATION_ALLOWLIST_ENV_VAR, raising=False)
|
||||
snapshot = list(_extra_allowed_modules)
|
||||
_extra_allowed_modules.clear()
|
||||
token = _current_context.set(_DeserializationContext())
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_extra_allowed_modules.clear()
|
||||
_extra_allowed_modules.extend(snapshot)
|
||||
_current_context.reset(token)
|
||||
|
||||
|
||||
class TestModuleMatches:
|
||||
def test_prefix_match_equal(self):
|
||||
assert _module_matches("haystack", "haystack")
|
||||
|
||||
def test_prefix_match_submodule(self):
|
||||
assert _module_matches("haystack.components.builders", "haystack")
|
||||
|
||||
def test_prefix_match_strips_trailing_wildcard(self):
|
||||
assert _module_matches("haystack.components", "haystack.*")
|
||||
assert _module_matches("haystack", "haystack.*")
|
||||
|
||||
def test_prefix_match_not_a_partial_word(self):
|
||||
assert not _module_matches("haystack_other", "haystack")
|
||||
|
||||
def test_trailing_star_matches_submodules(self):
|
||||
assert _module_matches("mypkg.components.foo", "mypkg.*")
|
||||
assert _module_matches("mypkg.foo.bar", "mypkg.*")
|
||||
|
||||
def test_trailing_star_does_not_match_unrelated(self):
|
||||
assert not _module_matches("other.foo", "mypkg.*")
|
||||
|
||||
def test_fnmatch_glob_in_middle(self):
|
||||
assert _module_matches("pkg.foo.utils", "pkg.*.utils")
|
||||
assert _module_matches("pkg.bar.utils", "pkg.*.utils")
|
||||
|
||||
def test_fnmatch_glob_in_middle_no_match(self):
|
||||
assert not _module_matches("pkg.foo.helpers", "pkg.*.utils")
|
||||
|
||||
def test_fnmatch_single_char(self):
|
||||
# `?` is an fnmatch wildcard for a single character.
|
||||
assert _module_matches("pkga", "pkg?")
|
||||
assert not _module_matches("pkgab", "pkg?")
|
||||
|
||||
def test_fnmatch_character_class(self):
|
||||
assert _module_matches("data_3", "data_[0-9]")
|
||||
assert not _module_matches("data_x", "data_[0-9]")
|
||||
|
||||
def test_trailing_star_with_wildcards_in_prefix_uses_fnmatch(self):
|
||||
# `j*on.*` has a `*` before the trailing `.*`, so it must NOT be short-circuited to a
|
||||
# prefix match against the literal `j*on`. It should fall through to fnmatch.
|
||||
assert _module_matches("json.tool", "j*on.*")
|
||||
assert _module_matches("jaeon.subpkg.foo", "j*on.*")
|
||||
# Pure fnmatch doesn't match the bare `json` for the pattern `j*on.*` (the `.*` requires
|
||||
# a `.X` part).
|
||||
assert not _module_matches("json", "j*on.*")
|
||||
|
||||
|
||||
class TestAllowlistDefaults:
|
||||
def test_haystack_allowed(self):
|
||||
assert _is_module_allowed("haystack")
|
||||
assert _is_module_allowed("haystack.components.builders.prompt_builder")
|
||||
|
||||
def test_haystack_integrations_allowed(self):
|
||||
assert _is_module_allowed("haystack_integrations.components.retrievers")
|
||||
|
||||
def test_haystack_experimental_allowed(self):
|
||||
assert _is_module_allowed("haystack_experimental")
|
||||
|
||||
def test_typing_allowed(self):
|
||||
assert _is_module_allowed("typing")
|
||||
|
||||
def test_collections_allowed(self):
|
||||
assert _is_module_allowed("collections")
|
||||
assert _is_module_allowed("collections.abc")
|
||||
|
||||
def test_builtins_allowed(self):
|
||||
assert _is_module_allowed("builtins")
|
||||
|
||||
def test_arbitrary_third_party_not_allowed(self):
|
||||
assert not _is_module_allowed("subprocess")
|
||||
assert not _is_module_allowed("os")
|
||||
|
||||
|
||||
class TestAllowDeserializationModule:
|
||||
def test_extends_allowlist(self):
|
||||
assert not _is_module_allowed("mypkg.components")
|
||||
allow_deserialization_module("mypkg")
|
||||
assert _is_module_allowed("mypkg")
|
||||
assert _is_module_allowed("mypkg.components")
|
||||
|
||||
def test_pattern_with_wildcard(self):
|
||||
allow_deserialization_module("mypkg.components.*")
|
||||
assert _is_module_allowed("mypkg.components.foo")
|
||||
|
||||
def test_duplicate_pattern_only_added_once(self):
|
||||
allow_deserialization_module("mypkg")
|
||||
allow_deserialization_module("mypkg")
|
||||
assert _extra_allowed_modules.count("mypkg") == 1
|
||||
|
||||
|
||||
class TestDeserializationContext:
|
||||
def test_extra_allowed_modules_via_context(self):
|
||||
assert not _is_module_allowed("mypkg.thing")
|
||||
with _deserialization_context(allowed_modules=["mypkg"]):
|
||||
assert _is_module_allowed("mypkg.thing")
|
||||
# The per-call extension is reset on exit.
|
||||
assert not _is_module_allowed("mypkg.thing")
|
||||
|
||||
def test_unsafe_bypasses_allowlist(self):
|
||||
assert not _is_module_allowed("subprocess")
|
||||
with _deserialization_context(unsafe=True):
|
||||
assert _is_module_allowed("subprocess")
|
||||
assert _is_module_allowed("any.arbitrary.module")
|
||||
assert not _is_module_allowed("subprocess")
|
||||
|
||||
|
||||
class TestEnvVar:
|
||||
def test_env_var_extends_allowlist(self, monkeypatch):
|
||||
monkeypatch.setenv(DESERIALIZATION_ALLOWLIST_ENV_VAR, "mypkg.components.*,otherpkg")
|
||||
assert _is_module_allowed("mypkg.components.foo")
|
||||
assert _is_module_allowed("otherpkg")
|
||||
assert _is_module_allowed("otherpkg.sub")
|
||||
assert not _is_module_allowed("yetanother")
|
||||
|
||||
def test_env_var_ignores_empty_entries(self, monkeypatch):
|
||||
monkeypatch.setenv(DESERIALIZATION_ALLOWLIST_ENV_VAR, ", ,mypkg,,")
|
||||
assert _is_module_allowed("mypkg.sub")
|
||||
|
||||
|
||||
class TestCheckModuleAllowed:
|
||||
def test_passes_silently_for_allowed_module(self):
|
||||
_check_module_allowed("haystack.foo")
|
||||
|
||||
def test_raises_for_disallowed_module(self):
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
_check_module_allowed("subprocess")
|
||||
|
||||
def test_error_message_suggests_remediations(self):
|
||||
with pytest.raises(DeserializationError) as exc_info:
|
||||
_check_module_allowed("mypkg.evil")
|
||||
message = str(exc_info.value)
|
||||
assert "allowed_modules" in message
|
||||
assert "allow_deserialization_module" in message
|
||||
assert DESERIALIZATION_ALLOWLIST_ENV_VAR in message
|
||||
assert "unsafe=True" in message
|
||||
|
||||
|
||||
class TestImportClassByNameAllowlist:
|
||||
def test_allowlisted_class(self):
|
||||
cls = import_class_by_name("haystack.core.pipeline.Pipeline")
|
||||
assert cls is Pipeline
|
||||
|
||||
def test_rejects_untrusted_module(self):
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
import_class_by_name("subprocess.Popen")
|
||||
|
||||
def test_per_call_extension(self):
|
||||
# subprocess is normally blocked
|
||||
with pytest.raises(DeserializationError):
|
||||
import_class_by_name("subprocess.Popen")
|
||||
# ... but extending the allowlist for a single call lets it through.
|
||||
with _deserialization_context(allowed_modules=["subprocess"]):
|
||||
cls = import_class_by_name("subprocess.Popen")
|
||||
assert cls is subprocess.Popen
|
||||
|
||||
# `type` is excluded: it is a valid class in this path (covered by test_allows_builtin_type),
|
||||
# even though it is denied as a *callable*. Every other denied builtin is a function, not a type.
|
||||
@pytest.mark.parametrize("name", sorted(_DENIED_BUILTIN_NAMES - {"type"}))
|
||||
def test_rejects_dangerous_builtin(self, name):
|
||||
# `builtins` passes the module allowlist, so a class reference must resolve to an actual
|
||||
# type — otherwise a nested `{"type": "builtins.compile"}` payload would resolve a function.
|
||||
with pytest.raises(DeserializationError, match="not a type"):
|
||||
import_class_by_name(f"builtins.{name}")
|
||||
|
||||
def test_allows_builtin_type(self):
|
||||
# Builtin types must still import as classes. `type` is a valid class here even though it
|
||||
# is blocked as a *callable* by deserialize_callable.
|
||||
assert import_class_by_name("builtins.dict") is dict
|
||||
assert import_class_by_name("builtins.type") is type
|
||||
|
||||
|
||||
class TestDefaultFromDictNestedBuiltins:
|
||||
def test_nested_dangerous_builtin_rejected(self):
|
||||
"""A nested `{"type": "builtins.<dangerous>"}` payload must be rejected via the same gate."""
|
||||
|
||||
class Container:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
data = {
|
||||
"type": generate_qualified_class_name(Container),
|
||||
"init_parameters": {"payload": {"type": "builtins.compile", "init_parameters": {}}},
|
||||
}
|
||||
with pytest.raises(DeserializationError, match="not a type"):
|
||||
default_from_dict(Container, data)
|
||||
|
||||
|
||||
class TestDeserializeCallableAllowlist:
|
||||
"""
|
||||
`deserialize_callable` walks progressively-shorter module prefixes when resolving a dotted
|
||||
name. The allowlist check must apply to "is *any* prefix on the allowlist?", not to each
|
||||
individual candidate — otherwise fnmatch patterns that match the actual module but not the
|
||||
full handle (e.g. `j*on` matches `json` but not `json.dumps`) would be wrongly rejected.
|
||||
"""
|
||||
|
||||
def test_fnmatch_pattern_matches_shorter_prefix(self):
|
||||
# `j*on` matches `json` (the actual module) but not `json.dumps` (the full handle).
|
||||
# The deferred allowlist check should still accept this.
|
||||
with _deserialization_context(allowed_modules=["j*on"]):
|
||||
fn = deserialize_callable("json.dumps")
|
||||
assert fn is json.dumps
|
||||
|
||||
def test_rejects_when_no_prefix_matches(self):
|
||||
# No prefix of `subprocess.Popen` matches the default allowlist (or `unrelated`). This also
|
||||
# covers the plain default-allowlist case, since the context only ever appends patterns.
|
||||
with _deserialization_context(allowed_modules=["unrelated"]):
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
deserialize_callable("subprocess.Popen")
|
||||
|
||||
|
||||
class TestDeniedBuiltins:
|
||||
"""
|
||||
`builtins` is on the default allowlist so harmless members (`builtins.print`, builtin types)
|
||||
round-trip, but the module-granular allowlist is too coarse to stop dangerous builtin callables
|
||||
like `eval`/`exec`. Those are blocked in the callable-resolution path regardless of the
|
||||
allowlist (parametrized over the canonical denied set so additions are covered automatically).
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("name", _DENIED_BUILTIN_NAMES)
|
||||
def test_dangerous_builtin_callable_rejected(self, name):
|
||||
with pytest.raises(DeserializationError, match="blocked because it can be used"):
|
||||
deserialize_callable(f"builtins.{name}")
|
||||
|
||||
def test_harmless_builtin_callable_still_resolves(self):
|
||||
# `serialize_callable(print)` emits "builtins.print"; harmless builtins must keep
|
||||
# round-tripping (only the denied set is blocked).
|
||||
assert deserialize_callable("builtins.print") is print
|
||||
assert deserialize_callable("builtins.len") is len
|
||||
assert deserialize_callable("builtins.sorted") is sorted
|
||||
|
||||
def test_alias_to_dangerous_builtin_rejected(self):
|
||||
# `io.open is builtins.open`, so the identity-based check catches it once `io` is allowed.
|
||||
with _deserialization_context(allowed_modules=["io"]):
|
||||
with pytest.raises(DeserializationError, match="blocked because it can be used"):
|
||||
deserialize_callable("io.open")
|
||||
|
||||
def test_type_blocked_as_callable_but_allowed_as_type(self):
|
||||
# `type` is a class-creation gadget as a callable, but a legitimate type annotation.
|
||||
with pytest.raises(DeserializationError, match="blocked because it can be used"):
|
||||
deserialize_callable("builtins.type")
|
||||
assert deserialize_type("type") is type
|
||||
|
||||
def test_unsafe_mode_bypasses_the_block(self):
|
||||
# `unsafe=True` disables all deserialization safety checks by design.
|
||||
with _deserialization_context(unsafe=True):
|
||||
assert deserialize_callable("builtins.eval") is eval
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _registered_untrusted_component():
|
||||
"""
|
||||
Set up a fake component class registered under a fully-qualified name in an untrusted module
|
||||
(`evilpkg.evilmod.EvilComponent`). Yields a dict payload referencing it. The fixture cleans
|
||||
up the registry on teardown.
|
||||
"""
|
||||
fake_type = "evilpkg.evilmod.EvilComponent"
|
||||
|
||||
@component_module
|
||||
class EvilComponent:
|
||||
@component_module.output_types(value=int)
|
||||
def run(self, value: int) -> dict[str, int]:
|
||||
return {"value": value}
|
||||
|
||||
registry = component_module.registry
|
||||
original = registry.get(fake_type)
|
||||
registry[fake_type] = EvilComponent
|
||||
try:
|
||||
yield {
|
||||
"fake_type": fake_type,
|
||||
"data": {
|
||||
"metadata": {},
|
||||
"components": {"evil": {"type": fake_type, "init_parameters": {}}},
|
||||
"connections": [],
|
||||
},
|
||||
}
|
||||
finally:
|
||||
if original is None:
|
||||
registry.pop(fake_type, None)
|
||||
else:
|
||||
registry[fake_type] = original
|
||||
|
||||
|
||||
class TestPipelineFromDictAllowlistBypass:
|
||||
def test_pre_registered_untrusted_component_is_rejected(self, _registered_untrusted_component):
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
Pipeline.from_dict(_registered_untrusted_component["data"])
|
||||
|
||||
def test_pre_registered_component_loadable_with_allowed_modules(self, _registered_untrusted_component):
|
||||
"""
|
||||
Counterpart to the bypass test: once the user opts the module into the allowlist, the
|
||||
load gets past the allowlist gate. (It still fails downstream because the fake type name
|
||||
doesn't match the test class's real qualified name — that's expected and proves the
|
||||
allowlist gate, not a downstream check, is what changed.)
|
||||
"""
|
||||
data = _registered_untrusted_component["data"]
|
||||
# Without allowed_modules, this is rejected as untrusted.
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
Pipeline.from_dict(data)
|
||||
# With the matching pattern, the allowlist gate passes; the failure now comes from
|
||||
# the qualified-name mismatch in default_from_dict — a downstream check.
|
||||
with pytest.raises(DeserializationError, match="can't be deserialized as"):
|
||||
Pipeline.from_dict(data, allowed_modules=["evilpkg.*"])
|
||||
|
||||
|
||||
class TestPipelineLoadAndLoadsPropagation:
|
||||
"""
|
||||
Verify that the security kwargs added to `Pipeline.from_dict` are propagated correctly
|
||||
through the `Pipeline.loads` (string) and `Pipeline.load` (file-like) entry points, and that
|
||||
they produce equivalent behavior to calling `from_dict` directly.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _yaml_for(data: dict) -> str:
|
||||
# We can't round-trip through `Pipeline.from_dict` + `dumps` because the registered
|
||||
# `EvilComponent`'s real qualified name doesn't match the fake type — the inner
|
||||
# `default_from_dict` would reject it. Build the YAML directly via the marshaller instead.
|
||||
return YamlMarshaller().marshal(data)
|
||||
|
||||
def test_loads_rejects_untrusted_by_default(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
Pipeline.loads(yaml_str)
|
||||
|
||||
def test_loads_propagates_allowed_modules(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
# With the matching pattern, the allowlist gate passes; downstream we get the type
|
||||
# mismatch — proving the kwarg reached the gate.
|
||||
with pytest.raises(DeserializationError, match="can't be deserialized as"):
|
||||
Pipeline.loads(yaml_str, allowed_modules=["evilpkg.*"])
|
||||
|
||||
def test_loads_propagates_unsafe(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
# `unsafe=True` bypasses the allowlist entirely; downstream we still get the type mismatch.
|
||||
with pytest.raises(DeserializationError, match="can't be deserialized as"):
|
||||
Pipeline.loads(yaml_str, unsafe=True)
|
||||
|
||||
def test_load_rejects_untrusted_by_default(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
with pytest.raises(DeserializationError, match="not on the trusted-module allowlist"):
|
||||
Pipeline.load(io.StringIO(yaml_str))
|
||||
|
||||
def test_load_propagates_allowed_modules(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
with pytest.raises(DeserializationError, match="can't be deserialized as"):
|
||||
Pipeline.load(io.StringIO(yaml_str), allowed_modules=["evilpkg.*"])
|
||||
|
||||
def test_load_propagates_unsafe(self, _registered_untrusted_component):
|
||||
yaml_str = self._yaml_for(_registered_untrusted_component["data"])
|
||||
with pytest.raises(DeserializationError, match="can't be deserialized as"):
|
||||
Pipeline.load(io.StringIO(yaml_str), unsafe=True)
|
||||
|
||||
def test_load_loads_from_dict_equivalent_on_rejection(self, _registered_untrusted_component):
|
||||
"""All three entry points produce the same rejection message for the same untrusted payload."""
|
||||
data = _registered_untrusted_component["data"]
|
||||
yaml_str = self._yaml_for(data)
|
||||
|
||||
def _capture(callable_: Callable[[], object]) -> str:
|
||||
with pytest.raises(DeserializationError) as exc_info:
|
||||
callable_()
|
||||
return str(exc_info.value)
|
||||
|
||||
from_dict_msg = _capture(lambda: Pipeline.from_dict(data))
|
||||
loads_msg = _capture(lambda: Pipeline.loads(yaml_str))
|
||||
load_msg = _capture(lambda: Pipeline.load(io.StringIO(yaml_str)))
|
||||
|
||||
assert "not on the trusted-module allowlist" in from_dict_msg
|
||||
assert from_dict_msg == loads_msg == load_msg
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user