From 4a4ee51df399fe96e2d18055530a701edbb70467 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 4 Sep 2023 11:30:34 -0700
Subject: [PATCH] working sagemaker support

---
 litellm/__init__.py              |  1 +
 litellm/llms/sagemaker.py        | 39 ++++++++++++++++++++++++--------
 litellm/main.py                  |  1 -
 litellm/tests/test_completion.py | 12 ++++++++++
 4 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index be21654bda..b5c84bab4c 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -265,6 +265,7 @@ provider_list = [
     "ai21",
     "baseten",
     "azure",
+    "sagemaker",
 ]
 
 models_by_provider = {
diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py
index 250383b2cf..8ddc35b669 100644
--- a/litellm/llms/sagemaker.py
+++ b/litellm/llms/sagemaker.py
@@ -5,6 +5,7 @@ import requests
 import time
 from typing import Callable
 from litellm.utils import ModelResponse
+import sys
 
 class SagemakerError(Exception):
     def __init__(self, status_code, message):
@@ -14,18 +15,32 @@ class SagemakerError(Exception):
             self.message
         )  # Call the base class constructor with the parameters it needs
 
+"""
+SAGEMAKER AUTH Keys/Vars
+os.environ['AWS_ACCESS_KEY_ID'] = ""
+os.environ['AWS_SECRET_ACCESS_KEY'] = ""
+"""
+
 def completion(
     model: str,
     messages: list,
     model_response: ModelResponse,
     print_verbose: Callable,
     encoding,
-    api_key,
     logging_obj,
     optional_params=None,
     litellm_params=None,
     logger_fn=None,
 ):
+    import sys
+    if 'boto3' not in sys.modules:
+        import boto3
+    
+    client = boto3.client(
+        "sagemaker-runtime", 
+        region_name="us-west-2"
+    )
+    
 
     model = model
     prompt = ""
@@ -42,7 +57,7 @@ def completion(
         else:
             prompt += f"{message['content']}"
     data = {
-        "prompt": prompt,
+        "inputs": prompt,
         # "instruction": prompt, # some baseten models require the prompt to be passed in via the 'instruction' kwarg
         **optional_params,
     }
@@ -50,26 +65,30 @@ def completion(
     ## LOGGING
     logging_obj.pre_call(
             input=prompt,
-            api_key=api_key,
+            api_key="",
             additional_args={"complete_input_dict": data},
         )
     ## COMPLETION CALL
-    response = requests.post(
-        "https://api.ai21.com/studio/v1/" + model + "/complete", headers=headers, data=json.dumps(data)
+    response = client.invoke_endpoint(
+        EndpointName=model,
+        ContentType="application/json",
+        Body=json.dumps(data),
+        CustomAttributes="accept_eula=true",
     )
+    response = response["Body"].read().decode("utf8")
     if "stream" in optional_params and optional_params["stream"] == True:
         return response.iter_lines()
     else:
         ## LOGGING
         logging_obj.post_call(
                 input=prompt,
-                api_key=api_key,
-                original_response=response.text,
+                api_key="",
+                original_response=response,
                 additional_args={"complete_input_dict": data},
             )
-        print_verbose(f"raw model_response: {response.text}")
+        print_verbose(f"raw model_response: {response}")
         ## RESPONSE OBJECT
-        completion_response = response.json()
+        completion_response = json.loads(response)
         if "error" in completion_response:
             raise SagemakerError(
                 message=completion_response["error"],
@@ -77,7 +96,7 @@ def completion(
             )
         else:
             try:
-                model_response["choices"][0]["message"]["content"] = completion_response["completions"][0]["data"]["text"]
+                model_response["choices"][0]["message"]["content"] = completion_response[0]["generation"]
             except:
                 raise SagemakerError(message=json.dumps(completion_response), status_code=response.status_code)
 
diff --git a/litellm/main.py b/litellm/main.py
index d1ae00fa0f..10cfca1062 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -692,7 +692,6 @@ def completion(
                 litellm_params=litellm_params,
                 logger_fn=logger_fn,
                 encoding=encoding,
-                api_key=ai21_key,
                 logging_obj=logging
             )
             
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 4a41aac2e9..1cd6e570a2 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -114,6 +114,7 @@ def test_completion_claude_stream():
         pytest.fail(f"Error occurred: {e}")
 
 
+
 # def test_completion_hf_api():
 #     try:
 #         user_message = "write some code to find the sum of two numbers"
@@ -391,6 +392,17 @@ def test_completion_together_ai():
         pytest.fail(f"Error occurred: {e}")
 
 
+# def test_completion_sagemaker():
+#     try:
+#         response = completion(
+#             model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", 
+#             messages=messages
+#         )
+#         # Add any assertions here to check the response
+#         print(response)
+#     except Exception as e:
+#         pytest.fail(f"Error occurred: {e}")
+
 # def test_vertex_ai():
 #     model_name = "chat-bison"
 #     try: