diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index cc8a78fb50d..80767ea93f0 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -5,7 +5,10 @@ from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConf from ..common_utils import VLLMModelInfo if TYPE_CHECKING: - from httpx import URL + from httpx import URL, Response + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import CostResponseTypes class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): @@ -30,3 +33,55 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): self.format_url(endpoint, base_target_url, request_query_params), base_target_url, ) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: "Response", + request_data: dict, + logging_obj: "LiteLLMLoggingObj", + endpoint: str, + ) -> Optional["CostResponseTypes"]: + from litellm import encoding + from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig + from litellm.types.utils import EmbeddingResponse, ModelResponse, Usage + + if "chat/completions" in endpoint: + return OpenAIGPTConfig().transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data=request_data, + encoding=encoding, + ) + + endpoint_name = endpoint.rstrip("/").rsplit("/", 1)[-1] + if endpoint_name not in {"pooling", "embeddings", "classify", "score", "rerank"}: + return None + + try: + response_json = httpx_response.json() + except ValueError: + return None + + if not isinstance(response_json, dict): + return None + + usage = response_json.get("usage") + if not isinstance(usage, dict): + return None + + prompt_tokens = usage.get("prompt_tokens", 0) or 0 + return EmbeddingResponse( + model=response_json.get("model") or model, + usage=Usage( + prompt_tokens=prompt_tokens, + total_tokens=usage.get("total_tokens") or prompt_tokens, + ), + ) diff --git a/tests/test_litellm/llms/vllm/passthrough/test_vllm_passthrough_transformation.py b/tests/test_litellm/llms/vllm/passthrough/test_vllm_passthrough_transformation.py new file mode 100644 index 00000000000..b0a38242b23 --- /dev/null +++ b/tests/test_litellm/llms/vllm/passthrough/test_vllm_passthrough_transformation.py @@ -0,0 +1,102 @@ +from typing import Optional +from unittest.mock import MagicMock + +from httpx import Response + +from litellm.llms.vllm.passthrough.transformation import VLLMPassthroughConfig +from litellm.types.utils import CostResponseTypes, EmbeddingResponse, ModelResponse + + +def _run_response( + response: Response, endpoint: str = "/vllm/pooling" +) -> Optional[CostResponseTypes]: + return VLLMPassthroughConfig().logging_non_streaming_response( + model="intfloat/e5-mistral-7b-instruct", + custom_llm_provider="vllm", + httpx_response=response, + request_data={}, + logging_obj=MagicMock(), + endpoint=endpoint, + ) + + +def _run( + response_json: object, endpoint: str = "/vllm/pooling" +) -> Optional[CostResponseTypes]: + return _run_response(Response(status_code=200, json=response_json), endpoint) + + +def test_usage_bearing_response_is_cost_trackable(): + result = _run( + { + "model": "intfloat/e5-mistral-7b-instruct", + "data": [{"index": 0, "data": [0.1, 0.2]}], + "usage": {"prompt_tokens": 11, "total_tokens": 11}, + } + ) + assert isinstance(result, EmbeddingResponse) + assert result.model == "intfloat/e5-mistral-7b-instruct" + assert result.usage.prompt_tokens == 11 + assert result.usage.total_tokens == 11 + + +def test_missing_total_tokens_defaults_to_prompt_tokens(): + result = _run( + { + "model": "intfloat/e5-mistral-7b-instruct", + "data": [{"index": 0, "data": [0.1, 0.2]}], + "usage": {"prompt_tokens": 11}, + } + ) + assert isinstance(result, EmbeddingResponse) + assert result.usage.prompt_tokens == 11 + assert result.usage.total_tokens == 11 + + +def test_chat_completion_preserves_completion_tokens(): + result = _run( + { + "id": "chatcmpl-vllm", + "created": 1, + "model": "example-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "done"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 11, + "completion_tokens": 7, + "total_tokens": 18, + }, + }, + endpoint="/vllm/v1/chat/completions", + ) + assert isinstance(result, ModelResponse) + assert result.usage.prompt_tokens == 11 + assert result.usage.completion_tokens == 7 + assert result.usage.total_tokens == 18 + + +def test_unsupported_endpoint_returns_none(): + assert ( + _run( + {"usage": {"prompt_tokens": 11, "total_tokens": 11}}, + endpoint="/vllm/models", + ) + is None + ) + + +def test_response_without_usage_returns_none(): + assert _run({"model": "m", "data": []}) is None + + +def test_non_dict_response_returns_none(): + assert _run(["not", "a", "dict"]) is None + + +def test_unparseable_body_returns_none(): + assert _run_response(Response(status_code=200, content=b"not json")) is None