From 52e878934c74d3c24d5e25546ac7b022d2870f1b Mon Sep 17 00:00:00 2001 From: rshemet Date: Fri, 17 Jul 2026 11:03:28 -0700 Subject: [PATCH] Fix Gemma 4 agentic turn rendering --- cactus-engine/src/tokenizer.cpp | 56 +++++++++++++++++--- cactus-engine/tests/test_llm.cpp | 60 +++++++++++++++++++++ python/tests/test_chat_template_golden.py | 64 ++++++++++++++++++++++- 3 files changed, 173 insertions(+), 7 deletions(-) diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index d792578db..0ba28bbdb 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -716,11 +716,21 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa if (tw == 0) tw = side; return static_cast((th / p / k) * (tw / p / k)); }; + enum class Gemma4MessageType { NONE, TOOL_CALL, TOOL_RESPONSE }; + bool in_model_turn = false; + Gemma4MessageType previous_message_type = Gemma4MessageType::NONE; + std::string previous_non_tool_role; std::vector pending_call_names; auto close_model_turn = [&]() { if (in_model_turn) { result += "\n"; in_model_turn = false; } }; + auto next_non_tool_role = [&](size_t index) -> std::string { + for (size_t j = index + 1; j < messages.size(); ++j) { + if (messages[j].role != "tool") return messages[j].role; + } + return ""; + }; for (size_t i = first_msg; i < messages.size(); i++) { const auto& msg = messages[i]; @@ -728,13 +738,29 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa : (msg.role == "developer") ? "system" : msg.role; if (role == "model") { - if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } + previous_message_type = Gemma4MessageType::NONE; + bool continue_same_model_turn = previous_non_tool_role == "assistant"; + if (!continue_same_model_turn) { + result += "<|turn>model\n"; + in_model_turn = true; + } result += msg.content; for (const auto& tc : msg.tool_calls) { result += format_tool_call_for_prompt(tc.name, tc.arguments, true); pending_call_names.push_back(tc.name); } - if (msg.tool_calls.empty()) close_model_turn(); + std::string next_role = next_non_tool_role(i); + bool has_following_tool_response = i + 1 < messages.size() + && messages[i + 1].role == "tool"; + bool continues_into_next = next_role == "assistant" + && (msg.tool_calls.empty() || has_following_tool_response); + if (!msg.tool_calls.empty()) { + previous_message_type = Gemma4MessageType::TOOL_CALL; + if (!has_following_tool_response) result += "<|tool_response>"; + } else if (!continues_into_next) { + close_model_turn(); + } + previous_non_tool_role = "assistant"; } else if (role == "tool") { if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } std::string fn = !msg.name.empty() ? msg.name @@ -742,8 +768,23 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa : std::string("unknown")); if (!pending_call_names.empty()) pending_call_names.erase(pending_call_names.begin()); result += format_tool_response_for_prompt(fn, msg.content, true); + previous_message_type = Gemma4MessageType::TOOL_RESPONSE; + + bool more_tool_responses = i + 1 < messages.size() + && messages[i + 1].role == "tool"; + std::string next_role = next_non_tool_role(i); + bool continues_into_assistant = previous_non_tool_role == "assistant" + && next_role == "assistant"; + if (!more_tool_responses && !continues_into_assistant && !next_role.empty()) { + close_model_turn(); + } } else { - close_model_turn(); + if (previous_message_type == Gemma4MessageType::TOOL_CALL) { + in_model_turn = false; + } else { + close_model_turn(); + } + previous_message_type = Gemma4MessageType::NONE; result += "<|turn>" + role + "\n"; for (const auto& image_path : msg.images) { size_t n = compute_soft_tokens(image_path); @@ -762,13 +803,16 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa result += ""; } result += "\n"; + previous_non_tool_role = msg.role; } } if (add_generation_prompt) { - if (!in_model_turn) result += "<|turn>model\n"; - } else { - close_model_turn(); + if (previous_message_type == Gemma4MessageType::TOOL_RESPONSE) { + if (enable_thinking_if_supported) result += "<|channel>thought\n"; + } else if (previous_message_type != Gemma4MessageType::TOOL_CALL && !in_model_turn) { + result += "<|turn>model\n"; + } } return result; diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index 8bcff7526..bf09347ce 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -410,6 +410,65 @@ bool test_partition_thinking_response() { "thought1\nthought2", "text1text2"); } +bool test_prompt_gemma4_agentic_turn_closure() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + bool ok = true; + + std::vector consecutive = { + {"user", "start", "", {}, {}, 0, {}}, + {"assistant", "first", "", {}, {}, 0, {}}, + {"assistant", "second", "", {}, {}, 0, {}} + }; + std::string consecutive_prompt = tok->format_chat_prompt(consecutive, true, "", false); + std::string consecutive_expected = "<|turn>model\nfirstsecond\n<|turn>model\n"; + if (consecutive_prompt.find(consecutive_expected) == std::string::npos) { + std::cerr << " consecutive assistant messages did not share a model turn\n"; + ok = false; + } + + std::vector tool_loop = { + {"user", "weather", "", {}, {}, 0, {}}, + {"assistant", "", "", {}, {}, 0, {{"get_weather", R"({"city":"Paris"})"}}}, + {"tool", "Sunny", "get_weather", {}, {}, 0, {}} + }; + std::string terminal_plain = tok->format_chat_prompt(tool_loop, true, "", false); + std::string terminal_response = + "<|tool_response>response:get_weather{value:<|\"|>Sunny<|\"|>}"; + if (terminal_plain.size() < terminal_response.size() + || terminal_plain.compare(terminal_plain.size() - terminal_response.size(), + terminal_response.size(), terminal_response) != 0) { + std::cerr << " terminal tool response was closed before generation\n"; + ok = false; + } + + std::string terminal_thinking = tok->format_chat_prompt(tool_loop, true, "", true); + std::string thought_primer = terminal_response + "<|channel>thought\n"; + if (terminal_thinking.size() < thought_primer.size() + || terminal_thinking.compare(terminal_thinking.size() - thought_primer.size(), + thought_primer.size(), thought_primer) != 0) { + std::cerr << " thinking primer missing after terminal tool response\n"; + ok = false; + } + + tool_loop.push_back({"assistant", "first", "", {}, {}, 0, {}}); + tool_loop.push_back({"assistant", "second", "", {}, {}, 0, {}}); + tool_loop.push_back({"user", "done", "", {}, {}, 0, {}}); + std::string continued = tok->format_chat_prompt(tool_loop, true, "", false); + std::string continued_expected = terminal_response + + "firstsecond\n<|turn>user\ndone\n<|turn>model\n"; + if (continued.find(continued_expected) == std::string::npos) { + std::cerr << " tool response did not continue into the following assistant turn\n"; + ok = false; + } + + cactus_destroy(model); + return ok; +} + bool test_prompt_gemma4_retains_thinking() { cactus_model_t model = load_gemma4_or_skip(); if (!model) return true; @@ -669,6 +728,7 @@ int main() { runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); runner.run_test("tool_constraint_clear_releases_bias", test_tool_constraint_clear_releases_bias()); runner.run_test("partition_thinking_response", test_partition_thinking_response()); + runner.run_test("prompt_gemma4_agentic_turn_closure", test_prompt_gemma4_agentic_turn_closure()); runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking()); runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean()); runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist()); diff --git a/python/tests/test_chat_template_golden.py b/python/tests/test_chat_template_golden.py index 78077dd2c..f2891c1d8 100644 --- a/python/tests/test_chat_template_golden.py +++ b/python/tests/test_chat_template_golden.py @@ -95,6 +95,19 @@ {"role": "user", "content": "Briefly explain entropy."}, ] +CONSECUTIVE_ASSISTANT = [ + {"role": "user", "content": "Give the answer in two parts."}, + {"role": "assistant", "content": "Part one."}, + {"role": "assistant", "content": "Part two."}, +] + +TOOL_RESPONSE_CONTINUATION = [ + *TOOL_CALL_RESPONSE, + {"role": "assistant", "content": "The weather is sunny."}, + {"role": "assistant", "content": "No rain is expected."}, + {"role": "user", "content": "Thanks."}, +] + CASES = [ ("system_chat", SYSTEM_CHAT, None, False), ("tools_user", TOOLS_USER, TOOLS, False), @@ -136,6 +149,23 @@ def _to_hf_messages(messages: list[dict], model_type: str) -> list[dict]: _LOADED: dict[str, tuple] = {} +def _cached_snapshot(hf_id: str) -> Path | None: + from huggingface_hub.constants import HF_HUB_CACHE + + repo_dir = Path(HF_HUB_CACHE) / f"models--{hf_id.replace('/', '--')}" + main_ref = repo_dir / "refs" / "main" + if main_ref.is_file(): + snapshot = repo_dir / "snapshots" / main_ref.read_text(encoding="utf-8").strip() + if snapshot.is_dir(): + return snapshot + snapshots = sorted( + (path for path in (repo_dir / "snapshots").glob("*") if path.is_dir()), + key=lambda path: path.stat().st_mtime, + reverse=True, + ) + return snapshots[0] if snapshots else None + + @pytest.fixture(scope="session") def load_family(): from cactus import cactus_destroy, cactus_init @@ -148,8 +178,11 @@ def _load(bundle_name: str, hf_id: str): ) if not _valid_bundle(bundle): pytest.skip(f"bundle not prepared: {bundle}") + snapshot = _cached_snapshot(hf_id) + if snapshot is None: + pytest.skip(f"HF tokenizer not cached for {hf_id}") try: - tokenizer = transformers.AutoTokenizer.from_pretrained(hf_id, local_files_only=True) + tokenizer = transformers.AutoTokenizer.from_pretrained(snapshot, local_files_only=True) except OSError as exc: pytest.skip(f"HF tokenizer not cached for {hf_id}: {exc}") _LOADED[bundle_name] = (cactus_init(str(bundle)), tokenizer, _read_model_type(bundle)) @@ -198,6 +231,35 @@ def test_render_matches_hf_template(load_family, bundle_name, hf_id, case_name, f"{bundle_name}/{case_name}") +@pytest.mark.parametrize( + ("case_name", "messages", "tools", "thinking"), + [ + ("consecutive_assistant", CONSECUTIVE_ASSISTANT, None, False), + ("terminal_tool_response_thinking", TOOL_CALL_RESPONSE, TOOLS, True), + ("tool_response_continuation", TOOL_RESPONSE_CONTINUATION, TOOLS, False), + ], +) +def test_gemma4_agentic_render_matches_updated_template( + load_family, case_name, messages, tools, thinking +): + from cactus import cactus_render_prompt + + bundle_name, hf_id = FAMILIES[0] + model, tokenizer, model_type = load_family(bundle_name, hf_id) + engine_render = cactus_render_prompt( + model, messages, options={"enable_thinking_if_supported": thinking}, tools=tools + ) + hf_render = tokenizer.apply_chat_template( + _to_hf_messages(messages, model_type), + tools=tools, + add_generation_prompt=True, + tokenize=False, + enable_thinking=thinking, + ) + _assert_render_matches(model, tokenizer, engine_render, hf_render, + f"{bundle_name}/{case_name}") + + def _monkey_soft_tokens(hf_id: str) -> int: from huggingface_hub.constants import HF_HUB_CACHE from PIL import Image