From ce6d5f6e9edaf11599c6a19e949b5d8392af1ece Mon Sep 17 00:00:00 2001 From: Tanishq <74068132+tanishqbaweja@users.noreply.github.com> Date: Sun, 25 Oct 2026 22:59:13 +0530 Subject: [PATCH] fix(vllm): log native pooling passthrough responses --- .../base_llm/passthrough/transformation.py | 2 +- .../llms/vllm/passthrough/transformation.py | 7 ++- .../test_vllm_passthrough_transformation.py | 48 +++++++++++++++++++ 3 files changed, 54 insertions(+), 3 deletions(-) diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 7e58cc6e71b..0acded15efc 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -222,7 +222,7 @@ class BasePassthroughConfig(BaseLLMModelInfo): request_data: dict, logging_obj: LiteLLMLoggingObj, endpoint: str, - ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | None: + ) -> LoggedRelayResponse | OCRResponse | StandardPassThroughResponseObject | Mapping[str, object] | None: pass def handle_logging_collected_chunks( diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py index e7da6e54410..20ae7ba51d7 100644 --- a/litellm/llms/vllm/passthrough/transformation.py +++ b/litellm/llms/vllm/passthrough/transformation.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Final from httpx import Response 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, @@ -52,7 +52,10 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig): request_data: Mapping[str, object], logging_obj: Logging, endpoint: str, - ) -> "LoggedRelayResponse | None": + ) -> "LoggedRelayResponse | Mapping[str, object] | None": + if endpoint.strip("/").rsplit("/", 1)[-1] in ("pooling", "score", "rerank", "classify"): + body: Final = relayed_json_object(httpx_response) + return dict(body) if body is not None else 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) diff --git a/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py index bc8b7111fbb..93e3c98a513 100644 --- a/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py +++ b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py @@ -1,5 +1,6 @@ import json from datetime import datetime +from typing import Final import httpx import pytest @@ -239,3 +240,50 @@ 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( + ("endpoint", "request_data", "native_fields"), + [ + ( + "pooling", + {"input": ["Hello"], "task": "token_classify"}, + {"object": "list", "data": [{"index": 0, "object": "pooling", "data": [[0.1, 0.9], [0.8, 0.2]]}]}, + ), + ( + "v1/score", + {"text_1": "Hello", "text_2": ["Hi"]}, + {"object": "list", "data": [{"index": 0, "object": "score", "score": 0.9}]}, + ), + ( + "v2/rerank", + {"query": "Hello", "documents": ["Hi"]}, + {"results": [{"index": 0, "document": {"text": "Hi"}, "relevance_score": 0.9}]}, + ), + ( + "/classify/", + {"input": ["Hello"]}, + {"object": "list", "data": [{"index": 0, "label": "greeting", "probs": [0.1, 0.9], "num_classes": 2}]}, + ), + ], +) +async def test_native_pooling_responses_reach_standard_logging_with_their_ids_usage_and_full_body( + endpoint: str, request_data: dict[str, object], native_fields: dict[str, object] +) -> None: + body: Final = { + "id": "vllm-native-response", + "model": MODEL, + "usage": {"prompt_tokens": 7, "total_tokens": 7}, + **native_fields, + } + logging_obj: Final = _relay_logging_obj(endpoint, {"model": MODEL, **request_data}) + + await logging_obj.async_success_handler(result=_relayed_json(endpoint, body)) + + payload: Final = logging_obj.model_call_details.get("standard_logging_object") + assert payload is not None + assert payload["id"] == "vllm-native-response" + assert payload["custom_llm_provider"] == "hosted_vllm" + assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (7, 0, 7) + assert payload["response"] == body