Merge b7945d6587 into e16d79d635

litellm/llms/vllm/passthrough/transformation.py

@@ -2,9 +2,10 @@ from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Final
 
 from httpx import Response
+from pydantic import ConfigDict, ValidationError
 
 from litellm.litellm_core_utils.litellm_logging import Logging
-from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
+from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, relayed_json_object
 from litellm.llms.openai_like.passthrough.transformation import (
     logged_openai_like_response,
     logged_openai_like_stream,
@@ -12,6 +13,8 @@ from litellm.llms.openai_like.passthrough.transformation import (
     logged_text_completion_stream,
     relays_text_completion,
 )
+from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.utils import ModelResponse, Usage
 
 from ..common_utils import VLLMModelInfo
 
@@ -21,6 +24,57 @@ if TYPE_CHECKING:
     from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
 
 
+class ReportedUsage(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+
+
+class UsageReportingBody(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    id: str | None = None
+    model: str | None = None
+    usage: ReportedUsage
+
+
+def logged_reported_usage(model: str, httpx_response: Response) -> ModelResponse | None:
+    """
+    The usage of a relayed body that has no OpenAI shape but reports one anyway.
+
+    vLLM's own endpoints such as `/score`, `/classify` and `/pooling`, and servers built on vLLM such as
+    vLLM Semantic Router's `/v1/systemone`, answer with a `usage` object (`prompt_tokens`/`completion_tokens`,
+    or `input_tokens`/`output_tokens`) but none of the shapes the OpenAI relay parses. Without this, those
+    calls are logged with no tokens at all.
+    """
+    body: Final = relayed_json_object(httpx_response)
+    if body is None:
+        return None
+    try:
+        relayed: Final = UsageReportingBody.model_validate(body)
+    except ValidationError:
+        return None
+    usage: Final = relayed.usage
+    prompt_tokens: Final = usage.prompt_tokens if usage.prompt_tokens is not None else usage.input_tokens
+    completion_tokens: Final = usage.completion_tokens if usage.completion_tokens is not None else usage.output_tokens
+    if prompt_tokens is None and completion_tokens is None:
+        return None
+    response: Final = ModelResponse(
+        model=relayed.model or model,
+        usage=Usage(
+            prompt_tokens=prompt_tokens or 0,
+            completion_tokens=completion_tokens or 0,
+            total_tokens=(prompt_tokens or 0) + (completion_tokens or 0),
+        ),
+    )
+    if relayed.id:
+        response.id = relayed.id
+    return response
+
+
 class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
     def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
         return "stream" in request_data
@@ -55,7 +109,10 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
     ) -> "LoggedRelayResponse | None":
         if relays_text_completion(endpoint):
             return logged_text_completion_response(httpx_response, logging_obj, endpoint)
-        return logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
+        openai_shaped: Final = logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
+        if openai_shaped is not None:
+            return openai_shaped
+        return logged_reported_usage(model, httpx_response)
 
     def handle_logging_collected_chunks(
         self,

tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py

@@ -239,3 +239,68 @@ async def test_a_relayed_text_completion_stream_without_a_usage_chunk_counts_eve
     assert payload["completion_tokens"] == sum(
         litellm.token_counter(model=MODEL, text=text, count_response_tokens=True) for text in ["Hello there", "Hey you"]
     )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("path", "body", "tokens"),
+    [
+        (
+            "v1/score",
+            {
+                "id": "score-vllm-7",
+                "object": "list",
+                "created": 1,
+                "model": MODEL,
+                "data": [{"index": 0, "object": "score", "score": 0.91}],
+                "usage": {"prompt_tokens": 15, "total_tokens": 15, "completion_tokens": 0},
+            },
+            (15, 0),
+        ),
+        (
+            "classify",
+            {
+                "id": "classify-vllm-8",
+                "object": "list",
+                "created": 1,
+                "model": MODEL,
+                "data": [{"index": 0, "label": "positive", "probs": [0.1, 0.9], "num_classes": 2}],
+                "usage": {"prompt_tokens": 9, "total_tokens": 9},
+            },
+            (9, 0),
+        ),
+        (
+            "v1/systemone",
+            {
+                "model": MODEL,
+                "answers": {"urgent": {"type": "noul", "noul": 0.85}},
+                "usage": {"input_tokens": 209, "output_tokens": 0},
+            },
+            (209, 0),
+        ),
+    ],
+)
+async def test_a_relayed_body_without_an_openai_shape_is_logged_with_the_usage_it_reports(path, body, tokens):
+    logging_obj = _relay_logging_obj(path, {"model": MODEL})
+
+    await logging_obj.async_success_handler(result=_relayed_json(path, body))
+
+    payload = logging_obj.model_call_details["standard_logging_object"]
+    assert payload["custom_llm_provider"] == "hosted_vllm"
+    assert (payload["prompt_tokens"], payload["completion_tokens"]) == tokens
+    assert payload["total_tokens"] == sum(tokens)
+
+
+def test_a_relayed_body_without_an_openai_shape_or_a_usage_is_not_given_one():
+    body = {"count": 3, "max_model_len": 4096, "tokens": [1, 2, 3]}
+
+    logged = VLLMPassthroughConfig().logging_non_streaming_response(
+        model=MODEL,
+        custom_llm_provider="hosted_vllm",
+        httpx_response=_relayed_json("tokenize", body),
+        request_data={"model": MODEL, "prompt": PROMPT},
+        logging_obj=_relay_logging_obj("tokenize", {"model": MODEL, "prompt": PROMPT}),
+        endpoint="tokenize",
+    )
+
+    assert logged is None