Create a ChatGenerator Protocol for Agent serde support

[sjrl]

While working on adding the Agent abstraction we once again ran into how to best pass a ChatGenerator as an init variable. @anakin87, @mathislucka and I came to the conclusion that directly passing the ChatGenerator component at init time would be the most user friendly as long as we can properly provide serde support. To this effect @anakin87 suggested we create a ChatGenerator Protocol that would allow us to more straightforwardly serialize and deserialize the ChatGenerator component passed to the Agent at init time. You can see our original discussion here.

So the request is to create a ChatGenerator protocol and use it in the Agent. Additionally, this protocol could be used elsewhere such as in the LLMMetadataExtractor, and LLMEvaluator (see issue here)

[mathislucka]

@sjrl @anakin87

I started working on this but I realized that it is not that simple.

With the ChatGenerators in Agent we mostly care about the run-method signature. We want it to accept messages, tools and a streaming_callback (needed for streaming support in the Agent) but we don't really care about other args.

First I thought we could do something like this:

from typing import Any, Dict, List, Optional, Protocol, Union, runtime_checkable, TypeVar

from haystack.dataclasses import ChatMessage
from haystack_experimental.tools import Tool
from haystack.tools import Tool as HaystackTool
from haystack_experimental.dataclasses.streaming_chunk import StreamingCallbackT


ChatGeneratorT = TypeVar('ChatGeneratorT', bound='ChatGenerator')

@runtime_checkable
class ChatGenerator(Protocol):
    # Note: When implementing this protocol as a Haystack component,
    # use the @component decorator from haystack
    """
    Protocol defining the interface for chat generation components.

    Implementations of this protocol should provide functionality to:
    1. Generate chat responses based on input messages
    2. Support streaming responses via callbacks
    3. Integrate with tools for function calling
    4. Serialize to and deserialize from dictionaries

    This protocol is designed to be compatible with Haystack's Component protocol.
    Classes implementing this protocol can be decorated with @component to make
    them usable within Haystack pipelines.
    """

    def run(
            self,
            messages: List[ChatMessage],
            streaming_callback: Optional[StreamingCallbackT] = None,
            tools: Optional[List[Tool]] = None,
            **kwargs
    ) -> Dict[str, Any]:
        """
        Invokes chat completion based on the provided messages and generation parameters.

        :param messages:
            A list of ChatMessage instances representing the input messages.
        :param streaming_callback:
            A callback function that is called when a new token is received from the stream.
            Cannot be a coroutine.
        :param tools:
            A list of tools for which the model can prepare calls.
        :param kwargs:
            Additional keyword arguments for text generation. These parameters will
            override the parameters passed during component initialization.

        :returns:
            A dictionary containing the generated responses as a list of ChatMessage objects.
        """
        ...

    async def run_async(
            self,
            messages: List[ChatMessage],
            streaming_callback: Optional[StreamingCallbackT] = None,
            tools: Optional[List[Tool]] = None,
            **kwargs
    ):
        """
        Asynchronously invokes chat completion based on the provided messages and generation parameters.

        :param messages:
            A list of ChatMessage instances representing the input messages.
        :param streaming_callback:
            A callback function that is called when a new token is received from the stream.
            Must be a coroutine.
        :param tools:
            A list of tools for which the model can prepare calls.
        :param kwargs:
            Additional keyword arguments for text generation. These parameters will
            override the parameters passed during component initialization.

        :returns:
            A dictionary containing the generated responses as a list of ChatMessage objects.
        """
        ...

    def to_dict(self) -> Dict[str, Any]:
        """
        Serialize this component to a dictionary.

        :returns:
            The serialized component as a dictionary.
        """
        ...

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> ChatGeneratorT:
        """
        Deserialize this component from a dictionary.

        :param data:
            The dictionary representation of this component.
        :returns:
            The deserialized component instance.
        """
        ...

Normally, you would enforce this protocol at runtime by doing isinstance(chat_generator, ChatGenerator). However, as per the comment here: https://github.com/python/cpython/blob/052cb717f5f97d08d2074f4118fd2c21224d3015/Lib/typing.py#L2066

The isinstance check does not actually check the function signature. Additionally, not all signatures of all ChatGenerators match 100% (legitimately so).

We could:

• add a @chat_generator decorator that similarly to the component decorator can check for the presence of some parameters in the signature
• write a small utility function that checks for a compatible signature in the appropriate places at runtime (e.g. for Agent) and skip the decorator (not sure what use the protocol would be then since the ChatGenerator protocol would be basically the same as the component protocol)

Any ideas?

[mathislucka]

Right now, I'm leaning towards a runtime check in Agent for tools, streaming_callback, and messages and not adding any protocol.

The type hint on Agent would then be:

class Agent:
  def __init__(chat_generator: Component, ...)

Serialization should work fine using the existing component_to_dict utility function from haystack as well as the _load_component helper that @tstadel implemented for the ParallelExecutor in custom nodes.

[mathislucka]

To be more explicit, we could define an additional ChatGenerator protocol but it would be the same as Component protocol. The only benefit would be that it might be more explicit for the user when looking at the type hints:

class Agent:
  def __init__(chat_generator: ChatGenerator, ...)