From 0e044a88aa9de14ee992f10aaa0a722d0f603c8a Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Wed, 14 Feb 2024 14:57:17 +0100 Subject: [PATCH] Reintroduce serialize_callback_handler and deserialize_callback_handler (#6988) --- haystack/components/generators/utils.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/haystack/components/generators/utils.py b/haystack/components/generators/utils.py index 43c104286..cab8e8ffe 100644 --- a/haystack/components/generators/utils.py +++ b/haystack/components/generators/utils.py @@ -1,4 +1,7 @@ +from typing import Optional, Callable + from haystack.dataclasses import StreamingChunk +from haystack.utils.callable_serialization import serialize_callable, deserialize_callable def print_streaming_chunk(chunk: StreamingChunk) -> None: @@ -7,3 +10,22 @@ def print_streaming_chunk(chunk: StreamingChunk) -> None: Prints the tokens of the first completion to stdout as soon as they are received """ print(chunk.content, flush=True, end="") + + +def serialize_callback_handler(streaming_callback: Callable[[StreamingChunk], None]) -> str: + """ + Serializes the streaming callback handler. + :param streaming_callback: The streaming callback handler function + :return: The full path of the streaming callback handler function + """ + return serialize_callable(streaming_callback) + + +def deserialize_callback_handler(callback_name: str) -> Optional[Callable[[StreamingChunk], None]]: + """ + Deserializes the streaming callback handler. + :param callback_name: The full path of the streaming callback handler function + :return: The streaming callback handler function + :raises DeserializationError: If the streaming callback handler function cannot be found + """ + return deserialize_callable(callback_name)