diff --git a/integrations/cohere/tests/test_cohere_chat_generator.py b/integrations/cohere/tests/test_cohere_chat_generator.py index c91ada419..edefc1a43 100644 --- a/integrations/cohere/tests/test_cohere_chat_generator.py +++ b/integrations/cohere/tests/test_cohere_chat_generator.py @@ -260,7 +260,10 @@ def test_live_run(self): @pytest.mark.integration def test_live_run_wrong_model(self, chat_messages): component = CohereChatGenerator(model="something-obviously-wrong", api_key=os.environ.get("COHERE_API_KEY")) - with pytest.raises(cohere.CohereAPIError, match="finetuned model something-obviously-wrong is not valid"): + with pytest.raises( + cohere.CohereAPIError, + match="model not found, make sure the correct model ID was used and that you have access to the model.", + ): component.run(chat_messages) @pytest.mark.skipif( diff --git a/integrations/cohere/tests/test_cohere_generators.py b/integrations/cohere/tests/test_cohere_generators.py index e2ce10405..90d4d3e28 100644 --- a/integrations/cohere/tests/test_cohere_generators.py +++ b/integrations/cohere/tests/test_cohere_generators.py @@ -164,7 +164,7 @@ def __init__(self): self.responses = "" def __call__(self, chunk): - self.responses += chunk.text + self.responses += chunk.content return chunk callback = Callback()