From 126830f08a29f9062c01c4358b993dc79db84cb1 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 4 Sep 2023 11:57:20 -0700
Subject: [PATCH] allow users to set AWS_REGION_NAME

---
 litellm/llms/sagemaker.py        | 14 ++++++++++----
 litellm/tests/test_completion.py |  1 +
 2 files changed, 11 insertions(+), 4 deletions(-)

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'] = <your-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():