Skip to content

Commit

Permalink
chore: Update Amazon Bedrock integration to use new generic callable …
Browse files Browse the repository at this point in the history
…(de)serializers for their callback handlers (#452)

* serialize the callback

* fix test

* lint
  • Loading branch information
ZanSara authored Feb 21, 2024
1 parent af346c5 commit 0ccf07a
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
},
}

Expand Down

0 comments on commit 0ccf07a

Please sign in to comment.