Merge ce6d5f6e9e into c6ca22ebc5

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(

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)

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