From a32639fa7958c10263905c334ecc62c34bb475b8 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 31 Jan 2024 19:09:54 -0800
Subject: [PATCH] fix(utils.py): support max token adjustment for sagemaker

---
 litellm/tests/test_completion_cost.py        |  2 +-
 litellm/tests/test_model_max_token_adjust.py | 28 ++++++++++++++++++++
 litellm/utils.py                             | 21 ++++++++-------
 3 files changed, 41 insertions(+), 10 deletions(-)
 create mode 100644 litellm/tests/test_model_max_token_adjust.py

diff --git a/litellm/tests/test_completion_cost.py b/litellm/tests/test_completion_cost.py
index b117223ab08..b55f9c9d6e0 100644
--- a/litellm/tests/test_completion_cost.py
+++ b/litellm/tests/test_completion_cost.py
@@ -13,7 +13,7 @@ import pytest
 def test_get_gpt3_tokens():
     max_tokens = get_max_tokens("gpt-3.5-turbo")
     print(max_tokens)
-    assert max_tokens == 4097
+    assert max_tokens == 4096
     # print(results)
 
 
diff --git a/litellm/tests/test_model_max_token_adjust.py b/litellm/tests/test_model_max_token_adjust.py
new file mode 100644
index 00000000000..b026b625677
--- /dev/null
+++ b/litellm/tests/test_model_max_token_adjust.py
@@ -0,0 +1,28 @@
+# What this tests?
+## Tests if max tokens get adjusted, if over limit
+
+import sys, os, time
+import traceback, asyncio
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import litellm
+from litellm import completion
+
+litellm.drop_params = True
+
+
+def test_completion_sagemaker():
+    response = completion(
+        model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
+        messages=[{"content": "Hello, how are you?", "role": "user"}],
+        temperature=0.2,
+        max_tokens=80000,
+        hf_model_name="meta-llama/Llama-2-70b-chat-hf",
+    )
+    print(f"response: {response}")
+
+
+# test_completion_sagemaker()
diff --git a/litellm/utils.py b/litellm/utils.py
index c9fccdc6c0e..43319246f12 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -2199,22 +2199,28 @@ def client(original_function):
                 )
             ):
                 try:
-                    max_output_tokens = get_max_tokens(model=model)
+                    base_model = model
+                    if kwargs.get("hf_model_name", None) is not None:
+                        base_model = f"huggingface/{kwargs.get('hf_model_name')}"
+                    max_output_tokens = (
+                        get_max_tokens(model=base_model) or 4096
+                    )  # assume min context window is 4k tokens
                     user_max_tokens = kwargs.get("max_tokens")
-                    ## Scenario 1: User limit > model limit
-                    if user_max_tokens > max_output_tokens:
-                        user_max_tokens = max_output_tokens
-                    ## Scenario 2: User limit + prompt > model limit
+                    ## Scenario 1: User limit + prompt > model limit
                     messages = None
                     if len(args) > 1:
                         messages = args[1]
                     elif kwargs.get("messages", None):
                         messages = kwargs["messages"]
-                    input_tokens = token_counter(model=model, messages=messages)
+                    input_tokens = token_counter(model=base_model, messages=messages)
+                    input_tokens += max(
+                        0.1 * input_tokens, 10
+                    )  # give at least a 10 token buffer. token counting can be imprecise.
                     if input_tokens > max_output_tokens:
                         pass  # allow call to fail normally
                     elif user_max_tokens + input_tokens > max_output_tokens:
                         user_max_tokens = max_output_tokens - input_tokens
+                    kwargs["max_tokens"] = user_max_tokens
                 except Exception as e:
                     print_verbose(f"Error while checking max token limit: {str(e)}")
             # MODEL CALL
@@ -4553,7 +4559,6 @@ def get_max_tokens(model: str):
     def _get_max_position_embeddings(model_name):
         # Construct the URL for the config.json file
         config_url = f"https://huggingface.co/{model_name}/raw/main/config.json"
-
         try:
             # Make the HTTP request to get the raw JSON file
             response = requests.get(config_url)
@@ -4561,10 +4566,8 @@ def get_max_tokens(model: str):
 
             # Parse the JSON response
             config_json = response.json()
-
             # Extract and return the max_position_embeddings
             max_position_embeddings = config_json.get("max_position_embeddings")
-
             if max_position_embeddings is not None:
                 return max_position_embeddings
             else: