[recipe] fix: normalize LangGraph chat template token IDs

langgraph_agent/chat_model.py

@@ -39,6 +39,7 @@
 from verl.experimental.agent_loop.agent_loop import AgentLoopOutput, AsyncLLMServerManager
 from verl.experimental.agent_loop.tool_parser import ToolParser
 from verl.experimental.agent_loop.utils import add_generation_prompt_for_gpt_oss, format_gpt_oss_tool_response_manually
+from verl.utils import normalize_token_ids
 
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
@@ -90,7 +91,9 @@ def bind_tools(self, tools, **kwargs) -> Runnable[LanguageModelInput, BaseMessag
         formatted_tools: list = [convert_to_openai_tool(tool) for tool in tools]
 
         # used to remove system prompt prefix when encoding tool response
-        system_prompt = self.tokenizer.apply_chat_template([{}], add_generation_prompt=False, tokenize=True)
+        system_prompt = normalize_token_ids(
+            self.tokenizer.apply_chat_template([{}], add_generation_prompt=False, tokenize=True)
+        )
         kwargs["system_prompt"] = system_prompt
 
         return self.bind(tools=formatted_tools, **kwargs)
@@ -190,6 +193,7 @@ async def _preprocess(self, messages: list[BaseMessage], **kwargs: Any) -> tuple
                     tokenize=True,
                 ),
             )
+            prompt_ids = normalize_token_ids(prompt_ids)
             return str(uuid.uuid4()), prompt_ids, []
 
         # Case 2: follow up chat completion with tool/human response: [system], human, ai, human|tool, ...
@@ -210,7 +214,7 @@ async def _preprocess(self, messages: list[BaseMessage], **kwargs: Any) -> tuple
                     messages, add_generation_prompt=True, tokenize=True
                 ),
             )
-            tool_response_ids = tool_response_ids[len(kwargs["system_prompt"]) :]
+            tool_response_ids = normalize_token_ids(tool_response_ids)[len(kwargs["system_prompt"]) :]
         elif self.tool_parser == "gpt-oss":
             # Format tool responses manually
             # since gpt-oss chat template requires tool call messages to parse tool response messages

langgraph_agent/test_chat_model_on_cpu.py

@@ -0,0 +1,70 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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 pytest
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from recipe.langgraph_agent.chat_model import ChatModel
+from transformers import BatchEncoding
+
+
+class _BatchEncodingTokenizer:
+    def apply_chat_template(self, *args, **kwargs):
+        del args, kwargs
+        return BatchEncoding({"input_ids": [101, 102]})
+
+
+def _make_chat_model() -> ChatModel:
+    return ChatModel.model_construct(
+        model_name="dummy-model",
+        client=None,
+        tokenizer=_BatchEncodingTokenizer(),
+        max_tokens=16,
+    )
+
+
+@pytest.mark.asyncio
+async def test_preprocess_normalizes_batch_encoding_to_token_ids():
+    model = _make_chat_model()
+
+    _, prompt_ids, response_mask = await model._preprocess([HumanMessage(content="hello")])
+
+    assert prompt_ids == [101, 102]
+    assert response_mask == []
+
+
+@pytest.mark.asyncio
+async def test_preprocess_normalizes_batch_encoding_for_tool_response():
+    model = _make_chat_model()
+    messages = [
+        HumanMessage(content="hello"),
+        AIMessage(
+            content="",
+            response_metadata={"request_id": "request-1", "prompt_ids": [1, 2], "response_mask": [1]},
+        ),
+        ToolMessage(content="tool result", tool_call_id="tool-call-1"),
+    ]
+
+    request_id, prompt_ids, response_mask = await model._preprocess(messages, system_prompt=[101])
+
+    assert request_id == "request-1"
+    assert prompt_ids == [1, 2, 102]
+    assert response_mask == [1, 0]
+
+
+def test_bind_tools_normalizes_system_prompt_to_token_ids():
+    model = _make_chat_model()
+
+    bound = model.bind_tools([])
+
+    assert bound.kwargs["system_prompt"] == [101, 102]