Fix/issue 10113 embeddings use non default tokenizer (#10629)

* fix(embeddings): use non default tokenizer when passing list of lists of tokens (int)

* feat(embeddings): allow for passthrough of list of lists of tokens to hosted_vllm models

* Revert "fix(embeddings): use non default tokenizer when passing list of lists of tokens (int)"

This reverts commit a48acd95f8.

* refactor(embeddings): use a list to verify if provider accept as input a list of tokens

* fix(embeddings): verify the model name before validating if provider accept a arrays of tokens as input

When passing a list of tokens as input, verify the provider of the model by going through the list of models (`llm_model_list`). First, it check for model name then get the provider and verify if it accept or not arrays of tokens. If yes, then pass, else decode.
Previously, it was verifying provider and model name at the same time resulting in decoding even if the current model checked was not the target one (looping onto `llm_model_list`)

* test(embedding): add unit test to bypass decode for some providers with input as array of tokens

Ref: https://github.com/BerriAI/litellm/issues/10113

litellm/constants.py

@@ -226,6 +226,12 @@ LITELLM_CHAT_PROVIDERS = [
     "nscale",
 ]
 
+LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [
+    "openai",
+    "azure",
+    "hosted_vllm"
+]
+
 
 OPENAI_CHAT_COMPLETION_PARAMS = [
     "functions",

litellm/proxy/proxy_server.py

@@ -28,6 +28,7 @@ from typing import (
 from litellm.constants import (
     DEFAULT_MAX_RECURSE_DEPTH,
     DEFAULT_SLACK_ALERTING_THRESHOLD,
+    LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS
 )
 from litellm.types.utils import (
     ModelResponse,
@@ -3844,23 +3845,26 @@ async def embeddings(  # noqa: PLR0915
             and isinstance(data["input"][0], list)
             and isinstance(data["input"][0][0], int)
         ):  # check if array of tokens passed in
-            # check if non-openai/azure model called - e.g. for langchain integration
+            # check if provider accept list of tokens as input - e.g. for langchain integration
             if llm_model_list is not None and data["model"] in router_model_names:
                 for m in llm_model_list:
-                    if m["model_name"] == data["model"] and (
-                        m["litellm_params"]["model"] in litellm.open_ai_embedding_models
-                        or m["litellm_params"]["model"].startswith("azure/")
-                    ):
-                        pass
-                    else:
-                        # non-openai/azure embedding model called with token input
-                        input_list = []
-                        for i in data["input"]:
-                            input_list.append(
-                                litellm.decode(model="gpt-3.5-turbo", tokens=i)
-                            )
-                        data["input"] = input_list
-                        break
+                    if m["model_name"] == data["model"]:
+                        if (m["litellm_params"]["model"] in litellm.open_ai_embedding_models
+                                or any(
+                                    m["litellm_params"]["model"].startswith(provider)
+                                    for provider in LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS
+                                )
+                        ):
+                            pass
+                        else:
+                            # non-openai/azure embedding model called with token input
+                            input_list = []
+                            for i in data["input"]:
+                                input_list.append(
+                                    litellm.decode(model="gpt-3.5-turbo", tokens=i)
+                                )
+                            data["input"] = input_list
+                            break
 
         ### CALL HOOKS ### - modify incoming data / reject request before calling the model
         data = await proxy_logging_obj.pre_call_hook(

tests/litellm/proxy/test_configs/test_config_no_auth.yaml

@@ -0,0 +1,6 @@
+model_list:
+- litellm_params:
+    model: hosted_vllm/embed_model
+  model_info:
+    description: this is a test embedding hosted_vllm model
+  model_name: vllm_embed_model

tests/litellm/proxy/test_proxy_server.py

@@ -1,9 +1,11 @@
+import asyncio
 import importlib
 import json
 import os
 import socket
 import subprocess
 import sys
+from unittest import mock
 from unittest.mock import AsyncMock, MagicMock, mock_open, patch
 
 import click
@@ -18,6 +20,51 @@ sys.path.insert(
 )  # Adds the parent directory to the system-path
 
 import litellm
+from litellm.proxy.proxy_server import app, initialize
+
+example_embedding_result = {
+    "object": "list",
+    "data": [
+        {
+            "object": "embedding",
+            "index": 0,
+            "embedding": [
+                -0.006929283495992422,
+                -0.005336422007530928,
+                -4.547132266452536e-05,
+                -0.024047505110502243,
+                -0.006929283495992422,
+                -0.005336422007530928,
+                -4.547132266452536e-05,
+                -0.024047505110502243,
+                -0.006929283495992422,
+                -0.005336422007530928,
+                -4.547132266452536e-05,
+                -0.024047505110502243,
+            ],
+        }
+    ],
+    "model": "text-embedding-3-small",
+    "usage": {"prompt_tokens": 5, "total_tokens": 5},
+}
+
+def mock_patch_aembedding():
+    return mock.patch(
+        "litellm.proxy.proxy_server.llm_router.aembedding",
+        return_value=example_embedding_result,
+    )
+
+@pytest.fixture(scope="function")
+def client_no_auth():
+    # Assuming litellm.proxy.proxy_server is an object
+    from litellm.proxy.proxy_server import cleanup_router_config_variables
+
+    cleanup_router_config_variables()
+    filepath = os.path.dirname(os.path.abspath(__file__))
+    config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
+    # initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
+    asyncio.run(initialize(config=config_fp, debug=True))
+    return TestClient(app)
 
 
 @pytest.mark.asyncio
@@ -189,3 +236,32 @@ def test_team_info_masking():
     print("Got exception: {}".format(exc_info.value))
     assert "secret-test-key" not in str(exc_info.value)
     assert "public-test-key" not in str(exc_info.value)
+
+
+@mock_patch_aembedding()
+def test_embedding_input_array_of_tokens(mock_aembedding, client_no_auth):
+    """
+    Test to bypass decoding input as array of tokens for selected providers
+
+    Ref: https://github.com/BerriAI/litellm/issues/10113
+    """
+    try:
+        test_data = {
+            "model": "vllm_embed_model",
+            "input": [[2046, 13269, 158208]],
+        }
+
+        response = client_no_auth.post("/v1/embeddings", json=test_data)
+
+        mock_aembedding.assert_called_once_with(
+            model="vllm_embed_model",
+            input=[[2046, 13269, 158208]],
+            metadata=mock.ANY,
+            proxy_server_request=mock.ANY,
+        )
+        assert response.status_code == 200
+        result = response.json()
+        print(len(result["data"][0]["embedding"]))
+        assert len(result["data"][0]["embedding"]) > 10  # this usually has len==1536 so
+    except Exception as e:
+        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")