From 404b31db4c42874cd62f4be8df69419e68ed96f2 Mon Sep 17 00:00:00 2001 From: liangming Date: Sun, 23 Aug 2026 13:13:00 +0800 Subject: [PATCH 1/2] Retry Google Auth transport failures --- google/genai/_api_client.py | 2 +- google/genai/tests/client/test_retries.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 3c4cbc247..93699ae6d 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -577,7 +577,7 @@ def retry_args(options: Optional[HttpRetryOptions]) -> _common.StringDict: retriable_codes = options.http_status_codes or _RETRY_HTTP_STATUS_CODES retry = tenacity.retry_if_exception( lambda e: (isinstance(e, errors.APIError) and e.code in retriable_codes) - or isinstance(e, _HTTPX_TRANSIENT_EXC), + or isinstance(e, _HTTPX_TRANSIENT_EXC + (auth_exceptions.TransportError,)), ) wait = tenacity.wait_exponential_jitter( initial=options.initial_delay or _RETRY_INITIAL_DELAY, diff --git a/google/genai/tests/client/test_retries.py b/google/genai/tests/client/test_retries.py index d02711313..e92cfe822 100644 --- a/google/genai/tests/client/test_retries.py +++ b/google/genai/tests/client/test_retries.py @@ -19,6 +19,8 @@ from collections.abc import Sequence import datetime from unittest import mock + +from google.auth import exceptions as auth_exceptions import pytest try: @@ -204,6 +206,12 @@ def test_retry_args_retries_httpx_transport_errors(): assert not retry.predicate(ValueError('not a transport error')) +def test_retry_args_retries_google_auth_transport_errors(): + args = api_client.retry_args(types.HttpRetryOptions()) + assert args['retry'].predicate(auth_exceptions.TransportError('refresh failed')) + assert not args['retry'].predicate(auth_exceptions.RefreshError('invalid credentials')) + + def _patch_auth_default(): return mock.patch( 'google.auth.default', From 4692a4a698a0b2490efa8588e9a752b93384b147 Mon Sep 17 00:00:00 2001 From: liangming Date: Sun, 23 Aug 2026 13:57:22 +0800 Subject: [PATCH 2/2] Add last-prompt editing and regeneration to ChatSession --- google/genai/chats.py | 75 +++++++++++++++++++ google/genai/tests/chats/test_send_message.py | 36 +++++++++ 2 files changed, 111 insertions(+) diff --git a/google/genai/chats.py b/google/genai/chats.py index be10298bc..e05e6d39b 100644 --- a/google/genai/chats.py +++ b/google/genai/chats.py @@ -182,6 +182,31 @@ def get_history(self, curated: bool = False) -> list[Content]: else: return self._comprehensive_history + def _get_last_user_input(self) -> Content: + for content in reversed(self._curated_history): + if content.role == "user": + return content + raise ValueError("Cannot regenerate a chat with no user prompt.") + + def _remove_last_turn(self) -> None: + last_user_index = next( + (index for index in range(len(self._curated_history) - 1, -1, -1) + if self._curated_history[index].role == "user"), + None, + ) + if last_user_index is None: + raise ValueError("Cannot update a chat with no user prompt.") + + del self._curated_history[last_user_index:] + + comprehensive_user_index = next( + (index for index in range(len(self._comprehensive_history) - 1, -1, -1) + if self._comprehensive_history[index].role == "user"), + None, + ) + if comprehensive_user_index is not None: + del self._comprehensive_history[comprehensive_user_index:] + def _is_part_type( contents: Union[list[PartUnionDict], PartUnionDict], @@ -218,6 +243,31 @@ def __init__( history=history, ) + def update_last_prompt( + self, + message: Union[list[PartUnionDict], PartUnionDict], + config: Optional[GenerateContentConfigOrDict] = None, + ) -> GenerateContentResponse: + """Replaces the last user prompt and generates a new response. + + The message accepts the same text, media, and audio part types as + :meth:`send_message`. + """ + if not _is_part_type(message): + raise ValueError( + f"Message must be a valid part type: {types.PartUnion} or" + f" {types.PartUnionDict}, got {type(message)}" + ) + self._remove_last_turn() + return self.send_message(message, config) + + def regenerate_last_turn( + self, + config: Optional[GenerateContentConfigOrDict] = None, + ) -> GenerateContentResponse: + """Regenerates the response for the last user prompt.""" + return self.update_last_prompt(self._get_last_user_input().parts, config) + def send_message( self, message: Union[list[PartUnionDict], PartUnionDict], @@ -602,6 +652,31 @@ def __init__( history=history, ) + async def update_last_prompt( + self, + message: Union[list[PartUnionDict], PartUnionDict], + config: Optional[GenerateContentConfigOrDict] = None, + ) -> GenerateContentResponse: + """Replaces the last user prompt and generates a new response. + + The message accepts the same text, media, and audio part types as + :meth:`send_message`. + """ + if not _is_part_type(message): + raise ValueError( + f"Message must be a valid part type: {types.PartUnion} or" + f" {types.PartUnionDict}, got {type(message)}" + ) + self._remove_last_turn() + return await self.send_message(message, config) + + async def regenerate_last_turn( + self, + config: Optional[GenerateContentConfigOrDict] = None, + ) -> GenerateContentResponse: + """Regenerates the response for the last user prompt.""" + return await self.update_last_prompt(self._get_last_user_input().parts, config) + async def send_message( self, message: Union[list[PartUnionDict], PartUnionDict], diff --git a/google/genai/tests/chats/test_send_message.py b/google/genai/tests/chats/test_send_message.py index 94d451582..7ff349b3d 100644 --- a/google/genai/tests/chats/test_send_message.py +++ b/google/genai/tests/chats/test_send_message.py @@ -16,6 +16,7 @@ import json import os import sys +from unittest.mock import Mock from pydantic import BaseModel from pydantic import ValidationError @@ -23,6 +24,7 @@ from .. import pytest_helper from ... import errors +from ...chats import Chat from ... import types try: @@ -253,6 +255,40 @@ def test_history(client): assert len(chat.get_history()) > 2 +def test_update_last_prompt_replaces_previous_turn(): + history = [ + types.Content(role='user', parts=[types.Part.from_text(text='old prompt')]), + types.Content(role='model', parts=[types.Part.from_text(text='old response')]), + ] + chat = Chat(modules=Mock(), model=MODEL_NAME, history=history) + response = Mock() + chat.send_message = Mock(return_value=response) + + result = chat.update_last_prompt( + [types.Part.from_text(text='updated prompt'), types.Part.from_text(text='audio')] + ) + + assert result is response + assert chat.get_history() == [] + chat.send_message.assert_called_once() + + +def test_regenerate_last_turn_reuses_previous_prompt(): + history = [ + types.Content(role='user', parts=[types.Part.from_text(text='keep this')]), + types.Content(role='model', parts=[types.Part.from_text(text='response')]), + ] + chat = Chat(modules=Mock(), model=MODEL_NAME, history=history) + response = Mock() + chat.send_message = Mock(return_value=response) + + result = chat.regenerate_last_turn() + + assert result is response + sent_parts = chat.send_message.call_args.args[0] + assert sent_parts[0].text == 'keep this' + + def test_send_2_messages(client): chat = client.chats.create(model=MODEL_NAME) chat.send_message('write a python function to check if a year is a leap year')