scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

fix(gemini): resend once with the placeholder signature when the other backend rejects a replayed thought signature (#49697)

Browse source 
* fix(gemini): resend once with the placeholder signature when the other backend rejects a replayed thought signature

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): gemini tool-call signature rejected by the backend is resent once with the placeholder

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(gemini): keep the placeholder resend inside the cache rebuild so either rejection order ends in an accepted body

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(gemini): resend a cached-prefix signature inline and cover native generateContent

When the rejected signature sits inside a cachedContents prefix, the
placeholder resend now carries the history inline, since replacing the
signature in the request body cannot reach the cache. The native
generateContent handlers get the same one-time resend.

* fix(gemini): send the history inline when Vertex rejects a signature while creating the cache

* fix(gemini): keep the context cache on a signature resend and go inline only when that is rejected too

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-02 17:24:55 -08:00 • committed by GitHub
parent 0bd677afef
commit 640a87f70d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 933 additions and 114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -1329,7 +1329,7 @@ def get_thought_signature_from_tool(tool: Mapping[str, object]) -> str | None:
_get_thought_signature_from_tool = get_thought_signature_from_tool
def _get_dummy_thought_signature() -> str:
def get_dummy_thought_signature() -> str:
"""Generate a dummy thought signature for models that require it.
This is used when transferring conversation history from older models
@ -1416,7 +1416,7 @@ def convert_to_gemini_tool_call_invoke(
# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures#parallel_function_calling_example
is_first_function_call = len(_parts_list) == 0
if not thought_signature and is_first_function_call and needs_dummy_signature:
thought_signature = _get_dummy_thought_signature()
thought_signature = get_dummy_thought_signature()
if thought_signature:
part_dict["thoughtSignature"] = thought_signature
@ -1440,7 +1440,7 @@ def convert_to_gemini_tool_call_invoke(
# If no signature found and model is gemini-3, use dummy signature
if not thought_signature and needs_dummy_signature:
thought_signature = _get_dummy_thought_signature()
thought_signature = get_dummy_thought_signature()
if thought_signature:
part_dict_function["thoughtSignature"] = thought_signature

2
litellm/llms/base_llm/google_genai/transformation.py
View file

@ -164,7 +164,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC):
tools: ToolConfigDict | None,
generate_content_config_dict: dict,
system_instruction: object | None = None,
) -> dict:
) -> dict[str, object]:
"""
Transform the request parameters for the generate content API.

96
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -455,6 +455,20 @@ def _google_genai_streaming_hidden_params(
}
def _log_generate_content_request(
logging_obj: LiteLLMLoggingObj,
contents: object,
api_base: str,
headers: Mapping[str, object],
body: Mapping[str, object],
) -> None:
logging_obj.pre_call(
input=contents,
api_key="",
additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers},
)
@lru_cache(maxsize=None)
def _responses_api_optional_request_param_names() -> frozenset[str]:
return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys())
@ -12443,6 +12457,10 @@ class BaseLLMHTTPHandler:
from litellm.google_genai.streaming_iterator import (
GoogleGenAIGenerateContentStreamingIterator,
)
from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import (
post_with_accepted_thought_signatures_sync,
with_placeholder_native_thought_signatures,
)
if _is_async:
return self.async_generate_content_handler(
@ -12494,27 +12512,19 @@ class BaseLLMHTTPHandler:
if extra_body:
data.update(extra_body)
## LOGGING
logging_obj.pre_call(
input=contents,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers)
log_request(data)
def post(body: Mapping[str, object]) -> httpx.Response:
return sync_httpx_client.post(
url=api_base, headers=headers, json=dict(body), timeout=timeout, stream=stream
)
try:
sent_body, response = post_with_accepted_thought_signatures_sync(
data, with_placeholder_native_thought_signatures, None, post, log_request
)
if stream:
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout,
stream=True,
)
# Return streaming iterator
return GoogleGenAIGenerateContentStreamingIterator(
response=response,
model=model,
@ -12522,7 +12532,7 @@ class BaseLLMHTTPHandler:
generate_content_provider_config=generate_content_provider_config,
litellm_metadata=litellm_metadata or {},
custom_llm_provider=custom_llm_provider,
request_body=data,
request_body=dict(sent_body),
hidden_params=_google_genai_streaming_hidden_params(
api_base=api_base,
litellm_params=litellm_params,
@ -12530,13 +12540,6 @@ class BaseLLMHTTPHandler:
response_headers=response.headers,
),
)
else:
response = sync_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise self._handle_error(
e=e,
@ -12574,6 +12577,10 @@ class BaseLLMHTTPHandler:
from litellm.google_genai.streaming_iterator import (
AsyncGoogleGenAIGenerateContentStreamingIterator,
)
from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import (
post_with_accepted_thought_signatures_async,
with_placeholder_native_thought_signatures,
)
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
@ -12609,27 +12616,19 @@ class BaseLLMHTTPHandler:
if extra_body:
data.update(extra_body)
## LOGGING
logging_obj.pre_call(
input=contents,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
log_request: Final = partial(_log_generate_content_request, logging_obj, contents, api_base, headers)
log_request(data)
async def post(body: Mapping[str, object]) -> httpx.Response:
return await async_httpx_client.post(
url=api_base, headers=headers, json=dict(body), timeout=timeout, stream=stream
)
try:
sent_body, response = await post_with_accepted_thought_signatures_async(
data, with_placeholder_native_thought_signatures, None, post, log_request
)
if stream:
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout,
stream=True,
)
# Return async streaming iterator
return AsyncGoogleGenAIGenerateContentStreamingIterator(
response=response,
model=model,
@ -12637,7 +12636,7 @@ class BaseLLMHTTPHandler:
generate_content_provider_config=generate_content_provider_config,
litellm_metadata=litellm_metadata or {},
custom_llm_provider=custom_llm_provider,
request_body=data,
request_body=dict(sent_body),
hidden_params=_google_genai_streaming_hidden_params(
api_base=api_base,
litellm_params=litellm_params,
@ -12645,13 +12644,6 @@ class BaseLLMHTTPHandler:
response_headers=response.headers,
),
)
else:
response = await async_httpx_client.post(
url=api_base,
headers=headers,
json=data,
timeout=timeout,
)
except Exception as e:
raise self._handle_error(
e=e,

27
litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py
View file

@ -1,4 +1,4 @@
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Mapping, Sequence
from copy import deepcopy
from functools import partial
from typing import Final, TypeVar
@ -17,16 +17,27 @@ Built: Final = TypeVar("Built")
Sent: Final = TypeVar("Sent")
def cache_rebuild(
def _build_from_copies(
build: Callable[..., Built], messages: Sequence[AllMessageValues], optional_params: Mapping[str, object]
) -> Built:
return build(messages=deepcopy(list(messages)), optional_params=independent_snapshot(dict(optional_params)))
def cache_rebuilds(
build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict
) -> Callable[[], Built] | None:
"""Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``."""
) -> tuple[Callable[[], Built], Callable[[], Built]] | tuple[None, None]:
"""Returns one build with a fresh cache and one with the history inline, both over a snapshot of the inputs
taken before the first transform, which pops keys from ``optional_params`` and ``messages``. Each call
copies the snapshot again, since every transform consumes the inputs it is given."""
if not any(map(is_cached_message, messages)):
return None
return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params))
return None, None
from_snapshot: Final = partial(
_build_from_copies, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params)
)
return partial(from_snapshot, build), partial(from_snapshot, partial(build, context_caching=False))
def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]:
def status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]:
match error:
case httpx.HTTPStatusError():
return error.response.status_code, error.response.text
@ -35,7 +46,7 @@ def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tupl
def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool:
return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error))
return body.get("cachedContent") is not None and names_rejected_cache(*status_and_message(error))
async def post_with_fresh_cache_async(

6
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -425,6 +425,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_auth_header: str | None,
extra_headers: dict | None = None,
cached_content: str | None = None,
context_caching: bool = True,
) -> tuple[list[AllMessageValues], dict, str | None]:
"""
Receives
@ -436,7 +437,7 @@ class ContextCachingEndpoints(VertexBase):
Follows - https://ai.google.dev/api/caching#request-body
"""
if cached_content is not None:
if cached_content is not None or not context_caching:
return messages, optional_params, cached_content
cached_messages, non_cached_messages = separate_cached_messages(messages=messages)
@ -573,6 +574,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_auth_header: str | None,
extra_headers: dict | None = None,
cached_content: str | None = None,
context_caching: bool = True,
) -> tuple[list[AllMessageValues], dict, str | None]:
"""
Receives
@ -584,7 +586,7 @@ class ContextCachingEndpoints(VertexBase):
Follows - https://ai.google.dev/api/caching#request-body
"""
if cached_content is not None:
if cached_content is not None or not context_caching:
return messages, optional_params, cached_content
cached_messages, non_cached_messages = separate_cached_messages(messages=messages)

235
litellm/llms/vertex_ai/gemini/rejected_thought_signature_retry.py Normal file
View file

@ -0,0 +1,235 @@
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from typing import Final, TypeVar
import httpx
from pydantic import TypeAdapter, ValidationError
from litellm import verbose_logger
from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.vertex_ai import ContentType, PartType, RequestBody
from ..context_caching.rejected_cache_retry import status_and_message
Body: Final = TypeVar("Body", bound=Mapping[str, object])
Sent: Final = TypeVar("Sent")
_REJECTED_SIGNATURE_MARKER: Final = "thought signature"
_SIGNATURE_KEYS: Final = frozenset({"thoughtSignature", "thought_signature"})
_ENTRIES: Final = TypeAdapter(list[dict[str, object]])
def names_rejected_thought_signature(status_code: int, message: str) -> bool:
"""A signature only validates on the backend that minted it. AI Studio answers a replayed Vertex signature
with ``400 Corrupted thought signature.`` and Vertex answers an AI Studio one with ``400 Invalid thought
signature.``"""
return status_code == 400 and _REJECTED_SIGNATURE_MARKER in message.lower()
def _is_rejected_signature(error: httpx.HTTPStatusError | BaseLLMException) -> bool:
return names_rejected_thought_signature(*status_and_message(error))
def _entries(value: object) -> Sequence[Mapping[str, object]]:
try:
return _ENTRIES.validate_python(value)
except ValidationError:
return []
def _foreign_signatures(body: Mapping[str, object], placeholder: str) -> Iterator[object]:
for content in _entries(body.get("contents")):
for part in _entries(content.get("parts")):
yield from (
value for key, value in part.items() if key in _SIGNATURE_KEYS and value and value != placeholder
)
def _foreign_signature_count(body: Mapping[str, object]) -> int:
return sum(1 for _ in _foreign_signatures(body, get_dummy_thought_signature()))
def _part_with_placeholder(part: PartType, placeholder: str) -> PartType:
if not part.get("thoughtSignature"):
return part
return {**part, "thoughtSignature": placeholder}
def _content_with_placeholders(content: ContentType, placeholder: str) -> ContentType:
return {**content, "parts": [_part_with_placeholder(part, placeholder) for part in content["parts"]]}
def with_placeholder_thought_signatures(body: RequestBody) -> RequestBody | None:
"""Swaps every thought signature in the history for Google's documented skip-validator placeholder, or
returns ``None`` when the body carries nothing the backend could have rejected."""
placeholder: Final = get_dummy_thought_signature()
if next(_foreign_signatures(body, placeholder), None) is None:
return None
return {**body, "contents": [_content_with_placeholders(content, placeholder) for content in body["contents"]]}
def _native_part_with_placeholder(part: Mapping[str, object], placeholder: str) -> Mapping[str, object]:
return {key: placeholder if key in _SIGNATURE_KEYS and value else value for key, value in part.items()}
def _native_content_with_placeholders(content: Mapping[str, object], placeholder: str) -> Mapping[str, object]:
parts: Final = _entries(content.get("parts"))
if not parts:
return content
return {**content, "parts": [_native_part_with_placeholder(part, placeholder) for part in parts]}
def with_placeholder_native_thought_signatures(body: Mapping[str, object]) -> Mapping[str, object] | None:
"""The same swap on a ``generateContent`` body as the client sent it, whose parts may spell the key
``thoughtSignature`` or ``thought_signature``."""
placeholder: Final = get_dummy_thought_signature()
if next(_foreign_signatures(body, placeholder), None) is None:
return None
return {
**body,
"contents": [
_native_content_with_placeholders(content, placeholder) for content in _entries(body.get("contents"))
],
}
async def _inline_history_async(body: Body, uncached: Callable[[], Awaitable[Body]] | None) -> Body | None:
if uncached is None or body.get("cachedContent") is None:
return None
inline: Final = await uncached()
return inline if _foreign_signature_count(inline) > _foreign_signature_count(body) else None
def _inline_history_sync(body: Body, uncached: Callable[[], Body] | None) -> Body | None:
if uncached is None or body.get("cachedContent") is None:
return None
inline: Final = uncached()
return inline if _foreign_signature_count(inline) > _foreign_signature_count(body) else None
def _resend_body(signed: Body, placeholders: Callable[[Body], Body | None]) -> Body | None:
resend_body: Final = placeholders(signed)
if resend_body is not None:
verbose_logger.warning(
"Gemini rejected a replayed thought signature, resending with the skip_thought_signature_validator"
" placeholder in place of the replayed signatures"
)
return resend_body
async def build_with_accepted_thought_signatures_async(
build: Callable[[], Awaitable[Body]],
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Awaitable[Body]] | None,
) -> Body:
"""Vertex validates the signatures in a cache's contents when it creates the cache, so a rejected one there
sends the history inline with the placeholder instead."""
try:
return await build()
except BaseLLMException as error:
if uncached is None or not _is_rejected_signature(error):
raise
inline_body: Final = _resend_body(await uncached(), placeholders)
if inline_body is None:
raise
return inline_body
def build_with_accepted_thought_signatures_sync(
build: Callable[[], Body],
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Body] | None,
) -> Body:
try:
return build()
except BaseLLMException as error:
if uncached is None or not _is_rejected_signature(error):
raise
inline_body: Final = _resend_body(uncached(), placeholders)
if inline_body is None:
raise
return inline_body
async def _post_inline_async(
body: Body,
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Awaitable[Body]] | None,
post: Callable[[Body], Awaitable[Sent]],
log_request: Callable[[Body], None],
rejection: httpx.HTTPStatusError | BaseLLMException,
) -> tuple[Body, Sent]:
inline: Final = await _inline_history_async(body, uncached)
resend_body: Final = None if inline is None else _resend_body(inline, placeholders)
if resend_body is None:
raise rejection
log_request(resend_body)
return resend_body, await post(resend_body)
def _post_inline_sync(
body: Body,
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Body] | None,
post: Callable[[Body], Sent],
log_request: Callable[[Body], None],
rejection: httpx.HTTPStatusError | BaseLLMException,
) -> tuple[Body, Sent]:
inline: Final = _inline_history_sync(body, uncached)
resend_body: Final = None if inline is None else _resend_body(inline, placeholders)
if resend_body is None:
raise rejection
log_request(resend_body)
return resend_body, post(resend_body)
async def post_with_accepted_thought_signatures_async(
body: Body,
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Awaitable[Body]] | None,
post: Callable[[Body], Awaitable[Sent]],
log_request: Callable[[Body], None],
) -> tuple[Body, Sent]:
"""Innermost retry around the raw post, so a cache rebuild (which re-transforms the messages and brings
the rejected signatures back) is itself wrapped by it. The resend keeps the cache and swaps the signatures
the body carries, and sends the history inline only when that is rejected too and the cache holds one.
Returns the body the backend accepted."""
try:
return body, await post(body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if not _is_rejected_signature(error):
raise
resend_body: Final = _resend_body(body, placeholders)
if resend_body is None:
return await _post_inline_async(body, placeholders, uncached, post, log_request, error)
log_request(resend_body)
try:
return resend_body, await post(resend_body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if not _is_rejected_signature(error):
raise
return await _post_inline_async(body, placeholders, uncached, post, log_request, error)
def post_with_accepted_thought_signatures_sync(
body: Body,
placeholders: Callable[[Body], Body | None],
uncached: Callable[[], Body] | None,
post: Callable[[Body], Sent],
log_request: Callable[[Body], None],
) -> tuple[Body, Sent]:
try:
return body, post(body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if not _is_rejected_signature(error):
raise
resend_body: Final = _resend_body(body, placeholders)
if resend_body is None:
return _post_inline_sync(body, placeholders, uncached, post, log_request, error)
log_request(resend_body)
try:
return resend_body, post(resend_body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if not _is_rejected_signature(error):
raise
return _post_inline_sync(body, placeholders, uncached, post, log_request, error)

4
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -1256,6 +1256,7 @@ def sync_transform_request_body(
vertex_project: str | None,
vertex_location: str | None,
vertex_auth_header: str | None,
context_caching: bool = True,
) -> RequestBody:
from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints
@ -1275,6 +1276,7 @@ def sync_transform_request_body(
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
context_caching=context_caching,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
@ -1390,6 +1392,7 @@ async def async_transform_request_body(
vertex_project: str | None,
vertex_location: str | None,
vertex_auth_header: str | None,
context_caching: bool = True,
) -> RequestBody:
from ..context_caching.vertex_ai_context_caching import ContextCachingEndpoints
@ -1409,6 +1412,7 @@ async def async_transform_request_body(
timeout=timeout,
extra_headers=extra_headers,
cached_content=optional_params.pop("cached_content", None),
context_caching=context_caching,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,

99
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -104,7 +104,7 @@ from ..common_utils import (
)
from ..context_caching.known_cached_contents import known_cached_contents
from ..context_caching.rejected_cache_retry import (
cache_rebuild,
cache_rebuilds,
post_with_fresh_cache_async,
post_with_fresh_cache_sync,
)
@ -135,6 +135,13 @@ from .long_decode import (
with_initial_prompt_tokens,
with_spans_shifted,
)
from .rejected_thought_signature_retry import (
build_with_accepted_thought_signatures_async,
build_with_accepted_thought_signatures_sync,
post_with_accepted_thought_signatures_async,
post_with_accepted_thought_signatures_sync,
with_placeholder_thought_signatures,
)
from .transformation import (
async_transform_request_body,
gemini_convert_messages_with_history,
@ -3500,23 +3507,17 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
rebuild: Final = cache_rebuild(
partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
request_body: Final = await async_transform_request_body(
transform: Final = partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
)
rebuild, uncached = cache_rebuilds(transform, messages, optional_params)
request_body: Final = await build_with_accepted_thought_signatures_async(
transform, with_placeholder_thought_signatures, uncached
)
## LOGGING
log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers)
@ -3545,7 +3546,14 @@ class VertexLLM(VertexBase):
async def post(body: RequestBody) -> ModelResponseIterator:
return await open_stream(client=client, data=json.dumps(body), request_body=body)
return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1]
accepted: Final = partial(
post_with_accepted_thought_signatures_async,
placeholders=with_placeholder_thought_signatures,
uncached=uncached,
post=post,
log_request=log_request,
)
return (await post_with_fresh_cache_async(request_body, rebuild, accepted, log_request))[1][1]
streaming_response: Final = CustomStreamWrapper(
completion_stream=None,
@ -3612,23 +3620,17 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
rebuild: Final = cache_rebuild(
partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
request_body: Final = await async_transform_request_body(
transform: Final = partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
)
rebuild, uncached = cache_rebuilds(transform, messages, optional_params)
request_body: Final = await build_with_accepted_thought_signatures_async(
transform, with_placeholder_thought_signatures, uncached
)
_async_client_params: Final = {}
if timeout:
@ -3669,7 +3671,14 @@ class VertexLLM(VertexBase):
)
return response
sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request)
accepted: Final = partial(
post_with_accepted_thought_signatures_async,
placeholders=with_placeholder_thought_signatures,
uncached=uncached,
post=post,
log_request=log_request,
)
_, (sent_body, first_response) = await post_with_fresh_cache_async(request_body, rebuild, accepted, log_request)
return _transform_long_decode_hops(
model=model,
hops=await collect_hops_async(parse_body(first_response), sent_body, post),
@ -3801,23 +3810,17 @@ class VertexLLM(VertexBase):
)
## TRANSFORMATION ##
rebuild: Final = cache_rebuild(
partial(
sync_transform_request_body,
**transform_request_params,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
data: Final = sync_transform_request_body(
transform: Final = partial(
sync_transform_request_body,
**transform_request_params,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
)
rebuild, uncached = cache_rebuilds(transform, messages, optional_params)
data: Final = build_with_accepted_thought_signatures_sync(
transform, with_placeholder_thought_signatures, uncached
)
## LOGGING
log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers)
@ -3845,7 +3848,14 @@ class VertexLLM(VertexBase):
def post(body: RequestBody) -> ModelResponseIterator:
return open_stream(client=client, data=json.dumps(body), request_body=body)
return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1]
accepted: Final = partial(
post_with_accepted_thought_signatures_sync,
placeholders=with_placeholder_thought_signatures,
uncached=uncached,
post=post,
log_request=log_request,
)
return post_with_fresh_cache_sync(data, rebuild, accepted, log_request)[1][1]
streaming_response: Final = CustomStreamWrapper(
completion_stream=None,
@ -3889,7 +3899,14 @@ class VertexLLM(VertexBase):
)
return response
sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request)
accepted: Final = partial(
post_with_accepted_thought_signatures_sync,
placeholders=with_placeholder_thought_signatures,
uncached=uncached,
post=post,
log_request=log_request,
)
_, (sent_body, first_response) = post_with_fresh_cache_sync(data, rebuild, accepted, log_request)
return _transform_long_decode_hops(
model=model,
hops=collect_hops_sync(parse_body(first_response), sent_body, post),

66
tests/integration/providers/test_gemini_thinking_replay_wire.py
View file

@ -2,6 +2,7 @@ import base64
import itertools
import json
import os
import re
import uuid
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Set as AbstractSet
@ -679,9 +680,70 @@ def test_gemini_own_tool_call_signature_is_still_replayed_on_the_function_call_p
assert _spend_row(rounds[1])["status"] == "success"
def test_responses_previous_response_id_replays_gemini_session_history(
gateway: Gateway, tmp_path: Path
_PLACEHOLDER_SIGNATURE: Final = base64.b64encode(b"skip_thought_signature_validator").decode()
_ISSUED_ID: Final = re.compile(r'"id":\s*"([^"]+)"')
_PLACEHOLDER_TOOL_TURN: Final = _model_turn(
_thought("I should count."),
{"function_call": {"name": "count_fruit", "args": {"kind": "apple"}}, "thoughtSignature": _PLACEHOLDER_SIGNATURE},
)
# Wording AI Studio answers with when a Vertex-minted signature is replayed to it, captured live on 2026-10-27
_FOREIGN_SIGNATURE_REJECTED: Final = Reply(
status=400,
body=json.dumps(
{"error": {"code": 400, "message": "Corrupted thought signature.", "status": "INVALID_ARGUMENT"}}
).encode(),
)
@pytest.mark.parametrize("stream", (False, True))
@pytest.mark.parametrize("endpoint", _ENDPOINTS)
def test_gemini_tool_call_signature_the_backend_rejects_is_resent_once_with_the_placeholder(
gateway: Gateway, endpoint: Endpoint, stream: bool
) -> None:
rounds: Final = (f"gemini-reply-{uuid.uuid4().hex}", f"gemini-reply-{uuid.uuid4().hex}")
calls: Final = itertools.count()
def respond(request: Request) -> Reply:
match next(calls):
case 0:
return _reply(rounds[0], stream=False, parts=_GEMINI_TOOL_REPLY_PARTS)
case 1:
return _FOREIGN_SIGNATURE_REJECTED
case _:
return _reply(rounds[1], stream=stream, parts=({"text": "Twelve apples."},))
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = _register(gateway, scenario, "gemini", wire)
first: Final = _tool_round_one(endpoint, model)
headers: Final = {"anthropic-version": "2023-06-01"}
opened: Final = gateway.request("POST", _path(endpoint), first, headers=headers)
assert opened.status_code == 200, opened.text
second: Final = {
**_tool_round_two(endpoint, first, _JSON_OBJECT.validate_json(opened.content)),
"stream": stream,
}
closed: Final = gateway.request("POST", _path(endpoint), second, headers=headers)
assert closed.status_code == 200, closed.text
assert "Twelve apples." in closed.text, closed.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [
("POST", _target("gemini", False)),
("POST", _target("gemini", stream)),
("POST", _target("gemini", stream)),
], received
rejected, resent = (_JSON_OBJECT.validate_json(request.body) for request in received[1:])
assert rejected == {
**_expected_body(endpoint, _user("Count the apples."), _GEMINI_TOOL_REPLAYED_TURN, _FUNCTION_RESPONSE_TURN),
**_DECLARATIONS,
}, rejected
assert resent == {
**_expected_body(endpoint, _user("Count the apples."), _PLACEHOLDER_TOOL_TURN, _FUNCTION_RESPONSE_TURN),
**_DECLARATIONS,
}, resent
assert _spend_row(rounds[1], *_ISSUED_ID.findall(closed.text))["status"] == "success"
def test_responses_previous_response_id_replays_gemini_session_history(gateway: Gateway, tmp_path: Path) -> None:
first_prompt: Final = f"gemini-session-first-{uuid.uuid4().hex}"
second_prompt: Final = f"gemini-session-second-{uuid.uuid4().hex}"
first_answer: Final = f"gemini-answer-first-{uuid.uuid4().hex}"

496
tests/unit/llms/vertex_ai/gemini/test_rejected_thought_signature_retry.py Normal file
View file

@ -0,0 +1,496 @@
import base64
import json
from collections.abc import Callable, Mapping, Sequence
from itertools import chain
from typing import Final
import httpx
import pytest
import litellm
from litellm.google_genai import agenerate_content, agenerate_content_stream, generate_content, generate_content_stream
from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.gemini.rejected_thought_signature_retry import with_placeholder_thought_signatures
from litellm.types.llms.vertex_ai import RequestBody
from litellm.types.utils import Choices, ModelResponse, ModelResponseStream
from litellm.utils import CustomStreamWrapper
MODEL: Final = "gemini/gemini-3.8-flash"
VERTEX_MINTED: Final = base64.b64encode(b"minted-by-vertex").decode()
VERTEX_MINTED_TEXT: Final = base64.b64encode(b"minted-by-vertex-for-text").decode()
AI_STUDIO_MINTED: Final = base64.b64encode(b"minted-by-ai-studio").decode()
HISTORY_WITH_FOREIGN_SIGNATURES: Final = (
{"role": "user", "content": "What is the weather in Paris?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
"provider_specific_fields": {"thought_signature": VERTEX_MINTED},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "18C and sunny"},
{
"role": "assistant",
"content": "It is 18C and sunny in Paris.",
"provider_specific_fields": {"thought_signatures": [VERTEX_MINTED_TEXT]},
},
{"role": "user", "content": "And in Berlin?"},
)
HISTORY_WITHOUT_SIGNATURES: Final = (
{"role": "user", "content": "What is the weather in Paris?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": "18C and sunny"},
{"role": "user", "content": "And in Berlin?"},
)
TOOLS: Final = (
{
"type": "function",
"function": {
"name": "get_weather",
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
},
},
)
ANSWER: Final = {
"candidates": [{"content": {"role": "model", "parts": [{"text": "Berlin is 12C."}]}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 40, "candidatesTokenCount": 6, "totalTokenCount": 46},
}
# Wording observed from Google on 2026-10-08 when a signature minted by the other backend is replayed
AI_STUDIO_REJECTION: Final = "Corrupted thought signature."
VERTEX_REJECTION: Final = "Invalid thought signature."
UNRELATED_REJECTION: Final = 'Invalid JSON payload received. Unknown name "foo".'
def _rejected(message: str) -> httpx.Response:
return httpx.Response(400, json={"error": {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}})
def _sse(*chunks: Mapping[str, object]) -> httpx.Response:
return httpx.Response(
200,
content="".join(f"data: {json.dumps(chunk)}\r\n\r\n" for chunk in chunks).encode(),
headers={"content-type": "text/event-stream"},
)
def _replay(
answers: Sequence[httpx.Response], sent: list[dict[str, object]]
) -> Callable[[httpx.Request], httpx.Response]:
def handler(request: httpx.Request) -> httpx.Response:
sent.append(json.loads(request.content))
return answers[len(sent) - 1]
return handler
def _signatures_of(body: Mapping[str, object], key: str = "thoughtSignature") -> list[str]:
contents: Final = body["contents"]
assert isinstance(contents, list)
parts: Final = chain.from_iterable(content["parts"] for content in contents)
return [part[key] for part in parts if key in part]
def _text_of(response: ModelResponse) -> str:
choice: Final = response.choices[0]
assert isinstance(choice, Choices), choice
assert isinstance(choice.message.content, str), choice.message
return choice.message.content
def _delta_text_of(chunk: object) -> str:
assert isinstance(chunk, ModelResponseStream), chunk
return chunk.choices[0].delta.content or ""
async def _complete(
messages: Sequence[Mapping[str, object]],
backend: Callable[[httpx.Request], httpx.Response],
is_async: bool,
stream: bool,
) -> str:
kwargs: Final = {
"model": MODEL,
"messages": list(messages),
"tools": list(TOOLS),
"stream": stream,
"api_key": "unit-test-key",
}
response: Final = (
await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(backend)))
if is_async
else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(backend))))
)
if isinstance(response, ModelResponse):
return _text_of(response)
assert isinstance(response, CustomStreamWrapper), response
if is_async:
return "".join([_delta_text_of(chunk) async for chunk in response])
return "".join(_delta_text_of(chunk) for chunk in response)
def _answer(stream: bool) -> httpx.Response:
return _sse(ANSWER) if stream else httpx.Response(200, json=ANSWER)
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("rejection", [AI_STUDIO_REJECTION, VERTEX_REJECTION])
@pytest.mark.asyncio
async def test_signature_rejected_by_the_other_backend_is_resent_once_with_the_placeholder(
is_async: bool, stream: bool, rejection: str
):
sent: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
HISTORY_WITH_FOREIGN_SIGNATURES, _replay((_rejected(rejection), _answer(stream)), sent), is_async, stream
)
assert answer == "Berlin is 12C."
assert len(sent) == 2
first_signatures, second_signatures = _signatures_of(sent[0]), _signatures_of(sent[1])
assert set(first_signatures) == {VERTEX_MINTED, VERTEX_MINTED_TEXT}
assert second_signatures == [get_dummy_thought_signature()] * len(first_signatures)
assert {**sent[1], "contents": None} == {**sent[0], "contents": None}
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_placeholder_rejected_too_is_raised_after_the_second_post(is_async: bool, stream: bool):
sent: Final[list[dict[str, object]]] = []
with pytest.raises(litellm.BadRequestError, match="thought signature"):
await _complete(
HISTORY_WITH_FOREIGN_SIGNATURES,
_replay((_rejected(AI_STUDIO_REJECTION), _rejected(AI_STUDIO_REJECTION)), sent),
is_async,
stream,
)
assert len(sent) == 2
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_unrelated_400_is_raised_without_a_resend(is_async: bool, stream: bool):
sent: Final[list[dict[str, object]]] = []
with pytest.raises(litellm.BadRequestError, match="Unknown name"):
await _complete(HISTORY_WITH_FOREIGN_SIGNATURES, _replay((_rejected(UNRELATED_REJECTION),), sent), is_async, stream)
assert len(sent) == 1
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_400_on_a_history_carrying_only_the_placeholder_is_raised_without_a_resend(
is_async: bool, stream: bool
):
sent: Final[list[dict[str, object]]] = []
with pytest.raises(litellm.BadRequestError, match="thought signature"):
await _complete(HISTORY_WITHOUT_SIGNATURES, _replay((_rejected(AI_STUDIO_REJECTION),), sent), is_async, stream)
assert len(sent) == 1
assert set(_signatures_of(sent[0])) <= {get_dummy_thought_signature()}
def test_placeholder_body_keeps_the_original_untouched():
body: Final[RequestBody] = {
"contents": [
{"role": "model", "parts": [{"text": "hi", "thoughtSignature": VERTEX_MINTED}, {"text": "unsigned"}]},
{"role": "user", "parts": [{"text": "again"}]},
],
"generationConfig": {"temperature": 0.1},
}
resend: Final = with_placeholder_thought_signatures(body)
assert resend is not None
assert _signatures_of(resend) == [get_dummy_thought_signature()]
assert _signatures_of(body) == [VERTEX_MINTED]
assert resend.get("generationConfig") == body.get("generationConfig")
def test_body_without_a_foreign_signature_has_nothing_to_resend():
body: Final[RequestBody] = {
"contents": [
{"role": "model", "parts": [{"text": "hi", "thoughtSignature": get_dummy_thought_signature()}]},
{"role": "user", "parts": [{"text": "again"}]},
]
}
assert with_placeholder_thought_signatures(body) is None
CACHE_CONTROL: Final = {"type": "ephemeral"}
def _cached_history(
tag: str, cached_signature: str = VERTEX_MINTED, later_signatures: tuple[str, ...] = ()
) -> tuple[Mapping[str, object], ...]:
later_answer: Final = (
{"role": "assistant", "content": "Noted.", "provider_specific_fields": {"thought_signatures": [*later_signatures]}}
if later_signatures
else {"role": "assistant", "content": "Noted."}
)
prefix: Final = " ".join(f"Reference paragraph {i} for {tag}: weather values are in Celsius." for i in range(600))
return (
{"role": "system", "content": [{"type": "text", "text": prefix, "cache_control": CACHE_CONTROL}]},
{
"role": "user",
"content": [{"type": "text", "text": "What is the weather in Paris?", "cache_control": CACHE_CONTROL}],
},
{
"role": "assistant",
"content": [{"type": "text", "text": "Calling the tool.", "cache_control": CACHE_CONTROL}],
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
"provider_specific_fields": {"thought_signature": cached_signature},
}
],
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": [{"type": "text", "text": "18C and sunny", "cache_control": CACHE_CONTROL}],
},
{"role": "user", "content": [{"type": "text", "text": "Got it.", "cache_control": CACHE_CONTROL}]},
later_answer,
{"role": "user", "content": "And in Berlin?"},
)
def _google_ai_studio(
stream: bool, generated: list[dict[str, object]], minted: frozenset[str] = frozenset()
) -> Callable[[httpx.Request], httpx.Response]:
"""Rejects a generateContent call whose own contents or referenced cache carry a signature AI Studio did not mint,
the way the live backend answered on 2026-11-02."""
caches: Final[dict[str, object]] = {}
def handler(request: httpx.Request) -> httpx.Response:
if request.url.path.endswith("/cachedContents"):
if request.method == "GET":
return httpx.Response(200, json={})
caches["cachedContents/prefix"] = json.loads(request.content)
return httpx.Response(
200, json={"name": "cachedContents/prefix", "usageMetadata": {"totalTokenCount": 9000}}
)
body: Final = json.loads(request.content)
generated.append(body)
cache: Final = caches.get(body.get("cachedContent"), {"contents": []})
assert isinstance(cache, dict)
signatures: Final = {*_signatures_of(body), *_signatures_of(cache)}
if signatures - {get_dummy_thought_signature(), *minted}:
return _rejected(AI_STUDIO_REJECTION)
return _answer(stream)
return handler
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_inside_a_cached_prefix_is_resent_inline_with_the_placeholder(is_async: bool, stream: bool):
generated: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
_cached_history(f"async={is_async} stream={stream}"),
_google_ai_studio(stream, generated),
is_async,
stream,
)
assert answer == "Berlin is 12C."
assert len(generated) == 2
assert generated[0]["cachedContent"] == "cachedContents/prefix"
assert _signatures_of(generated[0]) == []
assert "cachedContent" not in generated[1]
assert _signatures_of(generated[1]) == [get_dummy_thought_signature()]
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_after_a_cache_the_backend_accepts_is_resent_with_the_cache_kept(is_async: bool, stream: bool):
generated: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
_cached_history(
f"own cache async={is_async} stream={stream}", AI_STUDIO_MINTED, later_signatures=(VERTEX_MINTED_TEXT,)
),
_google_ai_studio(stream, generated, minted=frozenset({AI_STUDIO_MINTED})),
is_async,
stream,
)
assert answer == "Berlin is 12C."
assert [body.get("cachedContent") for body in generated] == ["cachedContents/prefix"] * 2
assert _signatures_of(generated[0]) == [VERTEX_MINTED_TEXT]
assert _signatures_of(generated[1]) == [get_dummy_thought_signature()]
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_inside_and_after_the_cached_prefix_goes_inline_once_the_cache_kept_resend_is_rejected(
is_async: bool, stream: bool
):
generated: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
_cached_history(f"both async={is_async} stream={stream}", later_signatures=(VERTEX_MINTED_TEXT,)),
_google_ai_studio(stream, generated),
is_async,
stream,
)
assert answer == "Berlin is 12C."
assert [body.get("cachedContent") for body in generated] == ["cachedContents/prefix"] * 2 + [None]
assert _signatures_of(generated[1]) == [get_dummy_thought_signature()]
assert _signatures_of(generated[2]) == [get_dummy_thought_signature()] * 2
def _backend_validating_cached_signatures(
stream: bool, created: list[dict[str, object]], generated: list[dict[str, object]]
) -> Callable[[httpx.Request], httpx.Response]:
"""Rejects creating a cache whose contents carry a signature it did not mint, the way Vertex answered on
2026-11-02."""
def handler(request: httpx.Request) -> httpx.Response:
body: Final = json.loads(request.content) if request.content else {}
if request.url.path.endswith("/cachedContents"):
if request.method == "GET":
return httpx.Response(200, json={})
created.append(body)
else:
generated.append(body)
if set(_signatures_of(body)) - {get_dummy_thought_signature()}:
return _rejected(VERTEX_REJECTION)
if request.url.path.endswith("/cachedContents"):
return httpx.Response(
200, json={"name": "cachedContents/prefix", "usageMetadata": {"totalTokenCount": 9000}}
)
return _answer(stream)
return handler
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_rejected_while_creating_the_cache_is_sent_inline_with_the_placeholder(
is_async: bool, stream: bool
):
created: Final[list[dict[str, object]]] = []
generated: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
_cached_history(f"cache creation async={is_async} stream={stream}"),
_backend_validating_cached_signatures(stream, created, generated),
is_async,
stream,
)
assert answer == "Berlin is 12C."
assert [_signatures_of(cache) for cache in created] == [[VERTEX_MINTED]]
assert len(generated) == 1
assert "cachedContent" not in generated[0]
assert _signatures_of(generated[0]) == [get_dummy_thought_signature()]
NATIVE_QUESTION: Final = {"role": "user", "parts": [{"text": "What is the weather in Paris?"}]}
NATIVE_TOOL_RESULT: Final = {
"role": "user",
"parts": [{"functionResponse": {"name": "get_weather", "response": {"result": "18C and sunny"}}}],
}
def _native_contents(signature_key: str) -> list[Mapping[str, object]]:
return [
NATIVE_QUESTION,
{
"role": "model",
"parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}, signature_key: VERTEX_MINTED}],
},
NATIVE_TOOL_RESULT,
]
async def _generate(
contents: list[Mapping[str, object]],
backend: Callable[[httpx.Request], httpx.Response],
is_async: bool,
stream: bool,
) -> str:
kwargs: Final = {"model": MODEL, "contents": contents, "api_key": "unit-test-key"}
sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(backend)))
async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(backend))
match is_async, stream:
case True, True:
chunks = await agenerate_content_stream(**kwargs, client=async_client)
return b"".join([chunk async for chunk in chunks]).decode()
case True, False:
return json.dumps((await agenerate_content(**kwargs, client=async_client)).model_dump(mode="json"))
case False, True:
return b"".join(generate_content_stream(**kwargs, client=sync_client)).decode()
case False, False:
return json.dumps(generate_content(**kwargs, client=sync_client).model_dump(mode="json"))
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.parametrize("signature_key", ["thoughtSignature", "thought_signature"])
@pytest.mark.asyncio
async def test_native_generate_content_resends_a_rejected_signature_with_the_placeholder(
is_async: bool, stream: bool, signature_key: str
):
sent: Final[list[dict[str, object]]] = []
output: Final = await _generate(
_native_contents(signature_key),
_replay((_rejected(AI_STUDIO_REJECTION), _answer(stream)), sent),
is_async,
stream,
)
assert "Berlin is 12C." in output
assert len(sent) == 2
assert _signatures_of(sent[0], signature_key) == [VERTEX_MINTED]
assert _signatures_of(sent[1], signature_key) == [get_dummy_thought_signature()]
assert {**sent[1], "contents": None} == {**sent[0], "contents": None}
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_native_generate_content_raises_an_unrelated_400_without_a_resend(is_async: bool, stream: bool):
sent: Final[list[dict[str, object]]] = []
with pytest.raises(litellm.BadRequestError, match="Unknown name"):
await _generate(
_native_contents("thoughtSignature"), _replay((_rejected(UNRELATED_REJECTION),), sent), is_async, stream
)
assert len(sent) == 1
def _history_signed_after_the_cached_prefix(tag: str) -> tuple[Mapping[str, object], ...]:
prefix: Final = " ".join(f"Reference paragraph {i} for {tag}: weather values are in Celsius." for i in range(600))
return (
{"role": "system", "content": [{"type": "text", "text": prefix, "cache_control": CACHE_CONTROL}]},
{"role": "user", "content": [{"type": "text", "text": "Hello.", "cache_control": CACHE_CONTROL}]},
*HISTORY_WITH_FOREIGN_SIGNATURES,
)
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_signature_after_the_cached_prefix_is_resent_with_the_cache_kept(is_async: bool, stream: bool):
generated: Final[list[dict[str, object]]] = []
answer: Final = await _complete(
_history_signed_after_the_cached_prefix(f"async={is_async} stream={stream}"),
_google_ai_studio(stream, generated),
is_async,
stream,
)
assert answer == "Berlin is 12C."
assert len(generated) == 2
assert generated[1]["cachedContent"] == generated[0]["cachedContent"] == "cachedContents/prefix"
assert set(_signatures_of(generated[0])) == {VERTEX_MINTED, VERTEX_MINTED_TEXT}
assert _signatures_of(generated[1]) == [get_dummy_thought_signature()] * 2

