fix(local_tokenizer): align TokensInfo roles with the texts actually tokenized

google/genai/local_tokenizer.py

@@ -59,6 +59,9 @@ class _TextsAccumulator:
   def __init__(self) -> None:
     self._texts: list[str] = []
 
+  def __len__(self) -> int:
+    return len(self._texts)
+
   def get_texts(self) -> Iterable[str]:
     return self._texts
 
@@ -369,14 +372,19 @@ def compute_tokens(
       # tokens_info=[TokensInfo(token_ids=[279, 329, 1313, 2508, 13], tokens=[b' What', b' is', b' your', b' name', b'?'], role='user')]
     """
     processed_contents = t.t_contents(contents)
-    roles = []
+    roles: list[Optional[str]] = []
 
     text_accumulator = _TextsAccumulator()
     for content in processed_contents:
+      texts_before = len(text_accumulator)
       text_accumulator.add_content(content)
-      if content.parts:
-        for _ in content.parts:
-          roles.append(content.role)
+      # A part does not map to exactly one tokenized text: a function_call
+      # or function_response part contributes the function name plus every
+      # key and string value of its args/response as separate texts, and a
+      # thought_signature-only part contributes none. Extend `roles` by the
+      # number of texts the accumulator actually added, so the zip() below
+      # stays aligned with `text_accumulator.get_texts()`.
+      roles.extend([content.role] * (len(text_accumulator) - texts_before))
 
     token_infos = []
     if self._tokenizer_name in loader.GEMMA_TOKENIZER_TO_MODEL_NAMES:

google/genai/tests/local_tokenizer/test_roles_alignment.py

@@ -0,0 +1,73 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+from sentencepiece import sentencepiece_model_pb2
+
+from ... import local_tokenizer
+from ... import types
+
+_FC_CONTENT = types.Content(
+    role='model',
+    parts=[types.Part(function_call=types.FunctionCall(
+        name='get_weather', args={'location': 'Boston'}))],
+)
+_USER_CONTENT = types.Content(role='user', parts=[types.Part(text='thanks')])
+# accumulator yields: 'get_weather', 'location', 'Boston', 'thanks'
+_EXPECTED_ROLES = ['model', 'model', 'model', 'user']
+
+
+class TestSentencePieceBranch(unittest.TestCase):
+
+  def setUp(self):
+    patch('genai._local_tokenizer_loader.load_model_proto').start()
+    m = patch('genai._local_tokenizer_loader.get_sentencepiece').start()
+    self.addCleanup(patch.stopall)
+    self.mock_tokenizer = MagicMock()
+    m.return_value = self.mock_tokenizer
+    self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-2.5-flash')
+    self.tokenizer._model_proto = sentencepiece_model_pb2.ModelProto(
+        pieces=[sentencepiece_model_pb2.ModelProto.SentencePiece(piece='x')] * 10)
+
+  def test_roles_align_with_texts(self):
+    def proto(i):
+      p = MagicMock(); p.pieces = [MagicMock(id=1, piece=f't{i}')]; return p
+    self.mock_tokenizer.EncodeAsImmutableProto.side_effect = (
+        lambda texts: [proto(i) for i in range(len(texts))])
+    r = self.tokenizer.compute_tokens([_FC_CONTENT, _USER_CONTENT])
+    self.assertEqual([t.role for t in r.tokens_info], _EXPECTED_ROLES)
+
+
+class TestHuggingFaceBranch(unittest.TestCase):
+
+  def setUp(self):
+    m = patch('genai._local_tokenizer_loader.get_huggingface_tokenizer').start()
+    self.addCleanup(patch.stopall)
+    self.mock_tokenizer = MagicMock()
+    m.return_value = self.mock_tokenizer
+    self.tokenizer = local_tokenizer.LocalTokenizer(model_name='gemini-3.5-flash')
+
+  def test_roles_align_with_texts(self):
+    self.mock_tokenizer.encode.side_effect = (
+        lambda texts: [[i] for i in range(len(texts))])
+    self.mock_tokenizer.convert_ids_to_tokens.side_effect = lambda ids: ['tok']
+    r = self.tokenizer.compute_tokens([_FC_CONTENT, _USER_CONTENT])
+    self.assertEqual([t.role for t in r.tokens_info], _EXPECTED_ROLES)
+
+
+if __name__ == '__main__':
+  unittest.main()