From 0ccf07af3335df40009ba62fb77bc6059057a9ab Mon Sep 17 00:00:00 2001 From: ZanSara Date: Wed, 21 Feb 2024 15:00:18 +0100 Subject: [PATCH] chore: Update Amazon Bedrock integration to use new generic callable (de)serializers for their callback handlers (#452) * serialize the callback * fix test * lint --- .../generators/amazon_bedrock/chat/chat_generator.py | 4 ++-- integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py index 6ce671e68..2af99e685 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py @@ -6,7 +6,7 @@ import boto3 from botocore.exceptions import BotoCoreError, ClientError from haystack import component, default_from_dict, default_to_dict -from haystack.components.generators.utils import deserialize_callback_handler +from haystack.components.generators.utils import deserialize_callback_handler, serialize_callback_handler from haystack.dataclasses import ChatMessage, StreamingChunk from haystack.utils.auth import Secret, deserialize_secrets_inplace @@ -248,7 +248,7 @@ def to_dict(self) -> Dict[str, Any]: model=self.model, stop_words=self.stop_words, generation_kwargs=self.model_adapter.generation_kwargs, - streaming_callback=self.streaming_callback, + streaming_callback=serialize_callback_handler(self.streaming_callback), ) @classmethod diff --git a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py index 574aab5cc..86a610811 100644 --- a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py +++ b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py @@ -58,7 +58,7 @@ def test_to_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): "model": "anthropic.claude-v2", "generation_kwargs": {"temperature": 0.7}, "stop_words": [], - "streaming_callback": print_streaming_chunk, + "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk", }, }