From aac424df1a6959b3193ca2005cfe7a69182b8124 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 23:35:28 +0000 Subject: [PATCH] fix(responses): carry logprobs across the chat completions bridge A /v1/responses request with top_logprobs on a chat-bridged deployment sent top_logprobs without logprobs, which OpenAI-style chat APIs and Vertex Gemini reject, and the bridge dropped the chat choice's logprobs on the way back. The bridge now sends logprobs with top_logprobs and copies choice.logprobs into output_text.logprobs, on the response and on the output_text.delta, output_text.done and content_part.done stream events Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../logprobs.py | 34 +++++ .../streaming_iterator.py | 26 +++- .../transformation.py | 17 ++- litellm/types/llms/base.py | 32 +++- litellm/types/llms/openai.py | 15 +- litellm/types/responses/main.py | 4 +- tests/integration/proxy_config.yaml | 7 + .../translation/responses/bases/vertex_ai.py | 35 +++++ .../basic/test_responses_basic_vertex_ai.py | 2 + .../test_responses_sampling_vertex_ai.py | 113 ++++++++++++++ .../test_litellm_completion_responses.py | 109 +++++++++++++- .../test_streaming_iterator_transformation.py | 141 ++++++++++-------- 12 files changed, 451 insertions(+), 84 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/logprobs.py create mode 100644 tests/integration/translation/responses/sampling/test_responses_sampling_vertex_ai.py diff --git a/litellm/responses/litellm_completion_transformation/logprobs.py b/litellm/responses/litellm_completion_transformation/logprobs.py new file mode 100644 index 00000000000..81953ccd602 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/logprobs.py @@ -0,0 +1,34 @@ +"""Chat ``choice.logprobs`` as Responses ``output_text.logprobs``, which share one token shape.""" + +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.types.llms.base import OutputTextLogprob, OutputTextTopLogprob +from litellm.types.utils import ChatCompletionTokenLogprob, ChoiceLogprobs, TopLogprob + +_CHOICE_LOGPROBS: Final = TypeAdapter(ChoiceLogprobs | None) + + +def _top_logprob(top: TopLogprob) -> OutputTextTopLogprob: + return {"token": top.token, "logprob": top.logprob, "bytes": top.bytes} + + +def _token_logprob(token: ChatCompletionTokenLogprob) -> OutputTextLogprob: + return { + "token": token.token, + "logprob": token.logprob, + "bytes": token.bytes, + "top_logprobs": [_top_logprob(top) for top in token.top_logprobs], + } + + +def output_text_logprobs(choice_logprobs: object) -> list[OutputTextLogprob] | None: + """The ``logprobs`` a Responses text part carries for a chat choice, ``None`` when the choice sent none.""" + try: + validated: Final = _CHOICE_LOGPROBS.validate_python(choice_logprobs) + except ValidationError: + return None + if validated is None or validated.content is None: + return None + return [_token_logprob(token) for token in validated.content] diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index ddcdc66ff20..0f84159c7dc 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -44,6 +44,7 @@ from litellm.types.llms.openai import ( OutputTextAnnotationAddedEvent, OutputTextDeltaEvent, OutputTextDoneEvent, + OutputTextLogprob, ReasoningSummaryPartDoneEvent, ReasoningSummaryTextDeltaEvent, ReasoningSummaryTextDoneEvent, @@ -71,6 +72,8 @@ from litellm.types.utils import ( TextCompletionResponse, ) +from .logprobs import output_text_logprobs + def _index_of_output_item_type(items: Sequence[object], item_type: str) -> int | None: return next( @@ -210,6 +213,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._sent_reasoning_summary_text_done_event: bool = False self._sent_reasoning_summary_part_done_event: bool = False self._reasoning_summary_text: str = "" + self._text_logprobs: tuple[OutputTextLogprob, ...] = () # -- GENERIC RESPONSE-EVENTS PENDING QUEUE as required by fix -- self._pending_response_events: list[BaseLiteLLMOpenAIResponseObject] = [] self._reasoning_active = False @@ -1014,8 +1018,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): output_index=self._message_index(), content_index=0, text=getattr(litellm_complete_object.choices[0].message, "content", "") or "", + logprobs=self._final_text_logprobs(litellm_complete_object), ) + def _final_text_logprobs(self, litellm_complete_object: ModelResponse) -> list[OutputTextLogprob] | None: + if self._text_logprobs: + return list(self._text_logprobs) + return output_text_logprobs(getattr(litellm_complete_object.choices[0], "logprobs", None)) + def create_output_content_part_done_event(self, litellm_complete_object: ModelResponse) -> ContentPartDoneEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{uuid.uuid4()}" @@ -1032,7 +1042,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type="output_text", text=text, annotations=response_annotations, - logprobs=None, + logprobs=self._final_text_logprobs(litellm_complete_object), ) return ContentPartDoneEvent( @@ -1056,8 +1066,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): annotations=annotations ) ) + logprobs: Final = self._final_text_logprobs(litellm_complete_object) text_parts: Final = ( - ({"type": "output_text", "text": text, "annotations": response_annotations},) + ( + { + "type": "output_text", + "text": text, + "annotations": response_annotations, + **({} if logprobs is None else {"logprobs": logprobs}), + }, + ) if self.sent_content_part_added_event else () ) @@ -1505,12 +1523,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._queue_message_item_added_events() self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) self._sequence_number += 1 + delta_logprobs: Final = output_text_logprobs(getattr(chunk.choices[0], "logprobs", None)) + if delta_logprobs is not None: + self._text_logprobs = (*self._text_logprobs, *delta_logprobs) text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, output_index=self._message_index(), content_index=0, delta=delta_content, + logprobs=delta_logprobs, ) text_delta_event.__dict__["sequence_number"] = self._sequence_number return text_delta_event diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 0f4ed0611be..66750d26eaa 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -111,6 +111,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, OpenAIWebSearchOptions, OpenAIWebSearchUserLocation, + OutputTextLogprob, OutputTokensDetails, Reasoning, ResponseAPIUsage, @@ -167,6 +168,7 @@ from .custom_tools import ( validated_tool_search_execution, without_malformed_tool_flags, ) +from .logprobs import output_text_logprobs from .reasoning_items import decode_thinking_blocks, encode_thinking_blocks, mint_reasoning_item_id from .tool_search import ( CompletionRoute, @@ -950,6 +952,7 @@ class LiteLLMCompletionResponsesConfig: "text", "tool_choice", "tools", + "top_logprobs", "top_p", "user", ] @@ -1289,6 +1292,8 @@ class LiteLLMCompletionResponsesConfig: ), "tools": tools, "top_p": responses_api_request.get("top_p"), + "top_logprobs": responses_api_request.get("top_logprobs"), + "logprobs": True if responses_api_request.get("top_logprobs") is not None else None, "user": responses_api_request.get("user"), "temperature": responses_api_request.get("temperature"), "parallel_tool_calls": responses_api_request.get("parallel_tool_calls"), @@ -3695,7 +3700,9 @@ class LiteLLMCompletionResponsesConfig: choice=choice, ) message_output_items.extend(image_generation_items) - elif content := LiteLLMCompletionResponsesConfig._message_output_content(choice.message): + elif content := LiteLLMCompletionResponsesConfig._message_output_content( + choice.message, output_text_logprobs(getattr(choice, "logprobs", None)) + ): message_output_items.append( GenericResponseOutputItem( type="message", @@ -3708,7 +3715,9 @@ class LiteLLMCompletionResponsesConfig: return message_output_items @staticmethod - def _message_output_content(message: Message) -> tuple[OutputText | ResponseOutputRefusal, ...]: + def _message_output_content( + message: Message, logprobs: list[OutputTextLogprob] | None = None + ) -> tuple[OutputText | ResponseOutputRefusal, ...]: refusal: Final = openai_chat_refusal_text(message) refusal_parts: Final = (ResponseOutputRefusal(type="refusal", refusal=refusal),) if refusal is not None else () has_text_part: Final = ( @@ -3717,7 +3726,7 @@ class LiteLLMCompletionResponsesConfig: or (message.content is not None and not refusal_parts) ) text_parts: Final = ( - (LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text(message),) + (LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text(message, logprobs),) if has_text_part else () ) @@ -3769,6 +3778,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_chat_message_to_response_output_text( message: Message, + logprobs: list[OutputTextLogprob] | None = None, ) -> OutputText: annotations: Final = getattr(message, "annotations", None) transformed_annotations: Final = ( @@ -3781,6 +3791,7 @@ class LiteLLMCompletionResponsesConfig: type="output_text", text=message.content or "", annotations=transformed_annotations, + logprobs=logprobs, ) @staticmethod diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index 5a2da071cec..d82b73705ef 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -1,7 +1,9 @@ -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from openai._models import BaseModel as OpenAIObject -from pydantic import ConfigDict +from pydantic import ConfigDict, SerializerFunctionWrapHandler, model_serializer +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base_model import LiteLLMBaseModel @@ -83,3 +85,29 @@ class CachedTokensDetails(LiteLLMBaseModel): text_tokens: int | None = None audio_tokens: int | None = None image_tokens: int | None = None + + +class OutputTextTopLogprob(TypedDict): + token: ReadOnly[str] + logprob: ReadOnly[float] + bytes: ReadOnly[list[int] | None] + + +class OutputTextLogprob(TypedDict): + token: ReadOnly[str] + logprob: ReadOnly[float] + bytes: ReadOnly[list[int] | None] + top_logprobs: ReadOnly[list[OutputTextTopLogprob]] + + +class CarriesOutputTextLogprobs(BaseLiteLLMOpenAIResponseObject): + """A Responses text payload carries ``logprobs`` only when the chat choice behind it sent them.""" + + logprobs: list[OutputTextLogprob] | None = None + + @model_serializer(mode="wrap") + def _without_absent_logprobs(self, handler: SerializerFunctionWrapHandler) -> Mapping[str, object]: + dumped: Final = cast("Mapping[str, object]", handler(self)) # cast-ok: handler dumps this model's fields + if self.logprobs is not None: + return dumped + return {key: value for key, value in dumped.items() if key != "logprobs"} diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d7e55b57a07..c079cba3aba 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -3,7 +3,7 @@ import math from collections.abc import Iterable, Mapping, Sequence from enum import Enum from os import PathLike -from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union +from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union, cast import httpx from openai import Omit @@ -86,7 +86,12 @@ from typing_extensions import ( ) from litellm._logging import verbose_logger -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, LiteLLMBaseModel +from litellm.types.llms.base import ( + BaseLiteLLMOpenAIResponseObject, + CarriesOutputTextLogprobs, + LiteLLMBaseModel, + OutputTextLogprob, +) from litellm.types.responses.main import ( CustomToolCallOutputItem, GenericResponseOutputItem, @@ -1863,7 +1868,7 @@ class ContentPartDonePartOutputText(BaseLiteLLMOpenAIResponseObject): type: Literal["output_text"] text: str annotations: list[BaseLiteLLMOpenAIResponseObject] - logprobs: list[OpenAIChatCompletionLogprobsContent] | None = None + logprobs: list[OutputTextLogprob] | None = None class ContentPartDonePartRefusal(BaseLiteLLMOpenAIResponseObject): @@ -1887,7 +1892,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): part: PART_UNION_TYPES -class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDeltaEvent(CarriesOutputTextLogprobs): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA] item_id: str output_index: int @@ -1904,7 +1909,7 @@ class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject): annotation: dict -class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDoneEvent(CarriesOutputTextLogprobs): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE] item_id: str output_index: int diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 86dc857cfd0..a8c8410e723 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -8,7 +8,7 @@ from openai.types.responses.response_output_refusal import ResponseOutputRefusal from pydantic import ConfigDict, PrivateAttr from typing_extensions import TypedDict -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject, CarriesOutputTextLogprobs Phase = Literal["commentary", "final_answer"] | None @@ -23,7 +23,7 @@ class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject title: str | None -class OutputText(BaseLiteLLMOpenAIResponseObject): +class OutputText(CarriesOutputTextLogprobs): """Text output content from an assistant message""" type: str | None # "output_text" diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 4e2cab56d27..3047e53f61b 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -501,6 +501,13 @@ model_list: vertex_project: synthetic-vertex-project vertex_location: global vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' + - model_name: vertex_ai/gemini-2.5-flash + litellm_params: + model: vertex_ai/gemini-2.5-flash + api_base: http://127.0.0.1:8191 + vertex_project: synthetic-vertex-project + vertex_location: global + vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' - model_name: vertex_ai/gemini-3.8-flash litellm_params: model: vertex_ai/gemini-3.8-flash diff --git a/tests/integration/translation/responses/bases/vertex_ai.py b/tests/integration/translation/responses/bases/vertex_ai.py index 59637e380ce..3b31b995aa8 100644 --- a/tests/integration/translation/responses/bases/vertex_ai.py +++ b/tests/integration/translation/responses/bases/vertex_ai.py @@ -1,4 +1,5 @@ import json +from dataclasses import replace from typing import Final from unittest.mock import ANY @@ -1858,3 +1859,37 @@ GEMINI_4_ARGON_TEST_CASE: Final = TranslationTestCase( "provider_specific_fields": {"traffic_type": "ON_DEMAND"}, }, ) + + +GEMINI_2_5_FLASH_TEST_CASE: Final = replace( + GEMINI_3_5_FLASH_TEST_CASE, + litellm_request={**GEMINI_3_5_FLASH_TEST_CASE.litellm_request, "model": "vertex_ai/gemini-2.5-flash"}, + expected_provider_endpoint=( + "/v1/projects/synthetic-vertex-project/locations/global/publishers/google/models/gemini-2.5-flash:generateContent" + ), + mock_provider_response={ + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hello."}]}, "finishReason": "STOP"}], + "usageMetadata": { + "promptTokenCount": 9, + "candidatesTokenCount": 2, + "totalTokenCount": 32, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 9}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 2}], + "thoughtsTokenCount": 21, + }, + "modelVersion": "gemini-2.5-flash", + "createTime": "2026-10-31T23:26:23.954987Z", + "responseId": "H3nmauukOrv46tkPw5vggAo", + }, + expected_litellm_response={ + **GEMINI_3_5_FLASH_TEST_CASE.expected_litellm_response, + "model": "vertex_ai/gemini-2.5-flash", + "usage": { + **GEMINI_3_5_FLASH_TEST_CASE.expected_litellm_response["usage"], + "output_tokens": 23, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 21, "text_tokens": 2}, + "total_tokens": 32, + }, + }, +) diff --git a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py index bd9118482d0..6d6e08835a3 100644 --- a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py +++ b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py @@ -11,6 +11,7 @@ from integration.translation.responses.bases.vertex_ai import ( CLAUDE_SONNET_4_6_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, CLAUDE_SONNET_5_TEST_CASE, + GEMINI_2_5_FLASH_TEST_CASE, GEMINI_3_1_PRO_PREVIEW_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, GEMINI_3_8_FLASH_TEST_CASE, @@ -30,6 +31,7 @@ from integration.translation.runner import assert_translation CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, + GEMINI_2_5_FLASH_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, GEMINI_3_8_FLASH_TEST_CASE, GEMINI_3_1_PRO_PREVIEW_TEST_CASE, diff --git a/tests/integration/translation/responses/sampling/test_responses_sampling_vertex_ai.py b/tests/integration/translation/responses/sampling/test_responses_sampling_vertex_ai.py new file mode 100644 index 00000000000..a89e8a867a7 --- /dev/null +++ b/tests/integration/translation/responses/sampling/test_responses_sampling_vertex_ai.py @@ -0,0 +1,113 @@ +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.responses.bases.vertex_ai import GEMINI_2_5_FLASH_TEST_CASE +from integration.translation.runner import assert_translation + +GENERATION_CONFIG: Final = object_value(GEMINI_2_5_FLASH_TEST_CASE.expected_provider_request["generationConfig"]) +USAGE: Final = object_value(GEMINI_2_5_FLASH_TEST_CASE.expected_litellm_response["usage"]) + +GEMINI_2_5_FLASH_TOP_LOGPROBS_TEST_CASE: Final = replace( + GEMINI_2_5_FLASH_TEST_CASE, + scenario="top_logprobs", + litellm_request={**GEMINI_2_5_FLASH_TEST_CASE.litellm_request, "top_logprobs": 2}, + expected_provider_request={ + **GEMINI_2_5_FLASH_TEST_CASE.expected_provider_request, + "generationConfig": {**GENERATION_CONFIG, "responseLogprobs": True, "logprobs": 2}, + }, + mock_provider_response={ + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "Hello."}]}, + "finishReason": "STOP", + "avgLogprobs": -2.20536901451851e-06, + "logprobsResult": { + "topCandidates": [ + { + "candidates": [ + {"token": "Hello", "logProbability": -1.3113013e-06, "tokenId": 9259}, + {"token": "Hi", "logProbability": -15.149978, "tokenId": 10979}, + ] + }, + { + "candidates": [ + {"token": ".", "logProbability": -3.0994368e-06, "tokenId": 236761}, + {"token": "", "logProbability": -13.81299, "tokenId": 106}, + ] + }, + ], + "chosenCandidates": [ + {"token": "Hello", "logProbability": -1.3113013e-06, "tokenId": 9259}, + {"token": ".", "logProbability": -3.0994368e-06, "tokenId": 236761}, + ], + }, + } + ], + "usageMetadata": { + "promptTokenCount": 9, + "candidatesTokenCount": 2, + "totalTokenCount": 30, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 9}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 2}], + "thoughtsTokenCount": 19, + }, + "modelVersion": "gemini-2.5-flash", + "createTime": "2026-10-31T23:20:20.755722Z", + "responseId": "tHfmaoqQLrKZ9LsPvuuB-AU", + }, + expected_litellm_response={ + **GEMINI_2_5_FLASH_TEST_CASE.expected_litellm_response, + "output": [ + { + "type": "message", + "id": GEMINI_2_5_FLASH_TEST_CASE.expected_litellm_response["output"][0]["id"], + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello.", + "annotations": [], + "logprobs": [ + { + "token": "Hello", + "logprob": -1.3113013e-06, + "bytes": None, + "top_logprobs": [ + {"token": "Hello", "logprob": -1.3113013e-06, "bytes": None}, + {"token": "Hi", "logprob": -15.149978, "bytes": None}, + ], + }, + { + "token": ".", + "logprob": -3.0994368e-06, + "bytes": None, + "top_logprobs": [ + {"token": ".", "logprob": -3.0994368e-06, "bytes": None}, + {"token": "", "logprob": -13.81299, "bytes": None}, + ], + }, + ], + } + ], + "phase": None, + } + ], + "usage": { + **USAGE, + "output_tokens": 21, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 19, "text_tokens": 2}, + "total_tokens": 30, + }, + }, +) + + +@pytest.mark.parametrize("case", [GEMINI_2_5_FLASH_TOP_LOGPROBS_TEST_CASE], ids=lambda case: case.id) +def test_responses_sampling_vertex_ai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b4cd606374a..c446fba115c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -18,7 +18,6 @@ from openai.types.responses.response_function_web_search import ( from pydantic import TypeAdapter, ValidationError import litellm -from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field from litellm.constants import CALLER_SESSION_ID_METADATA_KEY from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR @@ -26,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, anthropic_messages_pt, ) +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks from litellm.responses.litellm_completion_transformation.tool_search import tool_search_routing from litellm.responses.litellm_completion_transformation.transformation import ( @@ -38,12 +38,15 @@ from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams from litellm.types.responses.main import OutputToolSearchCall, build_web_search_call from litellm.types.utils import ( ChatCompletionMessageToolCall, + ChatCompletionTokenLogprob, + ChoiceLogprobs, Choices, CompletionTokensDetailsWrapper, Function, Message, ModelResponse, PromptTokensDetailsWrapper, + TopLogprob, Usage, ) @@ -6876,7 +6879,11 @@ def test_a_new_message_after_a_paused_search_turn_replays_only_that_turn_without [ {"role": "user", "content": "Plan a Tokyo trip."}, _reasoning_item("rs_tokyo", tokyo_thinking), - {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Tokyo is sunny."}]}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Tokyo is sunny."}], + }, _hosted_search("ws_srvtoolu_01Tokyo", "completed"), {"role": "user", "content": "Now Osaka."}, _reasoning_item("rs_osaka", _ANSWERED_THINKING), @@ -7216,7 +7223,11 @@ def test_a_tool_output_with_an_unhashable_call_id_is_a_400_naming_its_input_item ) -> None: history: Final = cast( ResponseInputParam, - [_WEATHER_ASK, _weather_call(call_id="call_weather"), {"type": output_type, "call_id": call_id, "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id="call_weather"), + {"type": output_type, "call_id": call_id, "output": "sunny"}, + ], ) with pytest.raises(litellm.BadRequestError) as raised: @@ -7320,7 +7331,11 @@ def test_a_non_string_role_still_converts_for_a_target_that_judges_it_itself() - id="dict_function_call_id", ), pytest.param( - [_WEATHER_ASK, _weather_call(call_id=7), {"type": "function_call_output", "call_id": "7", "output": "sunny"}], + [ + _WEATHER_ASK, + _weather_call(call_id=7), + {"type": "function_call_output", "call_id": "7", "output": "sunny"}, + ], "input[1].call_id", "int", id="int_function_call_id", @@ -7394,3 +7409,89 @@ def test_a_malformed_item_after_a_codex_additional_tools_item_is_named_by_its_in ) assert f"'{field_path}'" in raised.value.message + + +_HELLO_LOGPROBS: Final = ChoiceLogprobs( + content=[ + ChatCompletionTokenLogprob( + token="Hello", + logprob=-1.3113013e-06, + top_logprobs=[ + TopLogprob(token="Hello", logprob=-1.3113013e-06), + TopLogprob(token="Hi", logprob=-15.149978), + ], + ), + ChatCompletionTokenLogprob(token=".", logprob=-3.0994368e-06, bytes=[46]), + ] +) + + +def _hello_chat_response(logprobs: ChoiceLogprobs | None) -> ModelResponse: + return ModelResponse( + id="chatcmpl-hello", + created=1793489000, + model="gemini-3.5-flash", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", index=0, message=Message(content="Hello.", role="assistant"), logprobs=logprobs + ) + ], + ) + + +def test_bridged_chat_logprobs_become_output_text_logprobs() -> None: + response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Say hello.", + responses_api_request={"top_logprobs": 2}, + chat_completion_response=_hello_chat_response(_HELLO_LOGPROBS), + ) + + assert response.model_dump()["output"][0]["content"] == [ + { + "type": "output_text", + "text": "Hello.", + "annotations": [], + "logprobs": [ + { + "token": "Hello", + "logprob": -1.3113013e-06, + "bytes": None, + "top_logprobs": [ + {"token": "Hello", "logprob": -1.3113013e-06, "bytes": None}, + {"token": "Hi", "logprob": -15.149978, "bytes": None}, + ], + }, + {"token": ".", "logprob": -3.0994368e-06, "bytes": [46], "top_logprobs": []}, + ], + } + ] + + +def test_bridged_chat_response_without_logprobs_emits_no_logprobs_key() -> None: + response: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Say hello.", responses_api_request={}, chat_completion_response=_hello_chat_response(None) + ) + + assert response.model_dump()["output"][0]["content"] == [ + {"type": "output_text", "text": "Hello.", "annotations": []} + ] + + +def test_responses_top_logprobs_bridges_to_chat_logprobs() -> None: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini-3.5-flash", + input="Say hello.", + responses_api_request={"top_logprobs": 2}, + custom_llm_provider="vertex_ai", + ) + + assert (request["top_logprobs"], request["logprobs"]) == (2, True) + + +def test_responses_without_top_logprobs_bridges_without_chat_logprobs() -> None: + request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini-3.5-flash", input="Say hello.", responses_api_request={}, custom_llm_provider="vertex_ai" + ) + + assert "top_logprobs" not in request and "logprobs" not in request diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 684f5a89e80..4605857dcae 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -30,6 +30,8 @@ from litellm.types.llms.openai import ( ) from litellm.types.responses.main import build_web_search_call from litellm.types.utils import ( + ChatCompletionTokenLogprob, + ChoiceLogprobs, Delta, ModelResponse, ModelResponseStream, @@ -38,9 +40,7 @@ from litellm.types.utils import ( ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" -RESPONSE_ID_EVENT_TYPES = frozenset( - {"response.created", "response.in_progress", "response.completed"} -) +RESPONSE_ID_EVENT_TYPES = frozenset({"response.created", "response.in_progress", "response.completed"}) def _chunk(content: str, finish_reason: str | None = None) -> ModelResponseStream: @@ -93,11 +93,7 @@ def _build_iterator(chunks) -> LiteLLMCompletionStreamingIterator: def _response_ids(events) -> list[str]: - return [ - event.response.id - for event in events - if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] + return [event.response.id for event in events if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] def test_tool_call_delta_is_emitted_as_responses_events(): @@ -228,9 +224,7 @@ async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, to content=None, provider_specific_fields={ "web_search_results": [first_result], - "web_search_calls": [ - build_web_search_call(call_ids[0], {"query": "one"}, first_result) - ] + "web_search_calls": [build_web_search_call(call_ids[0], {"query": "one"}, first_result)] if tool_type != "function" and first_result["type"] == "web_search_tool_result" else [], }, @@ -329,9 +323,7 @@ async def test_web_search_stream_preserves_hosted_and_client_calls(sync_mode, to assert {event["output_index"] for event in search_events} == {completed_indexes[item["id"]]} assert search_events[0]["item"]["status"] == "in_progress" assert search_events[-1]["item"] == item - assert item["status"] == ( - "failed" if result_kind == "error" and call_id.endswith("01Search") else "completed" - ) + assert item["status"] == ("failed" if result_kind == "error" and call_id.endswith("01Search") else "completed") assert item["action"]["type"] == "search" assert item["action"]["query"] == ("one" if call_id.endswith("01Search") else "two") assert item["action"]["queries"] == [item["action"]["query"]] @@ -440,9 +432,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): ) # Create a chunk with a large arguments string that should be split - large_arguments = ( - '{"param1": "value1", "param2": "value2", "param3": "value3"}' # 67 chars - ) + large_arguments = '{"param1": "value1", "param2": "value2", "param3": "value3"}' # 67 chars chunk = ModelResponseStream( id="chunk-1", created=123, @@ -497,9 +487,7 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): # Keep draining pending events (expected: ceil(67 / 10) = 7 delta events) while iterator._pending_tool_events: - evt = iterator._transform_chat_completion_chunk_to_response_api_chunk( - empty_chunk - ) + evt = iterator._transform_chat_completion_chunk_to_response_api_chunk(empty_chunk) if evt and evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: delta_events.append(evt) @@ -552,20 +540,12 @@ def test_tool_call_delta_without_id_uses_index_mapping(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] streamed_arguments = "".join(evt.delta for evt in delta_events) assert streamed_arguments == '{"location": "New York"}' - output_item_added_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED - ] + output_item_added_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] assert len(output_item_added_events) == 1 assert output_item_added_events[0].item.id == "fc_call_abc123" assert output_item_added_events[0].item.call_id == "call_abc123" @@ -606,18 +586,10 @@ def test_parallel_tool_calls_without_ids_use_index_mapping(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - output_item_added_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED - ] + output_item_added_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] assert len(output_item_added_events) == 2 - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] arguments_by_call_id = {} for evt in delta_events: arguments_by_call_id.setdefault(evt.item_id, "") @@ -670,11 +642,7 @@ def test_reused_index_with_new_call_id_marks_fallback_ambiguous(): while iterator._pending_tool_events: all_events.append(iterator._pending_tool_events.pop(0)) - delta_events = [ - evt - for evt in all_events - if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA - ] + delta_events = [evt for evt in all_events if evt.type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] arguments_by_call_id = {} for evt in delta_events: arguments_by_call_id.setdefault(evt.item_id, "") @@ -741,22 +709,17 @@ def test_sync_streaming_events_share_the_chat_completion_response_id(): assert len(response_ids) == 3 assert len(set(response_ids)) == 1 assert ( - ResponsesAPIRequestUtils.decode_responses_api_response_id(response_ids[0])["response_id"] - == CHAT_COMPLETION_ID + ResponsesAPIRequestUtils.decode_responses_api_response_id(response_ids[0])["response_id"] == CHAT_COMPLETION_ID ) @pytest.mark.asyncio async def test_streaming_emits_every_chunk_after_priming_the_response_id(): - iterator = _build_iterator( - [_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")] - ) + iterator = _build_iterator([_chunk("Hel"), _chunk("lo"), _chunk("!", finish_reason="stop")]) events = [event async for event in iterator] - deltas = "".join( - event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta" - ) + deltas = "".join(event.delta for event in events if getattr(event, "type", None) == "response.output_text.delta") assert deltas == "Hello!" @@ -802,7 +765,9 @@ async def test_leading_empty_choices_chunk_does_not_kill_the_stream(): event_types = [getattr(event, "type", None) for event in events] assert event_types.count(ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) == 1 - assert "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + assert ( + "".join(event.delta for event in events if event.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Hello!" + ) assert event_types[-1] == ResponsesAPIStreamEvents.RESPONSE_COMPLETED @@ -986,9 +951,7 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: litellm_metadata={}, ) - response_events: Final = [ - event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES - ] + response_events: Final = [event for event in iterator if getattr(event, "type", None) in RESPONSE_ID_EVENT_TYPES] assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] @@ -1038,9 +1001,7 @@ def _signed_thinking_and_text_in_one_chunk() -> tuple[ModelResponseStream, ...]: chunk: Final = _chunk("Hello!", finish_reason="stop") signed_block: Final = {"type": "thinking", "thinking": "", "signature": "sig_only"} delta: Final = Delta(**{**chunk.choices[0].delta.model_dump(), "thinking_blocks": [signed_block]}) - return ( - chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}), - ) + return (chunk.model_copy(update={"choices": [StreamingChoices(index=0, delta=delta, finish_reason="stop")]}),) StreamChunks = Callable[[], tuple[ModelResponseStream, ...]] @@ -1560,9 +1521,7 @@ async def test_anthropic_tool_search_call_streams_as_tool_search_call_item(sync_ model="claude-sonnet-5-5", object="chat.completion.chunk", choices=[ - StreamingChoices( - index=0, delta=delta, finish_reason="tool_calls" if index == len(deltas) - 1 else None - ) + StreamingChoices(index=0, delta=delta, finish_reason="tool_calls" if index == len(deltas) - 1 else None) ], ) for index, delta in enumerate(deltas) @@ -1972,9 +1931,7 @@ def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_b custom_llm_provider="anthropic", litellm_metadata={}, ) - relayed: Final = tuple( - itertools.takewhile(lambda event: not isinstance(event, OutputTextDeltaEvent), iterator) - ) + relayed: Final = tuple(itertools.takewhile(lambda event: not isinstance(event, OutputTextDeltaEvent), iterator)) assert relayed, "the stream must relay at least one chunk before the client goes away" partial: Final = iterator.partial_response_on_disconnect() @@ -2026,3 +1983,55 @@ async def test_streamed_refusal_deltas_become_one_refusal_part_of_the_message(sy assert [event.refusal for event in refusal_done] == [expected_content[0]["refusal"]] assert [item.model_dump()["content"] for item in message_done] == [expected_content] assert [item.model_dump()["content"] for item in finished.response.output] == [expected_content] + + +def _logprobs_chunk(content: str, logprob: float, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="gemini-3.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + logprobs=ChoiceLogprobs(content=[ChatCompletionTokenLogprob(token=content, logprob=logprob)]), + ) + ], + ) + + +def test_streamed_chat_logprobs_reach_the_text_delta_and_done_events() -> None: + events: Final = [ + event.model_dump() + for event in _build_iterator( + [_logprobs_chunk("Hello", -0.25), _logprobs_chunk(".", -0.5, finish_reason="stop")] + ) + ] + hello: Final = {"token": "Hello", "logprob": -0.25, "bytes": None, "top_logprobs": []} + period: Final = {"token": ".", "logprob": -0.5, "bytes": None, "top_logprobs": []} + + deltas: Final = [event for event in events if event["type"] == "response.output_text.delta"] + done: Final = next(event for event in events if event["type"] == "response.output_text.done") + part_done: Final = next(event for event in events if event["type"] == "response.content_part.done") + item_done: Final = next( + event for event in events if event["type"] == "response.output_item.done" and event["item"]["type"] == "message" + ) + + assert [event["logprobs"] for event in deltas] == [[hello], [period]] + assert done["logprobs"] == [hello, period] + assert part_done["part"]["logprobs"] == [hello, period] + assert item_done["item"]["content"][0]["logprobs"] == [hello, period] + + +def test_streamed_chat_without_logprobs_emits_text_events_without_a_logprobs_key() -> None: + events: Final = [ + event.model_dump() for event in _build_iterator([_chunk("Hello"), _chunk("!", finish_reason="stop")]) + ] + + text_events: Final = [ + event for event in events if event["type"] in {"response.output_text.delta", "response.output_text.done"} + ] + + assert text_events and all("logprobs" not in event for event in text_events)