diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py index 09f22bf52..e518c3c6d 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/errors.py @@ -1,9 +1,9 @@ - class SagemakerError(Exception): """ Parent class for all exceptions raised by the Sagemaker component """ + class AWSConfigurationError(SagemakerError): """Exception raised when AWS is not configured correctly""" diff --git a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py index b64902f00..3cf71abe4 100644 --- a/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py +++ b/integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py @@ -44,6 +44,7 @@ class SagemakerGenerator: >>> and respond to natural human language in a way that is both meaningful and useful.'], 'meta': [{}]} ``` """ + model_generation_keys: ClassVar = ["generated_text", "generation"] def __init__(