|
| 1 | +# Copyright 2026 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Resumes a generation that the model paused with a continuation token.""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +from collections.abc import Iterable |
| 20 | +from collections.abc import Sequence |
| 21 | +import itertools |
| 22 | +import logging |
| 23 | + |
| 24 | +from google.genai import types |
| 25 | + |
| 26 | +logger = logging.getLogger('google_adk.' + __name__) |
| 27 | + |
| 28 | +_TEXT_FIELDS = frozenset({'text', 'thought', 'thought_signature'}) |
| 29 | +_MAX_RESUMES = 256 |
| 30 | +# A resumed request carries the output so far, which one failure would discard. |
| 31 | +_RESUME_RETRY_OPTIONS = types.HttpRetryOptions(attempts=8) |
| 32 | +# A resume resends the prompt, so these repeat on every request and describe the |
| 33 | +# turn only when taken from the first one. |
| 34 | +_PROMPT_USAGE_FIELDS = frozenset({ |
| 35 | + 'cache_tokens_details', |
| 36 | + 'cached_content_token_count', |
| 37 | + 'prompt_token_count', |
| 38 | + 'prompt_tokens_details', |
| 39 | +}) |
| 40 | +_TOTALLED_USAGE_FIELDS = ( |
| 41 | + 'prompt_token_count', |
| 42 | + 'candidates_token_count', |
| 43 | + 'thoughts_token_count', |
| 44 | + 'tool_use_prompt_token_count', |
| 45 | +) |
| 46 | + |
| 47 | + |
| 48 | +def continuation_token( |
| 49 | + response: types.GenerateContentResponse, |
| 50 | +) -> bytes | None: |
| 51 | + """Returns the token that resumes ``response``, or None if it is complete.""" |
| 52 | + if not response.candidates: |
| 53 | + return None |
| 54 | + candidate = response.candidates[0] |
| 55 | + if candidate.finish_reason != types.FinishReason.CONTINUATION: |
| 56 | + return None |
| 57 | + return candidate.continuation_token or None |
| 58 | + |
| 59 | + |
| 60 | +class Continuation: |
| 61 | + """Carries one generation across the requests that resume it.""" |
| 62 | + |
| 63 | + def __init__( |
| 64 | + self, |
| 65 | + contents: Sequence[types.Content], |
| 66 | + config: types.GenerateContentConfig | None, |
| 67 | + *, |
| 68 | + retry_resends: bool = False, |
| 69 | + ) -> None: |
| 70 | + self._contents = contents |
| 71 | + self._config = config |
| 72 | + self._retry_resends = retry_resends |
| 73 | + self._parts: list[types.Part] = [] |
| 74 | + self._usage: types.GenerateContentResponseUsageMetadata | None = None |
| 75 | + self._token: bytes | None = None |
| 76 | + self._resumes = 0 |
| 77 | + |
| 78 | + @property |
| 79 | + def resumed(self) -> bool: |
| 80 | + """Whether the generation took more than one request.""" |
| 81 | + return self._token is not None |
| 82 | + |
| 83 | + @property |
| 84 | + def content(self) -> types.Content | None: |
| 85 | + """Everything generated across the requests recorded so far.""" |
| 86 | + if not self._parts: |
| 87 | + return None |
| 88 | + return types.Content(role='model', parts=list(self._parts)) |
| 89 | + |
| 90 | + @property |
| 91 | + def usage(self) -> types.GenerateContentResponseUsageMetadata | None: |
| 92 | + """The token usage summed over the requests recorded so far.""" |
| 93 | + return self._usage |
| 94 | + |
| 95 | + def will_resume(self, token: bytes) -> bool: |
| 96 | + """Whether a request that ended with ``token`` is resumed.""" |
| 97 | + return token != self._token and self._resumes < _MAX_RESUMES |
| 98 | + |
| 99 | + def advance( |
| 100 | + self, |
| 101 | + *, |
| 102 | + token: bytes | None, |
| 103 | + parts: Sequence[types.Part] | None, |
| 104 | + usage: types.GenerateContentResponseUsageMetadata | None, |
| 105 | + ) -> tuple[list[types.Content], types.GenerateContentConfig] | None: |
| 106 | + """Records one request's output and returns the request that resumes it. |
| 107 | +
|
| 108 | + Args: |
| 109 | + token: The continuation token the request ended with, if any. |
| 110 | + parts: The parts the request generated. |
| 111 | + usage: The token usage the request reported. |
| 112 | +
|
| 113 | + Returns: |
| 114 | + The contents and config of the next request, or None once the |
| 115 | + generation is complete. |
| 116 | + """ |
| 117 | + self._usage = _add_usage(self._usage, usage) |
| 118 | + if token is None and not self.resumed: |
| 119 | + return None |
| 120 | + for index, part in enumerate(parts or []): |
| 121 | + _append_part(self._parts, part, at_pause=index == 0) |
| 122 | + if token is None: |
| 123 | + return None |
| 124 | + if ends_mid_call(self._parts): |
| 125 | + logger.warning( |
| 126 | + 'The model paused in the middle of a function call, which cannot be' |
| 127 | + ' resumed; returning the output generated so far.' |
| 128 | + ) |
| 129 | + return None |
| 130 | + if not self.will_resume(token): |
| 131 | + if token == self._token: |
| 132 | + # A token that did not change means the model made no progress. |
| 133 | + logger.warning( |
| 134 | + 'The model returned the same continuation token twice; returning' |
| 135 | + ' the output generated so far.' |
| 136 | + ) |
| 137 | + else: |
| 138 | + logger.warning( |
| 139 | + 'The model paused the generation %d times; returning the output' |
| 140 | + ' generated so far.', |
| 141 | + _MAX_RESUMES, |
| 142 | + ) |
| 143 | + return None |
| 144 | + self._token = token |
| 145 | + self._resumes += 1 |
| 146 | + logger.info('The model paused the generation; resuming it.') |
| 147 | + return self._next_request(token) |
| 148 | + |
| 149 | + def _next_request( |
| 150 | + self, token: bytes |
| 151 | + ) -> tuple[list[types.Content], types.GenerateContentConfig]: |
| 152 | + """Returns the contents and config of the request that resumes ``token``.""" |
| 153 | + contents = list(self._contents) |
| 154 | + if self._parts: |
| 155 | + contents.append(types.Content(role='model', parts=list(self._parts))) |
| 156 | + config = ( |
| 157 | + self._config.model_copy(update={'continuation_token': token}) |
| 158 | + if self._config |
| 159 | + else types.GenerateContentConfig(continuation_token=token) |
| 160 | + ) |
| 161 | + if self._retry_resends and ( |
| 162 | + config.http_options is None or config.http_options.retry_options is None |
| 163 | + ): |
| 164 | + config.http_options = ( |
| 165 | + config.http_options or types.HttpOptions() |
| 166 | + ).model_copy(update={'retry_options': _RESUME_RETRY_OPTIONS}) |
| 167 | + return contents, config |
| 168 | + |
| 169 | + |
| 170 | +def ends_mid_call(parts: Sequence[types.Part]) -> bool: |
| 171 | + """Whether ``parts`` ends with a function call the model has not finished. |
| 172 | +
|
| 173 | + Streaming a call's arguments splits it over parts that set ``will_continue`` |
| 174 | + until the last one, and a request whose final model part still sets it is |
| 175 | + rejected, so a generation paused there cannot be resumed. |
| 176 | + """ |
| 177 | + if not parts: |
| 178 | + return False |
| 179 | + function_call = parts[-1].function_call |
| 180 | + return bool(function_call and function_call.will_continue) |
| 181 | + |
| 182 | + |
| 183 | +def join_parts(parts: Sequence[types.Part]) -> list[types.Part]: |
| 184 | + """Returns ``parts`` with adjacent text joined back together. |
| 185 | +
|
| 186 | + Callers pass parts the transport split rather than the model: a stream gives |
| 187 | + each chunk of one text run its own part, and the aggregator ends a run at |
| 188 | + each resume. |
| 189 | + """ |
| 190 | + joined: list[types.Part] = [] |
| 191 | + for part in parts: |
| 192 | + _append_part(joined, part, at_pause=True) |
| 193 | + return joined |
| 194 | + |
| 195 | + |
| 196 | +def _append_part( |
| 197 | + parts: list[types.Part], part: types.Part, *, at_pause: bool |
| 198 | +) -> None: |
| 199 | + """Appends ``part``, joining it onto the text a pause split. |
| 200 | +
|
| 201 | + A pause ends its text with an empty part carrying the thought signature, so |
| 202 | + that part joins onto the text it closes and the later signature wins. The |
| 203 | + streaming aggregator keeps those boundaries instead, which is why a resumed |
| 204 | + generation runs its result through here. |
| 205 | +
|
| 206 | + Args: |
| 207 | + parts: The parts accumulated so far, appended to in place. |
| 208 | + part: The part to add. |
| 209 | + at_pause: Whether a pause separated ``part`` from the ones before it. |
| 210 | + """ |
| 211 | + previous = parts[-1] if parts else None |
| 212 | + if previous is not None and _can_join(previous, part, at_pause=at_pause): |
| 213 | + parts[-1] = types.Part( |
| 214 | + text=(previous.text or '') + (part.text or ''), |
| 215 | + thought=previous.thought, |
| 216 | + thought_signature=part.thought_signature or previous.thought_signature, |
| 217 | + ) |
| 218 | + elif part.text or part.thought_signature or not _only_has_text_fields(part): |
| 219 | + parts.append(part) |
| 220 | + |
| 221 | + |
| 222 | +def _can_join(first: types.Part, second: types.Part, *, at_pause: bool) -> bool: |
| 223 | + """Whether ``second`` continues the text ``first`` started.""" |
| 224 | + if not _only_has_text_fields(first) or not _only_has_text_fields(second): |
| 225 | + return False |
| 226 | + if not first.text: |
| 227 | + return False |
| 228 | + if not second.text: |
| 229 | + # A pause closes its text with a signature on an otherwise empty part. |
| 230 | + return bool(second.thought_signature) |
| 231 | + # Two text parts the model returned side by side are two parts, each with |
| 232 | + # its own signature; only a pause splits text that belongs together. |
| 233 | + return at_pause and bool(first.thought) == bool(second.thought) |
| 234 | + |
| 235 | + |
| 236 | +def _only_has_text_fields(part: types.Part) -> bool: |
| 237 | + """Whether ``part`` carries nothing besides text, thought and signature.""" |
| 238 | + return all( |
| 239 | + getattr(part, name) is None |
| 240 | + for name in type(part).model_fields |
| 241 | + if name not in _TEXT_FIELDS |
| 242 | + ) |
| 243 | + |
| 244 | + |
| 245 | +def _add_usage( |
| 246 | + total: types.GenerateContentResponseUsageMetadata | None, |
| 247 | + usage: types.GenerateContentResponseUsageMetadata | None, |
| 248 | +) -> types.GenerateContentResponseUsageMetadata | None: |
| 249 | + """Returns the combined token usage of two requests.""" |
| 250 | + if total is None: |
| 251 | + return usage |
| 252 | + if usage is None: |
| 253 | + return total |
| 254 | + combined: dict[str, object] = {} |
| 255 | + for name in type(usage).model_fields: |
| 256 | + if name == 'total_token_count': |
| 257 | + continue |
| 258 | + if name in _PROMPT_USAGE_FIELDS: |
| 259 | + first = getattr(total, name) |
| 260 | + combined[name] = first if first is not None else getattr(usage, name) |
| 261 | + continue |
| 262 | + values = [ |
| 263 | + value |
| 264 | + for value in (getattr(total, name), getattr(usage, name)) |
| 265 | + if value is not None |
| 266 | + ] |
| 267 | + if values and all(isinstance(value, int) for value in values): |
| 268 | + combined[name] = sum(values) |
| 269 | + elif values and all(isinstance(value, list) for value in values): |
| 270 | + combined[name] = _add_modality_counts(itertools.chain(*values)) |
| 271 | + counts = [ |
| 272 | + value |
| 273 | + for name in _TOTALLED_USAGE_FIELDS |
| 274 | + if isinstance(value := combined.get(name), int) |
| 275 | + ] |
| 276 | + if counts: |
| 277 | + combined['total_token_count'] = sum(counts) |
| 278 | + return usage.model_copy(update=combined) |
| 279 | + |
| 280 | + |
| 281 | +def _add_modality_counts( |
| 282 | + counts: Iterable[types.ModalityTokenCount], |
| 283 | +) -> list[types.ModalityTokenCount]: |
| 284 | + """Sums token counts that share a modality, keeping first-seen order.""" |
| 285 | + totals: dict[types.MediaModality | None, int] = {} |
| 286 | + for count in counts: |
| 287 | + totals[count.modality] = totals.get(count.modality, 0) + ( |
| 288 | + count.token_count or 0 |
| 289 | + ) |
| 290 | + return [ |
| 291 | + types.ModalityTokenCount(modality=modality, token_count=token_count) |
| 292 | + for modality, token_count in totals.items() |
| 293 | + ] |
0 commit comments