Reintroduce serialize_callback_handler and deserialize_callback_handler (#6988)

This commit is contained in:
Vladimir Blagojevic 2024-02-14 14:57:17 +01:00 committed by GitHub
parent 6a776e672f
commit 0e044a88aa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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)