diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index b1eed4e997..5d02247f29 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -4,7 +4,7 @@ from enum import Enum import requests import time from typing import Callable -from litellm.utils import ModelResponse +from litellm.utils import ModelResponse, get_secret import sys class SagemakerError(Exception): @@ -21,6 +21,8 @@ os.environ['AWS_ACCESS_KEY_ID'] = "" os.environ['AWS_SECRET_ACCESS_KEY'] = "" """ +# set os.environ['AWS_REGION_NAME'] = + def completion( model: str, messages: list, @@ -35,12 +37,16 @@ def completion( import sys if 'boto3' not in sys.modules: import boto3 - + + region_name = ( + get_secret("AWS_REGION_NAME") or + "us-west-2" # default to us-west-2 + ) + client = boto3.client( "sagemaker-runtime", - region_name="us-west-2" + region_name=region_name ) - model = model prompt = "" diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index e582eee59c..c0ed591d9e 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -404,6 +404,7 @@ def test_completion_sagemaker(): print(response) except Exception as e: pytest.fail(f"Error occurred: {e}") +test_completion_sagemaker() # def test_vertex_ai():