deepset-ai--haystack/haystack/components/generators/chat/fallback.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import Any

from haystack import component, default_from_dict, default_to_dict, logging
from haystack.components.generators.chat.types import ChatGenerator
from haystack.components.generators.utils import _normalize_messages
from haystack.core.serialization import component_to_dict
from haystack.dataclasses import ChatMessage, StreamingCallbackT
from haystack.tools import ToolsType
from haystack.utils.async_utils import _execute_component_async
from haystack.utils.deserialization import deserialize_component_inplace

logger = logging.getLogger(__name__)


@component
class FallbackChatGenerator:
    """
    A chat generator wrapper that tries multiple chat generators sequentially.

    It forwards all parameters transparently to the underlying chat generators and returns the first successful result.
    Calls chat generators sequentially until one succeeds. Falls back on any exception raised by a generator.
    If all chat generators fail, it raises a RuntimeError with details.

    Timeout enforcement is fully delegated to the underlying chat generators. The fallback mechanism will only
    work correctly if the underlying chat generators implement proper timeout handling and raise exceptions
    when timeouts occur. For predictable latency guarantees, ensure your chat generators:
    - Support a `timeout` parameter in their initialization
    - Implement timeout as total wall-clock time (shared deadline for both streaming and non-streaming)
    - Raise timeout exceptions (e.g., TimeoutError, asyncio.TimeoutError, httpx.TimeoutException) when exceeded

    Note: Most well-implemented chat generators (OpenAI, Anthropic, Cohere, etc.) support timeout parameters
    with consistent semantics. For HTTP-based LLM providers, a single timeout value (e.g., `timeout=30`)
    typically applies to all connection phases: connection setup, read, write, and pool. For streaming
    responses, read timeout is the maximum gap between chunks. For non-streaming, it's the time limit for
    receiving the complete response.

    Fail over is automatically triggered when a generator raises any exception, including:
    - Timeout errors (if the generator implements and raises them)
    - Rate limit errors (429)
    - Authentication errors (401)
    - Context length errors (400)
    - Server errors (500+)
    - Any other exception
    """

    def __init__(self, chat_generators: list[ChatGenerator]) -> None:
        """
        Creates an instance of FallbackChatGenerator.

        :param chat_generators: A non-empty list of chat generator components to try in order.
        """
        if not chat_generators:
            msg = "'chat_generators' must be a non-empty list"
            raise ValueError(msg)

        self.chat_generators = list(chat_generators)

    def to_dict(self) -> dict[str, Any]:
        """Serialize the component, including nested chat generators."""
        return default_to_dict(
            self,
            chat_generators=[
                component_to_dict(gen, name=f"chat_generator_{idx}") for idx, gen in enumerate(self.chat_generators)
            ],
        )

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> FallbackChatGenerator:
        """Rebuild the component from a serialized representation, restoring nested chat generators."""
        # Reconstruct nested chat generators from their serialized dicts
        init_params = data.get("init_parameters", {})
        serialized = init_params.get("chat_generators") or []
        deserialized: list[Any] = []
        for g in serialized:
            # Use the generic component deserializer available in Haystack
            holder = {"component": g}
            deserialize_component_inplace(holder, key="component")
            deserialized.append(holder["component"])
        init_params["chat_generators"] = deserialized
        data["init_parameters"] = init_params
        return default_from_dict(cls, data)

    def warm_up(self) -> None:
        """Warm up all underlying chat generators."""
        for gen in self.chat_generators:
            if hasattr(gen, "warm_up"):
                gen.warm_up()

    async def warm_up_async(self) -> None:
        """Warm up all underlying chat generators on the serving event loop."""
        for gen in self.chat_generators:
            if hasattr(gen, "warm_up_async"):
                await gen.warm_up_async()
            elif hasattr(gen, "warm_up"):
                gen.warm_up()

    def close(self) -> None:
        """Release the underlying chat generators' resources."""
        for gen in self.chat_generators:
            if hasattr(gen, "close"):
                gen.close()

    async def close_async(self) -> None:
        """Release the underlying chat generators' async resources."""
        for gen in self.chat_generators:
            if hasattr(gen, "close_async"):
                await gen.close_async()
            elif hasattr(gen, "close"):
                gen.close()

    def _run_single_sync(
        self,
        gen: Any,
        messages: list[ChatMessage],
        generation_kwargs: dict[str, Any] | None,
        tools: ToolsType | None,
        streaming_callback: StreamingCallbackT | None,
    ) -> dict[str, Any]:
        return gen.run(
            messages=messages, generation_kwargs=generation_kwargs, tools=tools, streaming_callback=streaming_callback
        )

    async def _run_single_async(
        self,
        gen: Any,
        messages: list[ChatMessage],
        generation_kwargs: dict[str, Any] | None,
        tools: ToolsType | None,
        streaming_callback: StreamingCallbackT | None,
    ) -> dict[str, Any]:
        return await _execute_component_async(
            gen,
            messages=messages,
            generation_kwargs=generation_kwargs,
            tools=tools,
            streaming_callback=streaming_callback,
        )

    @component.output_types(replies=list[ChatMessage], meta=dict[str, Any])
    def run(
        self,
        messages: list[ChatMessage] | str,
        generation_kwargs: dict[str, Any] | None = None,
        tools: ToolsType | None = None,
        streaming_callback: StreamingCallbackT | None = None,
    ) -> dict[str, list[ChatMessage] | dict[str, Any]]:
        """
        Execute chat generators sequentially until one succeeds.

        :param messages: The conversation history as a list of ChatMessage instances.
        :param generation_kwargs: Optional parameters for the chat generator (e.g., temperature, max_tokens).
        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for function calling capabilities.
        :param streaming_callback: Optional callable for handling streaming responses.
        :returns: A dictionary with:
            - "replies": Generated ChatMessage instances from the first successful generator.
            - "meta": Execution metadata including successful_chat_generator_index, successful_chat_generator_class,
              total_attempts, failed_chat_generators, plus any metadata from the successful generator.
        :raises RuntimeError: If all chat generators fail.
        """
        self.warm_up()

        messages = _normalize_messages(messages)

        failed: list[str] = []
        last_error: BaseException | None = None

        for idx, gen in enumerate(self.chat_generators):
            gen_name = gen.__class__.__name__
            try:
                result = self._run_single_sync(gen, messages, generation_kwargs, tools, streaming_callback)
                replies = result.get("replies", [])
                meta = dict(result.get("meta", {}))
                meta.update(
                    {
                        "successful_chat_generator_index": idx,
                        "successful_chat_generator_class": gen_name,
                        "total_attempts": idx + 1,
                        "failed_chat_generators": failed,
                    }
                )
                return {"replies": replies, "meta": meta}
            except Exception as e:  # noqa: BLE001 - fallback logic should handle any exception
                logger.warning(
                    "ChatGenerator {chat_generator} failed with error: {error}", chat_generator=gen_name, error=e
                )
                failed.append(gen_name)
                last_error = e

        failed_names = ", ".join(failed)
        msg = (
            f"All {len(self.chat_generators)} chat generators failed. "
            f"Last error: {last_error}. Failed chat generators: [{failed_names}]"
        )
        raise RuntimeError(msg)

    @component.output_types(replies=list[ChatMessage], meta=dict[str, Any])
    async def run_async(
        self,
        messages: list[ChatMessage] | str,
        generation_kwargs: dict[str, Any] | None = None,
        tools: ToolsType | None = None,
        streaming_callback: StreamingCallbackT | None = None,
    ) -> dict[str, list[ChatMessage] | dict[str, Any]]:
        """
        Asynchronously execute chat generators sequentially until one succeeds.

        :param messages: The conversation history as a list of ChatMessage instances.
        :param generation_kwargs: Optional parameters for the chat generator (e.g., temperature, max_tokens).
        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for function calling capabilities.
        :param streaming_callback: Optional callable for handling streaming responses.
        :returns: A dictionary with:
            - "replies": Generated ChatMessage instances from the first successful generator.
            - "meta": Execution metadata including successful_chat_generator_index, successful_chat_generator_class,
              total_attempts, failed_chat_generators, plus any metadata from the successful generator.
        :raises RuntimeError: If all chat generators fail.
        """
        await self.warm_up_async()

        messages = _normalize_messages(messages)

        failed: list[str] = []
        last_error: BaseException | None = None

        for idx, gen in enumerate(self.chat_generators):
            gen_name = gen.__class__.__name__
            try:
                result = await self._run_single_async(gen, messages, generation_kwargs, tools, streaming_callback)
                replies = result.get("replies", [])
                meta = dict(result.get("meta", {}))
                meta.update(
                    {
                        "successful_chat_generator_index": idx,
                        "successful_chat_generator_class": gen_name,
                        "total_attempts": idx + 1,
                        "failed_chat_generators": failed,
                    }
                )
                return {"replies": replies, "meta": meta}
            except Exception as e:  # noqa: BLE001 - fallback logic should handle any exception
                logger.warning(
                    "ChatGenerator {chat_generator} failed with error: {error}", chat_generator=gen_name, error=e
                )
                failed.append(gen_name)
                last_error = e

        failed_names = ", ".join(failed)
        msg = (
            f"All {len(self.chat_generators)} chat generators failed. "
            f"Last error: {last_error}. Failed chat generators: [{failed_names}]"
        )
        raise RuntimeError(msg)