diff --git a/haystack/__init__.py b/haystack/__init__.py index 730f8b19f5..282479f0c3 100644 --- a/haystack/__init__.py +++ b/haystack/__init__.py @@ -14,7 +14,7 @@ from haystack.core.component import component from haystack.core.errors import ComponentError, DeserializationError from haystack.core.pipeline import Pipeline -from haystack.core.serialization import default_from_dict, default_to_dict +from haystack.core.serialization import coerce_pipeline_inputs, default_from_dict, default_to_dict from haystack.core.super_component.super_component import SuperComponent, super_component from haystack.dataclasses import Answer, Document, ExtractedAnswer, GeneratedAnswer from haystack.version import __version__ # noqa: F401 @@ -34,6 +34,7 @@ "Pipeline", "SuperComponent", "super_component", + "coerce_pipeline_inputs", "component", "default_from_dict", "default_to_dict", diff --git a/haystack/core/serialization.py b/haystack/core/serialization.py index fec9780eed..ade1b13d30 100644 --- a/haystack/core/serialization.py +++ b/haystack/core/serialization.py @@ -4,8 +4,10 @@ import inspect from collections.abc import Callable, Iterable -from dataclasses import dataclass -from typing import Any, TypeVar +from dataclasses import dataclass, is_dataclass +from typing import TYPE_CHECKING, Any, TypeVar, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter from haystack.core.component.component import _hook_component_init from haystack.core.errors import DeserializationError, SerializationError @@ -16,7 +18,10 @@ from haystack.core.serialization_security import allow_deserialization_module as allow_deserialization_module from haystack.utils.auth import Secret from haystack.utils.device import ComponentDevice -from haystack.utils.type_serialization import _import_class_by_name +from haystack.utils.type_serialization import _import_class_by_name, _is_union_type + +if TYPE_CHECKING: + from haystack.core.pipeline.base import PipelineBase T = TypeVar("T") @@ -379,3 +384,143 @@ def import_class_by_name(fully_qualified_name: str) -> type[object]: :raises DeserializationError: If the module is not on the deserialization allowlist. """ return _import_class_by_name(fully_qualified_name) + + +def _coercible_classes(type_: Any) -> set[type]: + """ + Collect the distinct classes involved in `type_` that can deserialize a plain dictionary. + + A class is coercible if it exposes a `from_dict` method (Haystack dataclasses and Components), is a Pydantic + model, or is a standard-library dataclass. Follows `list[T]` and optionals/unions of those. + + :param type_: The type to inspect. + :returns: The set of coercible classes `type_` involves (empty if none). + """ + if _is_union_type(type_): + classes: set[type] = set() + for arm in get_args(type_): + classes |= _coercible_classes(arm) + return classes + if get_origin(type_) is list: + args = get_args(type_) + return _coercible_classes(args[0]) if args else set() + if isinstance(type_, type) and (hasattr(type_, "from_dict") or issubclass(type_, BaseModel) or is_dataclass(type_)): + return {type_} + return set() + + +def _deserialize_one(value: dict[str, Any], target_class: Any) -> Any: + """ + Deserialize a single dictionary into an instance of `target_class`. + + `from_dict`-capable classes take priority so Haystack objects keep their native deserialization. Pydantic + models and standard-library dataclasses are deserialized with Pydantic. + + :param value: The dictionary to deserialize. + :param target_class: The coercible class resolved for the value's socket type. + :returns: The deserialized instance. + """ + if hasattr(target_class, "from_dict"): + return target_class.from_dict(value) + if issubclass(target_class, BaseModel): + return target_class.model_validate(value) + return TypeAdapter(target_class).validate_python(value) + + +def _deserialize_from_dict(value: Any, target_class: Any) -> Any: + """ + Deserialize `value` into instances of `target_class`. + + Dictionaries are deserialized directly, lists element-wise (leaving non-dictionary items untouched). Any other + value is returned unchanged. + + :param value: The value to deserialize. + :param target_class: The coercible class resolved for the value's socket type. + :returns: The deserialized value. + """ + if isinstance(value, dict): + return _deserialize_one(value, target_class) + if isinstance(value, list): + return [_deserialize_one(item, target_class) if isinstance(item, dict) else item for item in value] + return value + + +def _needs_coercion(value: Any) -> bool: + """Whether `value` carries dictionaries that could be deserialized (a dict, or a list containing one).""" + return isinstance(value, dict) or (isinstance(value, list) and any(isinstance(item, dict) for item in value)) + + +def _coerce_input_value(value: Any, socket_types: list[Any]) -> Any: + if not _needs_coercion(value): + return value + for type_ in socket_types: + classes = _coercible_classes(type_) + if len(classes) > 1: + names = ", ".join(sorted(cls.__name__ for cls in classes)) + raise DeserializationError( + f"Cannot coerce input for socket type '{type_}': it has multiple deserializable members " + f"({names}) and the serialized payload carries no type information to choose between them. " + f"Provide this input as an already-deserialized object." + ) + if classes: + return _deserialize_from_dict(value, next(iter(classes))) + return value + + +def coerce_pipeline_inputs(pipeline: "PipelineBase", data: dict[str, Any]) -> dict[str, Any]: + """ + Deserialize serialized Haystack objects in pipeline input data, based on the pipeline's input socket types. + + For every provided value whose input socket type involves a coercible class, plain dictionaries are + converted into instances of that class. A class is coercible if it exposes a `from_dict` method (such as + `ChatMessage` or `Document`), is a Pydantic model, or is a standard-library dataclass. Socket types of the + form `T`, `list[T]`, and optionals/unions of those are supported. Values that are already deserialized, or + whose socket types involve no coercible class, are returned unchanged. A socket type that involves more than + one distinct coercible class (such as `GeneratedAnswer | ExtractedAnswer`) is ambiguous, since the payload + carries no type information to select one; coercing a dictionary against it raises a `DeserializationError`. + + The dictionaries are expected to be in the format produced by the class's serialization: `to_dict` for + `from_dict`-capable classes, `model_dump` for Pydantic models, and `dataclasses.asdict` for standard-library + dataclasses. A Pydantic model or dataclass may nest `from_dict`-capable Haystack objects: Pydantic + (de)serializes them via their fields, so a `model_dump` payload round-trips through `model_validate` without + going through `to_dict`/`from_dict`. + + Like `Pipeline.run`, `data` accepts the nested format (`{"component_name": {"input_name": value}}`) and + the flat format (`{"input_name": value}`). The same format detection rules apply and the returned + dictionary keeps the format of `data`. + + Usage example: + ```python + serialized_inputs = {"agent": {"messages": [message.to_dict() for message in messages]}} + inputs = coerce_pipeline_inputs(pipeline, serialized_inputs) + result = pipeline.run(inputs) + ``` + + :param pipeline: The pipeline whose input socket types drive the coercion. + :param data: The pipeline input data, in nested or flat format. + :returns: A new dictionary with the same structure as `data` and deserialized values. + :raises DeserializationError: If a socket type involves more than one distinct coercible class and the + corresponding value is a serialized dictionary. + :raises Exception: Any error raised while deserializing a dictionary that does not match the expected format, + such as a `from_dict` error or a Pydantic `ValidationError`. + """ + # mirrors the format detection in Pipeline.run: nested if all values are dictionaries + if all(isinstance(value, dict) for value in data.values()): + available_inputs = pipeline.inputs(include_components_with_connected_inputs=True) + coerced: dict[str, Any] = {} + for component_name, component_inputs in data.items(): + sockets = available_inputs.get(component_name, {}) + coerced[component_name] = { + input_name: _coerce_input_value(value, [sockets[input_name]["type"]] if input_name in sockets else []) + for input_name, value in component_inputs.items() + } + return coerced + + # flat format: input names are matched across components like in Pipeline.run + available_inputs = pipeline.inputs() + return { + input_name: _coerce_input_value( + value, [sockets[input_name]["type"] for sockets in available_inputs.values() if input_name in sockets] + ) + for input_name, value in data.items() + } diff --git a/haystack/tools/component_tool.py b/haystack/tools/component_tool.py index 2e1682a44c..11eaa46371 100644 --- a/haystack/tools/component_tool.py +++ b/haystack/tools/component_tool.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from collections.abc import Callable -from typing import Any, get_args, get_origin +from typing import Any from pydantic import Field, TypeAdapter, create_model @@ -11,6 +11,8 @@ from haystack.components.agents.state.state import State from haystack.core.component import Component from haystack.core.serialization import ( + _coercible_classes, + _deserialize_from_dict, component_from_dict, component_to_dict, generate_qualified_class_name, @@ -31,7 +33,6 @@ _serialize_outputs_to_state, _serialize_outputs_to_string, ) -from haystack.utils.type_serialization import _is_union_type logger = logging.getLogger(__name__) @@ -391,31 +392,10 @@ def _convert_param(self, param_value: Any, param_type: type) -> Any: :returns: The converted parameter value. """ - # We unwrap optional types so we can support types like messages: list[ChatMessage] | None - unwrapped_param_type = _unwrap_optional(param_type) - - # We handle union types (e.g. list[ChatMessage] | str) by extracting the list[T] arm that has - # from_dict, so the conversion below works the same as for plain list[T]. Other arms like str - # need no special handling and fall through to Pydantic or the plain return at the end. - if _is_union_type(unwrapped_param_type): - list_arms = [ - a - for a in get_args(unwrapped_param_type) - if get_origin(a) is list and get_args(a) and hasattr(get_args(a)[0], "from_dict") - ] - unwrapped_param_type = list_arms[0] if list_arms else unwrapped_param_type - - # We support calling from_dict on target types that have it, even if they are wrapped in a list. - # This allows us to support lists of dataclasses as well as single dataclass inputs. - target_type = ( - get_args(unwrapped_param_type)[0] if get_origin(unwrapped_param_type) is list else unwrapped_param_type - ) - if hasattr(target_type, "from_dict"): - if isinstance(param_value, list): - return [target_type.from_dict(item) if isinstance(item, dict) else item for item in param_value] - if isinstance(param_value, dict): - return target_type.from_dict(param_value) - return param_value - - # Use the original type for pydantic validation + # Types involving a single coercible class (e.g. list[ChatMessage] | None) are deserialized element-wise + # via that class; everything else is validated with Pydantic. + classes = _coercible_classes(param_type) + if len(classes) == 1: + return _deserialize_from_dict(param_value, next(iter(classes))) + return TypeAdapter(param_type).validate_python(param_value) diff --git a/releasenotes/notes/coerce-pipeline-inputs-utility-d38e32c0ca3b76cd.yaml b/releasenotes/notes/coerce-pipeline-inputs-utility-d38e32c0ca3b76cd.yaml new file mode 100644 index 0000000000..ec2de08f21 --- /dev/null +++ b/releasenotes/notes/coerce-pipeline-inputs-utility-d38e32c0ca3b76cd.yaml @@ -0,0 +1,19 @@ +--- +features: + - | + Added the ``coerce_pipeline_inputs`` utility function (importable from ``haystack``). + It deserializes serialized Haystack objects in pipeline input data based on the pipeline's input + socket types: for every provided value whose socket type involves a class exposing ``from_dict`` + (such as ``ChatMessage`` or ``Document``), plain dictionaries are converted into instances of that + class. Socket types of the form ``T``, ``list[T]``, and optionals/unions of those are supported, + and both the nested (``{"component_name": {"input_name": value}}``) and flat + (``{"input_name": value}``) input formats are accepted. This removes the need for API layers to + guess which pipeline inputs contain serialized ``ChatMessage`` objects before calling + ``Pipeline.run``: + + .. code-block:: python + + from haystack import coerce_pipeline_inputs + + inputs = coerce_pipeline_inputs(pipeline, serialized_inputs) + result = pipeline.run(inputs) diff --git a/test/core/test_serialization.py b/test/core/test_serialization.py index 8422ab2339..ac6e5446cf 100644 --- a/test/core/test_serialization.py +++ b/test/core/test_serialization.py @@ -4,15 +4,22 @@ import sys from copy import deepcopy +from dataclasses import dataclass from typing import Any from unittest.mock import Mock import pytest +from pydantic import BaseModel +from haystack import Document +from haystack.components.converters.txt import TextFileToDocument +from haystack.components.generators.chat import OpenAIChatGenerator +from haystack.components.joiners import BranchJoiner from haystack.core.component import component from haystack.core.errors import DeserializationError, SerializationError from haystack.core.pipeline import Pipeline from haystack.core.serialization import ( + coerce_pipeline_inputs, component_from_dict, component_to_dict, default_from_dict, @@ -20,6 +27,20 @@ generate_qualified_class_name, import_class_by_name, ) +from haystack.dataclasses import ( + ByteStream, + ChatMessage, + ChatRole, + ExtractedAnswer, + FileContent, + GeneratedAnswer, + ImageContent, + ReasoningContent, + SparseEmbedding, + TextContent, + ToolCall, + ToolCallResult, +) from haystack.document_stores.in_memory import InMemoryDocumentStore from haystack.testing import factory from haystack.utils import ComponentDevice, Secret @@ -644,3 +665,271 @@ def test_default_from_dict_rejects_unknown_nested_parameter(): # 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 + + +@component +class TypedEcho: + """Echoes back a single input whose socket type is set at construction time.""" + + def __init__(self, type_: Any) -> None: + component.set_input_types(self, value=type_) + component.set_output_types(self, value=Any) + + def run(self, **kwargs: Any) -> dict[str, Any]: + return {"value": kwargs.get("value")} + + +@component +class MultiInputEcho: + """A component with one coercible socket alongside plain, non-coercible sockets.""" + + @component.output_types(messages=list[ChatMessage]) + def run(self, messages: list[ChatMessage], top_k: int = 3, config: dict[str, Any] | None = None) -> dict[str, Any]: + return {"messages": messages} + + +@dataclass +class PlainDataclass: + name: str + count: int + + +class PydanticModel(BaseModel): + a: int + b: str + + +class Conversation(BaseModel): + messages: list[ChatMessage] + documents: list[Document] + + +class SerializationEnvelope(BaseModel): + value: Any + + +def serialize_with_haystack(instance: Any) -> Any: + """Serialize the way a client using Haystack's own methods does: via `to_dict`.""" + return instance.to_dict() + + +def serialize_with_pydantic(instance: Any) -> Any: + """ + Serialize the way a FastAPI app does: instances arrive inside a loosely-typed Pydantic field, so Pydantic + (not `to_dict`) produces the JSON payload. + """ + return SerializationEnvelope(value=instance).model_dump(mode="json")["value"] + + +TOOL_CALL = ToolCall(tool_name="t", arguments={"a": 1}, id="1") + +# A fully loaded ByteStream. It round-trips through `to_dict` (a list of ints) but not through Pydantic, which +# serializes the bytes field as a JSON string that ByteStream.from_dict cannot rebuild. +BYTE_STREAM = ByteStream(data=b"hello", mime_type="text/plain", meta={"k": "v"}) + +# Instances serializable by both a Haystack `to_dict` and a Pydantic `model_dump`. Fully loaded where possible: +# a Document with all non-bytes fields and a ChatMessage with every content-part type. +SHARED_ROUND_TRIP_INSTANCES = [ + pytest.param( + Document( + content="Hello", + meta={"k": "v"}, + score=0.5, + embedding=[0.1, 0.2], + sparse_embedding=SparseEmbedding(indices=[0, 2], values=[0.1, 0.2]), + ), + id="Document", + ), + pytest.param( + ChatMessage( + _role=ChatRole.ASSISTANT, + _content=[ + TextContent(text="hi"), + TOOL_CALL, + ToolCallResult(result="r", origin=TOOL_CALL, error=False), + ImageContent(base64_image="aGVsbG8=", mime_type="image/png"), + ReasoningContent(reasoning_text="because"), + FileContent(base64_data="aGVsbG8=", mime_type="application/pdf"), + ], + _name="n", + _meta={"k": "v"}, + ), + id="ChatMessage", + ), + pytest.param(GeneratedAnswer(data="answer", query="q", documents=[Document(content="d")]), id="GeneratedAnswer"), + pytest.param(ImageContent(base64_image="aGVsbG8=", mime_type="image/png"), id="ImageContent"), +] + + +def pytest_generate_tests(metafunc): + # Drives the round-trip tests off each subclass's `instances` attribute, so the two serialization methods + # can cover different sets of instances while sharing the same test bodies. + if "coercible_instance" in metafunc.fixturenames and metafunc.cls is not None: + metafunc.parametrize("coercible_instance", metafunc.cls.instances) + + +class CoercePipelineInputsTests: + """ + Shared `coerce_pipeline_inputs` tests. Subclasses set `serialize` (the client-side serialization method) and + `instances` (the round-trip cases valid for that method). + """ + + @staticmethod + def serialize(instance: Any) -> Any: + raise NotImplementedError + + instances: list[Any] = [] + + @pytest.fixture + def pipeline(self): + pipe = Pipeline() + pipe.add_component("echo", MultiInputEcho()) + return pipe + + def test_single_socket(self, coercible_instance): + pipe = Pipeline() + pipe.add_component("echo", TypedEcho(type(coercible_instance))) + coerced = coerce_pipeline_inputs(pipe, {"echo": {"value": self.serialize(coercible_instance)}}) + assert coerced["echo"]["value"] == coercible_instance + + def test_list_socket(self, coercible_instance): + pipe = Pipeline() + pipe.add_component("echo", TypedEcho(list[type(coercible_instance)])) # type: ignore[misc] + coerced = coerce_pipeline_inputs(pipe, {"echo": {"value": [self.serialize(coercible_instance)]}}) + assert coerced["echo"]["value"] == [coercible_instance] + + def test_optional_socket(self, coercible_instance): + pipe = Pipeline() + pipe.add_component("echo", TypedEcho(type(coercible_instance) | None)) + coerced = coerce_pipeline_inputs(pipe, {"echo": {"value": self.serialize(coercible_instance)}}) + assert coerced["echo"]["value"] == coercible_instance + + def test_nested_format(self, pipeline): + messages = [ChatMessage.from_user("Hi"), ChatMessage.from_assistant("Hello")] + data: dict[str, Any] = {"echo": {"messages": [self.serialize(message) for message in messages], "top_k": 5}} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced == {"echo": {"messages": messages, "top_k": 5}} + # the input data is not mutated + assert isinstance(data["echo"]["messages"][0], dict) + + def test_flat_format(self, pipeline): + messages = [ChatMessage.from_user("Hi")] + data = {"messages": [self.serialize(message) for message in messages], "top_k": 5} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced == {"messages": messages, "top_k": 5} + + def test_instances_pass_through(self, pipeline): + messages = [ChatMessage.from_user("Hi")] + data = {"echo": {"messages": messages}} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced["echo"]["messages"][0] is messages[0] + + def test_mixed_list(self, pipeline): + message = ChatMessage.from_user("Hi") + other = ChatMessage.from_assistant("Hello") + data = {"echo": {"messages": [message, self.serialize(other)]}} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced["echo"]["messages"] == [message, other] + + def test_non_coercible_values_untouched(self, pipeline): + data = {"echo": {"messages": [], "top_k": 5, "config": {"a": 1}}} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced == data + + def test_unknown_inputs_untouched(self, pipeline): + message = self.serialize(ChatMessage.from_user("Hi")) + data = {"unknown_input": [message], "top_k": 5} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced == {"unknown_input": [message], "top_k": 5} + + def test_unknown_component_untouched(self, pipeline): + message = self.serialize(ChatMessage.from_user("Hi")) + data = {"unknown_component": {"messages": [message]}} + coerced = coerce_pipeline_inputs(pipeline, data) + assert coerced == {"unknown_component": {"messages": [message]}} + + def test_variadic_socket(self): + pipe = Pipeline() + pipe.add_component("joiner", BranchJoiner(list[ChatMessage])) + messages = [ChatMessage.from_user("Hi")] + data = {"value": [self.serialize(message) for message in messages], "unrelated": 1} + coerced = coerce_pipeline_inputs(pipe, data) + assert coerced["value"] == messages + + def test_chat_generator_messages_socket(self): + # OpenAIChatGenerator's `messages` socket is `list[ChatMessage] | str`: the list of serialized messages + # is coerced, since ChatMessage is the only coercible member of the union. + pipe = Pipeline() + pipe.add_component("generator", OpenAIChatGenerator(api_key=Secret.from_token("test-key"))) + messages = [ChatMessage.from_user("Hi"), ChatMessage.from_system("Sys")] + data = {"generator": {"messages": [self.serialize(message) for message in messages]}} + coerced = coerce_pipeline_inputs(pipe, data) + assert coerced == {"generator": {"messages": messages}} + + def test_chat_generator_messages_socket_string_passes_through(self): + # The `str` arm of `list[ChatMessage] | str` is not coercible, so a bare string is left untouched. + pipe = Pipeline() + pipe.add_component("generator", OpenAIChatGenerator(api_key=Secret.from_token("test-key"))) + data = {"generator": {"messages": "just a string"}} + coerced = coerce_pipeline_inputs(pipe, data) + assert coerced == {"generator": {"messages": "just a string"}} + + def test_converter_sources_socket(self): + # A converter's `sources` socket is `list[str | Path | ByteStream]`: ByteStream is the only coercible + # member, so a serialized ByteStream is coerced while plain string paths pass through. ByteStream is + # serialized with `to_dict` regardless of the client, since it has no lossless Pydantic JSON form. + pipe = Pipeline() + pipe.add_component("converter", TextFileToDocument()) + byte_stream = ByteStream(data=b"hello", mime_type="text/plain") + data = {"converter": {"sources": ["/path/to/file.txt", byte_stream.to_dict()]}} + coerced = coerce_pipeline_inputs(pipe, data) + assert coerced["converter"]["sources"] == ["/path/to/file.txt", byte_stream] + + def test_ambiguous_union_socket_raises(self): + # A socket that involves more than one coercible class cannot be disambiguated from the payload. + pipe = Pipeline() + pipe.add_component("echo", TypedEcho(GeneratedAnswer | ExtractedAnswer)) + answer = GeneratedAnswer(data="a", query="q", documents=[]) + data = {"echo": {"value": self.serialize(answer)}} + with pytest.raises(DeserializationError, match="multiple deserializable members"): + coerce_pipeline_inputs(pipe, data) + + def test_ambiguous_union_socket_passes_through_instances(self): + # An already-deserialized value needs no coercion, so an ambiguous socket does not raise for it. + pipe = Pipeline() + pipe.add_component("echo", TypedEcho(GeneratedAnswer | ExtractedAnswer)) + answer = GeneratedAnswer(data="a", query="q", documents=[]) + coerced = coerce_pipeline_inputs(pipe, {"echo": {"value": answer}}) + assert coerced["echo"]["value"] is answer + + +class TestCoercePipelineInputsHaystackSerialization(CoercePipelineInputsTests): + """Inputs serialized with Haystack's own `to_dict`.""" + + serialize = staticmethod(serialize_with_haystack) + instances = [*SHARED_ROUND_TRIP_INSTANCES, pytest.param(BYTE_STREAM, id="ByteStream")] + + +class TestCoercePipelineInputsPydanticSerialization(CoercePipelineInputsTests): + """Inputs serialized with a Pydantic `model_dump`, as a FastAPI app does.""" + + serialize = staticmethod(serialize_with_pydantic) + instances = [ + *SHARED_ROUND_TRIP_INSTANCES, + pytest.param( + BYTE_STREAM, + id="ByteStream", + marks=pytest.mark.xfail( + reason="Pydantic serializes the bytes field as a JSON string that ByteStream.from_dict cannot rebuild", + raises=TypeError, + strict=True, + ), + ), + # These types have no `to_dict`, so they only arise under Pydantic serialization. + pytest.param(PydanticModel(a=1, b="x"), id="pydantic-model"), + pytest.param(PlainDataclass(name="x", count=2), id="stdlib-dataclass"), + pytest.param( + Conversation(messages=[ChatMessage.from_user("Hi")], documents=[Document(content="D")]), + id="pydantic-model-nesting-haystack", + ), + ]