Skip to content

Fix Gemma 4 agentic turn rendering #769

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
rshemet wants to merge 1 commit into from
Closed

Fix Gemma 4 agentic turn rendering #769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 50 additions & 6 deletions cactus-engine/src/tokenizer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -716,34 +716,75 @@ std::string Tokenizer::format_gemma4_style(const std::vector<ChatMessage>& messa
if (tw == 0) tw = side;
return static_cast<size_t>((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<std::string> pending_call_names;
auto close_model_turn = [&]() {
if (in_model_turn) { result += "<turn|>\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];
const std::string role = (msg.role == "assistant") ? "model"
: (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
: (!pending_call_names.empty() ? pending_call_names.front()
: 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);
Expand All @@ -762,13 +803,16 @@ std::string Tokenizer::format_gemma4_style(const std::vector<ChatMessage>& messa
result += "<audio|>";
}
result += "<turn|>\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;
Expand Down
60 changes: 60 additions & 0 deletions cactus-engine/tests/test_llm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<CactusModelHandle*>(model);
auto* tok = handle->model->get_tokenizer();
bool ok = true;

std::vector<ChatMessage> 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<turn|>\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<ChatMessage> 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<|\"|>}<tool_response|>";
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<turn|>\n<|turn>user\ndone<turn|>\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;
Expand Down Expand Up @@ -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());
Expand Down
64 changes: 63 additions & 1 deletion python/tests/test_chat_template_golden.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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
Expand All @@ -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))
Expand Down Expand Up @@ -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
Expand Down
Loading