diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index bd239922fd3..17094d18c85 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -144,6 +144,7 @@ _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) +_THOUGHT_SIGNATURES_ADAPTER: Final = TypeAdapter(list[str]) _RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) @@ -1421,15 +1422,70 @@ class LiteLLMCompletionResponsesConfig: # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: return [] + signature_fields: Final = ( + LiteLLMCompletionResponsesConfig._message_thought_signature_fields( + input_item.get("provider_specific_fields") + ) + if _input_item_role(input_item) == "assistant" + else None + ) + signed_text: Final = ( + LiteLLMCompletionResponsesConfig._signed_message_text(content) if signature_fields else None + ) return [ GenericChatCompletionMessage( role=_input_item_role(input_item), - content=LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( - content + content=( + signed_text + if isinstance(signed_text, str) + else LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content( + content + ) + ), + **( + MappingProxyType({"provider_specific_fields": signature_fields}) + if signature_fields + else MappingProxyType({}) ), ) ] + @staticmethod + def _signed_message_text(content: object) -> str | None: + if not isinstance(content, list): + return None + blocks: Final = _OBJECT_LIST_ADAPTER.validate_python(content) + if len(blocks) != 1 or not isinstance(blocks[0], Mapping): + return None + block: Final = _STR_KEY_DICT_ADAPTER.validate_python(blocks[0]) + text: Final = block.get("text") + return text if block.get("type") == "output_text" and isinstance(text, str) else None + + @staticmethod + def _with_message_thought_signatures( + item: GenericResponseOutputItem, + fields: object, + ) -> GenericResponseOutputItem: + signatures: Final = LiteLLMCompletionResponsesConfig._message_thought_signature_fields(fields) + return ( + item.model_copy(update=MappingProxyType({"provider_specific_fields": signatures})) if signatures else item + ) + + @staticmethod + def _message_thought_signature_fields( + fields: object, + ) -> dict[str, list[str]] | None: # mutable-ok: Gemini replay requires JSON dictionaries and arrays + if not isinstance(fields, Mapping): + return None + signatures: Final[object] = _STR_KEY_DICT_ADAPTER.validate_python(fields).get("thought_signatures") + if not signatures: + return None + try: + validated: Final = _THOUGHT_SIGNATURES_ADAPTER.validate_python(signatures, strict=True) + except ValidationError: + return None + return {"thought_signatures": validated} # mutable-ok: Gemini replay requires JSON arrays + @staticmethod def _reasoning_text_from_content(input_item: Mapping[str, object]) -> str | None: """ @@ -2719,18 +2775,21 @@ class LiteLLMCompletionResponsesConfig: message_output_items.extend(image_generation_items) elif choice.message.content is not None: message_output_items.append( - GenericResponseOutputItem( - type="message", - id=f"msg_{uuid.uuid4()}", - status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( - choice.finish_reason + LiteLLMCompletionResponsesConfig._with_message_thought_signatures( + GenericResponseOutputItem( + type="message", + id=f"msg_{uuid.uuid4()}", + status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status( + choice.finish_reason + ), + role=choice.message.role, + content=[ + LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text( + choice.message + ) + ], ), - role=choice.message.role, - content=[ - LiteLLMCompletionResponsesConfig._transform_chat_message_to_response_output_text( - choice.message - ) - ], + choice.message.provider_specific_fields, ) ) return message_output_items 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 7b9de4644b4..708c76a9998 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 @@ -5172,3 +5172,27 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool ) assert exc_info.value.status_code == 400 assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) + + +@pytest.mark.parametrize("signature", [None, "c2lnbmVkLXRleHQ="]) +def test_gemini_text_signature_survives_responses_tool_replay(signature: str | None) -> None: + from litellm.llms.vertex_ai.gemini.transformation import _gemini_convert_messages_with_history + + text: Final = " Let me check.\n" + fields: Final = {"thought_signatures": [signature]} if signature else None + response: Final = ModelResponse(choices=[Choices(message=Message(content=text, provider_specific_fields=fields))]) + output: Final = LiteLLMCompletionResponsesConfig._extract_message_output_items(response, response.choices) + replayed: Final = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=[ + {"role": "user", "content": "Check the weather"}, + output[0].model_dump(exclude_none=True), + {"type": "function_call", "call_id": "call_weather", "name": "weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "sunny"}, + ], + replay_reasoning=True, + ) + contents: Final = _gemini_convert_messages_with_history(replayed, model="gemini-2.5-flash") + parts: Final = contents[1]["parts"] + assert parts[0] == ({"text": text, "thoughtSignature": signature} if signature else {"text": text}) + assert parts[1]["function_call"] == {"name": "weather", "args": {}} + assert contents[2]["parts"][0]["function_response"]["response"] == {"content": "sunny"}