Add last-prompt editing and regeneration to ChatSession

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,

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],

google/genai/tests/chats/test_send_message.py

@@ -16,13 +16,15 @@
 import json
 import os
 import sys
+from unittest.mock import Mock
 
 from pydantic import BaseModel
 from pydantic import ValidationError
 import pytest
 
 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')

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',