From b88e56ebde79fb3473edf5ecb500887922176afc Mon Sep 17 00:00:00 2001 From: Camille Farineau Date: Thu, 15 May 2025 07:55:49 +0200 Subject: [PATCH] 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 a48acd95f860c4fc85853e20668eabffff07cae7. * 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 | 6 ++ litellm/proxy/proxy_server.py | 34 +++++---- .../test_configs/test_config_no_auth.yaml | 6 ++ tests/litellm/proxy/test_proxy_server.py | 76 +++++++++++++++++++ 4 files changed, 107 insertions(+), 15 deletions(-) create mode 100644 tests/litellm/proxy/test_configs/test_config_no_auth.yaml diff --git a/litellm/constants.py b/litellm/constants.py index 70db1133b3..4a5a00705f 100644 --- a/litellm/constants.py +++ b/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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 04872c6dae..b38e429855 100644 --- a/litellm/proxy/proxy_server.py +++ b/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( diff --git a/tests/litellm/proxy/test_configs/test_config_no_auth.yaml b/tests/litellm/proxy/test_configs/test_config_no_auth.yaml new file mode 100644 index 0000000000..1b6b9ad198 --- /dev/null +++ b/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 diff --git a/tests/litellm/proxy/test_proxy_server.py b/tests/litellm/proxy/test_proxy_server.py index 919a00d670..ba787bc1e1 100644 --- a/tests/litellm/proxy/test_proxy_server.py +++ b/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)}")