Fix ChatSampler.conversation crash for non-Gemma4 tokenizers (breaks ToolSampler multi-turn)

gemma/gm/text/_chat_sampler.py

@@ -430,7 +430,13 @@ def initialize_stream(
   @property
   def conversation(self) -> dialog.Conversation:
     """Returns the conversation."""
-    return dialog.Conversation(''.join(t.text for t in self.turns))
+    text = ''.join(t.text for t in self.turns)
+    if self.tokenizer.FORMAT is not dialog.Format.GEMMA4:
+      # Turns are stored formatted with the tokenizer's legacy format (e.g.
+      # `<start_of_turn>` tags for Gemma 2/3/3n), but `dialog.Conversation`
+      # only parses the canonical `<|turn>` tags.
+      text = self.tokenizer.FORMAT.to_gemma4(text)
+    return dialog.Conversation(text)
 
 
 # Legacy Gemma 3 tokens to detect in user prompts.

gemma/gm/text/_sampler_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from unittest import mock
+import dialog
 from gemma import gm
 from gemma.gm.text import _sampler
 from gemma.gm.text import _sampler_loop
@@ -145,3 +146,103 @@ def test_chat_sampler_non_gemma4_dispatch():
     assert isinstance(output, str)
   # Verify sampler.sample was called (not gemma4_sampler.sample).
   mock_sample.assert_called_once()
+
+
+def test_chat_sampler_conversation_legacy_format():
+  """Tests that `conversation` parses turns stored in a legacy format.
+
+  Non-Gemma4 tokenizers (Gemma 2/3/3n) store turns formatted with the legacy
+  `<start_of_turn>` tags, while `dialog.Conversation` only parses the canonical
+  `<|turn>` tags. The `conversation` property must convert the stored text to
+  the canonical format before parsing.
+  """
+  model = gm.testing.DummyGemma()
+  params = model.init(
+      jax.random.PRNGKey(0),
+      jnp.zeros((5,), dtype=jnp.int32),
+  )
+  params = params['params']
+  tokenizer = gm.testing.DummyTokenizer()
+  assert tokenizer.FORMAT is not dialog.Format.GEMMA4  # Legacy format.
+  chat_sampler = gm.text.ChatSampler(
+      model=model,  # pyrefly: ignore[bad-argument-type]
+      params=params,
+      tokenizer=tokenizer,
+      cache_length=128,
+      max_out_length=128,
+  )
+
+  mock_sample = mock.MagicMock(
+      return_value=_sampler.SamplerOutput(
+          text='hi there<end_of_turn>',
+          state=mock.MagicMock(),
+      )
+  )
+  with mock.patch.object(
+      type(chat_sampler),
+      'sampler',
+      new_callable=lambda: property(
+          lambda self: mock.MagicMock(sample=mock_sample)
+      ),
+  ):
+    chat_sampler.chat('Hello world')
+
+  conv = chat_sampler.conversation
+  assert len(conv) == 2
+  assert isinstance(conv[0], dialog.User)
+  assert isinstance(conv[1], dialog.Model)
+  assert 'Hello world' in conv[0].as_text()
+  assert 'hi there' in conv[1].as_text()
+
+
+def test_tool_sampler_multi_turn_legacy_format():
+  """Tests `ToolSampler.chat` across turns with a legacy-format tokenizer.
+
+  `ToolSampler.chat` evaluates `self.conversation` on every call (including
+  the internal recursive call after a tool executes), so the property must
+  not fail once turns are stored (regression test).
+  """
+
+  class _NoOpToolHandler(gm.tools.ToolHandlerBase):
+
+    def _tools(self):
+      return []
+
+    def _call_tool(self, name, arguments):
+      return None
+
+  model = gm.testing.DummyGemma()
+  params = model.init(
+      jax.random.PRNGKey(0),
+      jnp.zeros((5,), dtype=jnp.int32),
+  )
+  params = params['params']
+  tokenizer = gm.testing.DummyTokenizer()
+  tool_sampler = gm.text.ToolSampler(
+      model=model,  # pyrefly: ignore[bad-argument-type]
+      params=params,
+      tokenizer=tokenizer,
+      tool_handler=_NoOpToolHandler(),
+      multi_turn=True,
+      cache_length=128,
+      max_out_length=128,
+  )
+
+  mock_sample = mock.MagicMock(
+      return_value=_sampler.SamplerOutput(
+          text='hi there<end_of_turn>',
+          state=mock.MagicMock(),
+      )
+  )
+  with mock.patch.object(
+      type(tool_sampler),
+      'sampler',
+      new_callable=lambda: property(
+          lambda self: mock.MagicMock(sample=mock_sample)
+      ),
+  ):
+    output = tool_sampler.chat('Hello world')
+    assert isinstance(output, str)
+    # Second turn reads `self.conversation` and must not fail.
+    output = tool_sampler.chat('How are you?')
+    assert isinstance(output, str)