mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 03:06:22 +00:00
fix vllm passthrough
This commit is contained in:
@@ -11,7 +11,21 @@ from litellm.utils import _add_path_to_api_base
|
||||
|
||||
|
||||
class VLLMError(BaseLLMException):
|
||||
pass
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
request: Optional[httpx.Request] = None,
|
||||
response: Optional[httpx.Response] = None,
|
||||
headers: Optional[Union[httpx.Headers, dict]] = None,
|
||||
):
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
request=request,
|
||||
response=response,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
class VLLMModelInfo(BaseLLMModelInfo):
|
||||
@@ -25,7 +39,8 @@ class VLLMModelInfo(BaseLLMModelInfo):
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Google AI Studio sends api key in query params"""
|
||||
if api_key is not None:
|
||||
headers["x-api-key"] = api_key
|
||||
return headers
|
||||
|
||||
@staticmethod
|
||||
@@ -53,7 +68,7 @@ class VLLMModelInfo(BaseLLMModelInfo):
|
||||
endpoint = "/v1/models"
|
||||
if api_base is None or api_key is None:
|
||||
raise ValueError(
|
||||
"GEMINI_API_BASE or GEMINI_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint."
|
||||
"VLLM_API_BASE or VLLM_API_KEY is not set. Please set the environment variable, to query VLLM's `/models` endpoint."
|
||||
)
|
||||
|
||||
url = _add_path_to_api_base(api_base, endpoint)
|
||||
|
||||
@@ -233,14 +233,16 @@ async def get_request_body(request: Request) -> Dict[str, Any]:
|
||||
"""
|
||||
Read the request body and parse it as JSON.
|
||||
"""
|
||||
if request.headers.get("content-type") == "application/json":
|
||||
return await _read_request_body(request)
|
||||
elif (
|
||||
request.headers.get("content-type") == "multipart/form-data"
|
||||
or request.headers.get("content-type") == "application/x-www-form-urlencoded"
|
||||
):
|
||||
return await get_form_data(request)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported content type: {request.headers.get('content-type')}"
|
||||
)
|
||||
if request.method == "POST":
|
||||
if request.headers.get("content-type", "") == "application/json":
|
||||
return await _read_request_body(request)
|
||||
elif (
|
||||
"multipart/form-data" in request.headers.get("content-type", "")
|
||||
or "application/x-www-form-urlencoded" in request.headers.get("content-type", "")
|
||||
):
|
||||
return await get_form_data(request)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported content type: {request.headers.get('content-type')}"
|
||||
)
|
||||
return {}
|
||||
@@ -108,7 +108,16 @@ async def llm_passthrough_factory_proxy_route(
|
||||
|
||||
# Construct the full target URL using httpx
|
||||
base_url = httpx.URL(base_target_url)
|
||||
updated_url = base_url.copy_with(path=encoded_endpoint)
|
||||
# Join paths correctly by removing trailing/leading slashes as needed
|
||||
if not base_url.path or base_url.path == "/":
|
||||
# If base URL has no path, just use the new path
|
||||
updated_url = base_url.copy_with(path=encoded_endpoint)
|
||||
else:
|
||||
# Otherwise, combine the paths
|
||||
base_path = base_url.path.rstrip("/")
|
||||
clean_path = encoded_endpoint.lstrip("/")
|
||||
full_path = f"{base_path}/{clean_path}"
|
||||
updated_url = base_url.copy_with(path=full_path)
|
||||
|
||||
# Add or update query parameters
|
||||
provider_api_key = passthrough_endpoint_router.get_credentials(
|
||||
@@ -130,7 +139,11 @@ async def llm_passthrough_factory_proxy_route(
|
||||
is_streaming_request = False
|
||||
# anthropic is streaming when 'stream' = True is in the body
|
||||
if request.method == "POST":
|
||||
_request_body = await request.json()
|
||||
if "multipart/form-data" not in request.headers.get("content-type", ""):
|
||||
_request_body = await request.json()
|
||||
else:
|
||||
_request_body = dict(request._form)
|
||||
|
||||
if _request_body.get("stream"):
|
||||
is_streaming_request = True
|
||||
|
||||
|
||||
@@ -19,6 +19,8 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
BaseOpenAIPassThroughHandler,
|
||||
RouteChecks,
|
||||
create_pass_through_route,
|
||||
llm_passthrough_factory_proxy_route,
|
||||
vllm_proxy_route,
|
||||
vertex_discovery_proxy_route,
|
||||
vertex_proxy_route,
|
||||
bedrock_llm_proxy_route,
|
||||
@@ -914,3 +916,135 @@ class TestBedrockLLMProxyRoute:
|
||||
# For regular models, model should be just the model ID
|
||||
assert call_kwargs["model"] == "anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
assert result == "success"
|
||||
|
||||
|
||||
class TestLLMPassthroughFactoryProxyRoute:
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_passthrough_factory_proxy_route_success(self):
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.json = AsyncMock(return_value={"stream": False})
|
||||
mock_fastapi_response = MagicMock(spec=Response)
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_model_info"
|
||||
) as mock_get_provider, patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials"
|
||||
) as mock_get_creds, patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_create_route:
|
||||
mock_provider_config = MagicMock()
|
||||
mock_provider_config.get_api_base.return_value = "https://example.com/v1"
|
||||
mock_provider_config.validate_environment.return_value = {
|
||||
"Authorization": "Bearer test-key"
|
||||
}
|
||||
mock_get_provider.return_value = mock_provider_config
|
||||
mock_get_creds.return_value = "test-api-key"
|
||||
|
||||
mock_endpoint_func = AsyncMock(return_value="success")
|
||||
mock_create_route.return_value = mock_endpoint_func
|
||||
|
||||
result = await llm_passthrough_factory_proxy_route(
|
||||
custom_llm_provider="custom_provider",
|
||||
endpoint="/chat/completions",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_fastapi_response,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
mock_get_provider.assert_called_once_with(
|
||||
provider=litellm.LlmProviders("custom_provider"), model=None
|
||||
)
|
||||
mock_get_creds.assert_called_once_with(
|
||||
custom_llm_provider="custom_provider", region_name=None
|
||||
)
|
||||
mock_create_route.assert_called_once_with(
|
||||
endpoint="/chat/completions",
|
||||
target="https://example.com/v1/chat/completions",
|
||||
custom_headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
mock_endpoint_func.assert_awaited_once()
|
||||
|
||||
|
||||
class TestVLLMProxyRoute:
|
||||
@pytest.mark.asyncio
|
||||
async def test_vllm_proxy_route_with_router_model(self):
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.headers = {"content-type": "application/json"}
|
||||
mock_request.query_params = {}
|
||||
|
||||
mock_fastapi_response = MagicMock(spec=Response)
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
|
||||
return_value={"model": "router-model", "stream": False},
|
||||
), patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
|
||||
return_value=True,
|
||||
) as mock_is_router, patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.llm_router"
|
||||
) as mock_llm_router:
|
||||
mock_llm_router.allm_passthrough_route = AsyncMock()
|
||||
mock_response = httpx.Response(200, json={"response": "success"})
|
||||
mock_llm_router.allm_passthrough_route.return_value = mock_response
|
||||
|
||||
await vllm_proxy_route(
|
||||
endpoint="/chat/completions",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_fastapi_response,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
mock_is_router.assert_called_once()
|
||||
mock_llm_router.allm_passthrough_route.assert_awaited_once_with(
|
||||
model="router-model",
|
||||
method="POST",
|
||||
endpoint="/chat/completions",
|
||||
request_query_params={},
|
||||
request_headers={"content-type": "application/json"},
|
||||
stream=False,
|
||||
content=None,
|
||||
data=None,
|
||||
files=None,
|
||||
json={"model": "router-model", "stream": False},
|
||||
params=None,
|
||||
headers=None,
|
||||
cookies=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vllm_proxy_route_fallback_to_factory(self):
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_fastapi_response = MagicMock(spec=Response)
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
|
||||
return_value={"model": "other-model"},
|
||||
), patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
|
||||
return_value=False,
|
||||
), patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.llm_passthrough_factory_proxy_route"
|
||||
) as mock_factory_route:
|
||||
mock_factory_route.return_value = "factory_success"
|
||||
|
||||
result = await vllm_proxy_route(
|
||||
endpoint="/chat/completions",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_fastapi_response,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
assert result == "factory_success"
|
||||
mock_factory_route.assert_awaited_once_with(
|
||||
endpoint="/chat/completions",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_fastapi_response,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
custom_llm_provider="vllm",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user