From d56f47285edb9b993e630969c3c7c6c2a760c432 Mon Sep 17 00:00:00 2001 From: Vladimir Blagojevic Date: Thu, 7 Mar 2024 14:51:55 +0100 Subject: [PATCH] Add first live integration test --- .../tests/test_chat_generator.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/integrations/amazon_bedrock/tests/test_chat_generator.py b/integrations/amazon_bedrock/tests/test_chat_generator.py index 196ddb36c..04e3cebd9 100644 --- a/integrations/amazon_bedrock/tests/test_chat_generator.py +++ b/integrations/amazon_bedrock/tests/test_chat_generator.py @@ -2,7 +2,7 @@ import pytest from haystack.components.generators.utils import print_streaming_chunk -from haystack.dataclasses import ChatMessage +from haystack.dataclasses import ChatMessage, ChatRole from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator from haystack_integrations.components.generators.amazon_bedrock.chat.adapters import ( @@ -228,3 +228,25 @@ def test_get_responses(self) -> None: assert isinstance(message, ChatMessage) assert response_message == [ChatMessage.from_assistant(expected_response)] + + @pytest.mark.parametrize("model_name", [ + "anthropic.claude-3-sonnet-20240229-v1:0", + "anthropic.claude-v2:1", + "meta.llama2-13b-chat-v1" + ]) + @pytest.mark.integration + def test_default_inference_params(self, model_name): + messages = [ + ChatMessage.from_system("\\nYou are a helpful assistant, be super brief in your responses."), + ChatMessage.from_user("What's the capital of France?"), + ] + + client = AmazonBedrockChatGenerator(model=model_name) + response = client.run(messages) + assert response["replies"] + assert isinstance(response["replies"], list) + assert len(response["replies"]) > 0 + assert isinstance(response["replies"][0], ChatMessage) + assert response["replies"][0].content + assert ChatMessage.is_from(response["replies"][0], ChatRole.ASSISTANT) + assert "paris" in response["replies"][0].content.lower() \ No newline at end of file