diff --git a/litellm/constants.py b/litellm/constants.py index 70db1133b30..4a5a00705f6 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 04872c6dae0..b38e4298558 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 00000000000..1b6b9ad198e --- /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 919a00d6703..ba787bc1e1e 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)}")