fix(vllm): track spend for passthrough endpoints that report usage

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,
+            ),
+        )

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