remove stream arg

deepset-ai · Sep 28, 2024 · 839adda · 839adda
1 parent 907c10b
commit 839adda
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/...ock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/...ock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py
@@ -173,7 +173,7 @@ def run(
         generation_kwargs = generation_kwargs.copy()
 
         streaming_callback = streaming_callback or self.streaming_callback
-        generation_kwargs["stream"] = streaming_callback is not None
+        is_streaming_enabled = streaming_callback is not None
 
         # check if the prompt is a list of ChatMessage objects
         if not (
@@ -188,7 +188,7 @@ def run(
             messages=messages, **{"stop_words": self.stop_words, **generation_kwargs}
         )
         try:
-            if streaming_callback:
+            if is_streaming_enabled:
                 response = self.client.invoke_model_with_response_stream(
                     body=json.dumps(body), modelId=self.model, accept="application/json", contentType="application/json"
                 )

diff --git a/integrations/amazon_bedrock/tests/test_chat_generator.py b/integrations/amazon_bedrock/tests/test_chat_generator.py
@@ -246,7 +246,7 @@ def test_long_prompt_is_not_truncated_when_truncate_false(mock_boto3_session):
         mock_ensure_token_limit.assert_not_called(),
 
         # Check the prompt passed to prepare_body
-        generator.model_adapter.prepare_body.assert_called_with(messages=messages, stop_words=[], stream=False)
+        generator.model_adapter.prepare_body.assert_called_with(messages=messages, stop_words=[])
 
 
 @pytest.mark.parametrize(