From a8644d8a7d67fddc6c315050eb8d9e1944bbfce7 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 12 Aug 2024 16:28:22 -0700
Subject: [PATCH] fix(utils.py): if openai model, don't check hf tokenizers

---
 litellm/tests/test_token_counter.py | 13 +++++++++++++
 litellm/utils.py                    |  9 ++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/litellm/tests/test_token_counter.py b/litellm/tests/test_token_counter.py
index 6bd001fcc8..b6dca32c64 100644
--- a/litellm/tests/test_token_counter.py
+++ b/litellm/tests/test_token_counter.py
@@ -11,7 +11,9 @@ import pytest
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
+from unittest.mock import AsyncMock, MagicMock, patch
 
+import litellm
 from litellm import (
     create_pretrained_tokenizer,
     decode,
@@ -343,3 +345,14 @@ def test_empty_tools():
     )
 
     print(result)
+
+
+def test_gpt_4o_token_counter():
+    with patch.object(
+        litellm.utils, "openai_token_counter", new=MagicMock()
+    ) as mock_client:
+        token_counter(
+            model="gpt-4o-2024-05-13", messages=[{"role": "user", "content": "Hey!"}]
+        )
+
+        mock_client.assert_called()
diff --git a/litellm/utils.py b/litellm/utils.py
index 5c6a4406b3..ef77b5a69a 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1610,10 +1610,17 @@ def _select_tokenizer(model: str):
     # default - tiktoken
     else:
         tokenizer = None
+        if (
+            model in litellm.open_ai_chat_completion_models
+            or model in litellm.open_ai_text_completion_models
+            or model in litellm.open_ai_embedding_models
+        ):
+            return {"type": "openai_tokenizer", "tokenizer": encoding}
+
         try:
             tokenizer = Tokenizer.from_pretrained(model)
             return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
-        except:
+        except Exception:
             return {"type": "openai_tokenizer", "tokenizer": encoding}