diff --git a/tinker_cookbook/recipes/tool_use/search/search_env.py b/tinker_cookbook/recipes/tool_use/search/search_env.py
index 64cc5112..1f1254ce 100644
--- a/tinker_cookbook/recipes/tool_use/search/search_env.py
+++ b/tinker_cookbook/recipes/tool_use/search/search_env.py
@@ -32,7 +32,7 @@
SEARCH_TOOL_SYSTEM_PROMPT = """
You are an expert assistant who solves tasks using a Wikipedia search tool.
-Tool calling. Execute the tool by wrapping calls in ...
+Tool calling. Execute the tool by wrapping calls in ...
The search tool you are given has the following schema:
```
@@ -69,9 +69,9 @@
“Between 2020 and 2025, which year did New York City see the most population growth and how did San Francisco population change in that year?”
1. Think step by step: In order to answer this question, I need to know the population of New York City and San Francisco between 2020 and 2025. I will search for the population of New York City in each year
-2. Calling search tool: {"name": "search", "args": {"query_list": ["Population New York city between 2020 and 2025"]}} (Output omitted for brevity)
+2. Calling search tool: {"name": "search", "args": {"query_list": ["Population New York city between 2020 and 2025"]}} (Output omitted for brevity)
3. Think step by step again: I have the population of New York City in each year, and I see that the population of New York City grew the most in 2024. I need to know the population of San Francisco in 2024. I will search for the population of San Francisco in each year.
-{"name": "search", "args": {"query_list": ["Population San Francisco between 2023 and 2024"]}} (Output omitted for brevity)
+{"name": "search", "args": {"query_list": ["Population San Francisco between 2023 and 2024"]}} (Output omitted for brevity)
4. Answer: The population of New York City grew the most in 2024, and the population of San Francisco changed by XXXX in 2024.
"""
diff --git a/tinker_cookbook/renderers.py b/tinker_cookbook/renderers.py
index 9f03020d..44259c78 100644
--- a/tinker_cookbook/renderers.py
+++ b/tinker_cookbook/renderers.py
@@ -481,13 +481,20 @@ def _render_message(self, idx: int, message: Message) -> tuple[list[int], list[i
ob_str = f"{maybe_newline}<|im_start|>{message['role']}\n"
ac_content = message["content"]
# Observation (prompt) part
- ac_str = f"{ac_content}<|im_end|>"
+ if "tool_calls" in message:
+ ac_content += "\n".join(
+ [
+ f"\n{json.dumps(tool_call)}\n"
+ for tool_call in message["tool_calls"]
+ ]
+ )
+ ac_content += "<|im_end|>"
# Action part
ac_tail_str = "" # No action tail needed for Qwen format
# Action part that's only included in the last message in SFT
return (
self.tokenizer.encode(ob_str, add_special_tokens=False),
- self.tokenizer.encode(ac_str, add_special_tokens=False),
+ self.tokenizer.encode(ac_content, add_special_tokens=False),
self.tokenizer.encode(ac_tail_str, add_special_tokens=False),
)