10
tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
View file

@ -2824,7 +2824,7 @@ def test_gemini_25_text_signature_survives_replay_to_gemini_3():
"""gemini-2.5 history (signed text, unsigned tool call) replayed to gemini-3 keeps the real
text signature; the dummy signature synthesized for the unsigned tool call must not suppress it."""
from litellm.litellm_core_utils.prompt_templates.factory import (
_get_dummy_thought_signature,
get_dummy_thought_signature,
)
from litellm.llms.vertex_ai.gemini.transformation import (
gemini_convert_messages_with_history,
@ -2848,7 +2848,7 @@ def test_gemini_25_text_signature_survives_replay_to_gemini_3():
assert parts[0]["text"] == "I will list the directory."
assert parts[0]["thoughtSignature"] == "real_25_signature"
assert "function_call" in parts[1]
assert parts[1]["thoughtSignature"] == _get_dummy_thought_signature()
assert parts[1]["thoughtSignature"] == get_dummy_thought_signature()
def test_gemini_function_call_signature_round_trip_no_duplicate():
@ -2990,7 +2990,7 @@ def test_thinking_block_signature_is_not_forwarded_to_gemini() -> None:
def test_anthropic_messages_history_replays_to_gemini_without_claude_signature() -> None:
from litellm.litellm_core_utils.prompt_templates.factory import _get_dummy_thought_signature
from litellm.litellm_core_utils.prompt_templates.factory import get_dummy_thought_signature
from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter
thinking: Final = "I should look the weather up before answering."
@ -3013,6 +3013,6 @@ def test_anthropic_messages_history_replays_to_gemini_without_claude_signature()
parts: Final = _parts_of(gemini_convert_messages_with_history(messages=chat_messages, model="gemini-3.8-flash"))
signatures: Final = [part["thoughtSignature"] for part in parts if "thoughtSignature" in part]
assert signatures == [_get_dummy_thought_signature()]
assert signatures == [get_dummy_thought_signature()]
assert [part for part in parts if part.get("text") == thinking] == [{"thought": True, "text": thinking}]
assert next(part for part in parts if "function_call" in part)["thoughtSignature"] == _get_dummy_thought_signature()
assert next(part for part in parts if "function_call" in part)["thoughtSignature"] == get_dummy_thought_signature()