Skip to content

Commit 58f0ca6

Browse files
ochafikochafik
authored
Merge pull request #16 from google/tool-call
chat_template.hpp improvements
2 parents 202aa2f + bbc02d6 commit 58f0ca6

File tree

5 files changed

+212
-55
lines changed

5 files changed

+212
-55
lines changed

‎include/minja/chat-template.hpp

+148-34
Original file line numberDiff line numberDiff line change
@@ -21,47 +21,114 @@ class chat_template {
2121
public:
2222

2323
private:
24-
bool _supports_tools = true;
24+
bool supports_tools_ = true;
2525
// Meta-Llama-3.1-8B-Instruct's template expects arguments to be an object.
2626
// Most other templates (and OpenAI's API) expect the arguments object to be stringified.
27-
bool _requires_object_arguments = false;
28-
bool _supports_system_role = true;
29-
std::string _source;
30-
std::string _bos_token;
31-
std::string _eos_token;
32-
std::shared_ptr<minja::TemplateNode> _template_root;
27+
bool requires_object_arguments_ = false;
28+
bool supports_system_role_ = true;
29+
bool supports_parallel_tool_calls_ = false;
30+
std::string source_;
31+
std::string bos_token_;
32+
std::string eos_token_;
33+
std::shared_ptr<minja::TemplateNode> template_root_;
34+
35+
std::string try_render(
36+
const nlohmann::ordered_json & messages,
37+
const nlohmann::ordered_json & tools,
38+
bool add_generation_prompt,
39+
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
40+
{
41+
try {
42+
auto prompt = apply(messages, tools, add_generation_prompt, extra_context);
43+
// fprintf(stderr, "Prompt: %s\n", prompt.c_str());
44+
return prompt;
45+
} catch (const std::exception & e) {
46+
// fprintf(stderr, "Error: %s\n", e.what());
47+
return "";
48+
}
49+
}
3350

3451
public:
3552
chat_template(const std::string & source, const std::string & bos_token, const std::string & eos_token)
36-
: _source(source), _bos_token(bos_token), _eos_token(eos_token)
53+
: source_(source), bos_token_(bos_token), eos_token_(eos_token)
3754
{
38-
_supports_tools = source.find("tools") != std::string::npos;
39-
_requires_object_arguments =
40-
source.find("tool_call.arguments | items") != std::string::npos
41-
|| source.find("tool_call.arguments | tojson") != std::string::npos;
42-
_supports_system_role = source.find("System role not supported") == std::string::npos;
43-
44-
_template_root = minja::Parser::parse(_source, {
55+
template_root_ = minja::Parser::parse(source_, {
4556
/* .trim_blocks = */ true,
4657
/* .lstrip_blocks = */ true,
4758
/* .keep_trailing_newline = */ false,
4859
});
60+
supports_tools_ = source.find("tools") != std::string::npos;
61+
62+
auto renders_string_arguments =
63+
try_render({
64+
{
65+
{"role", "user"},
66+
{"content", "Hey"}
67+
},
68+
{
69+
{"role", "assistant"},
70+
{"tool_calls", json::array({
71+
{
72+
{"id", "call_1___"},
73+
{"type", "function"},
74+
{"function", {
75+
{"arguments", "{\"code\": \"print('Hello, World!')\"}"},
76+
{"name", "ipython"},
77+
}},
78+
},
79+
})},
80+
}
81+
}, {}, false).find("{\"code\": \"print") != std::string::npos;
82+
if (!renders_string_arguments) {
83+
auto renders_object_arguments =
84+
try_render({
85+
{
86+
{"role", "user"},
87+
{"content", "Hey"}
88+
},
89+
{
90+
{"role", "assistant"},
91+
{"tool_calls", json::array({
92+
{
93+
{"id", "call_1___"},
94+
{"type", "function"},
95+
{"function", {
96+
{"arguments", {
97+
{"code", "print('Hello, World!')"},
98+
}},
99+
{"name", "ipython"},
100+
}},
101+
},
102+
})},
103+
}
104+
}, {}, false).find("{\"code\": \"print") != std::string::npos;
105+
requires_object_arguments_ = renders_object_arguments;
106+
}
107+
supports_parallel_tool_calls_ = source.find("tool_call_id") != std::string::npos;
108+
109+
supports_system_role_ = try_render({
110+
{{"role", "system"}, {"content", "<System Needle>"}},
111+
{{"role", "user"}, {"content", "Hey"}}
112+
}, {}, false).find("<System Needle>") != std::string::npos;
49113
}
50114

51-
const std::string & source() const { return _source; }
52-
bool supports_tools() const { return _supports_tools; }
115+
const std::string & source() const { return source_; }
116+
bool supports_tools() const { return supports_tools_; }
117+
bool supports_parallel_tool_calls() const { return supports_parallel_tool_calls_; }
53118

54119
std::string apply(
55120
const nlohmann::ordered_json & messages,
56121
const nlohmann::ordered_json & tools,
57122
bool add_generation_prompt,
58123
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
59124
{
60-
auto actual_messages = messages;
125+
json actual_messages;
61126

62127
// First, "fix" messages so they have a chance to be rendered correctly by the template
63128

64-
if (_requires_object_arguments || !_supports_system_role) {
129+
if (requires_object_arguments_ || !supports_system_role_ || !supports_tools_) {
130+
actual_messages = json::array();
131+
65132
std::string pending_system;
66133
auto flush_sys = [&]() {
67134
if (!pending_system.empty()) {
@@ -72,13 +139,66 @@ class chat_template {
72139
pending_system.clear();
73140
}
74141
};
75-
for (auto & message : actual_messages) {
142+
for (const auto & message_ : messages) {
143+
auto message = message_;
76144
if (!message.contains("role") || !message.contains("content")) {
77145
throw std::runtime_error("message must have 'role' and 'content' fields: " + message.dump());
78146
}
79147
std::string role = message.at("role");
80148

81-
if (!message["content"].is_null() && !_supports_system_role) {
149+
if (message.contains("tool_calls")) {
150+
if (requires_object_arguments_ || !supports_tools_) {
151+
for (auto & tool_call : message.at("tool_calls")) {
152+
if (tool_call["type"] == "function") {
153+
auto & function = tool_call.at("function");
154+
std::string arguments = function.at("arguments");
155+
function["arguments"] = json::parse(arguments);
156+
}
157+
}
158+
}
159+
if (!supports_tools_) {
160+
auto content = message.at("content");
161+
auto tool_calls = json::array();
162+
for (const auto & tool_call : message.at("tool_calls")) {
163+
if (tool_call.at("type") != "function") {
164+
continue;
165+
}
166+
const auto & function = tool_call.at("function");
167+
auto tc = json {
168+
{"name", function.at("name")},
169+
{"arguments", function.at("arguments")},
170+
};
171+
if (tool_call.contains("id")) {
172+
tc["id"] = tool_call["id"];
173+
}
174+
tool_calls.push_back(tc);
175+
}
176+
auto obj = json {
177+
{"tool_calls", tool_calls},
178+
};
179+
if (!content.is_null() && content != "") {
180+
obj["content"] = content;
181+
}
182+
message["content"] = obj.dump(2);
183+
message.erase("tool_calls");
184+
}
185+
}
186+
if (!supports_tools_ && role == "tool") {
187+
message["role"] = "user";
188+
auto obj = json {
189+
{"tool_response", {
190+
{"tool", message.at("name")},
191+
{"content", message.at("content")},
192+
}},
193+
};
194+
if (message.contains("tool_call_id")) {
195+
obj["tool_response"]["tool_call_id"] = message.at("tool_call_id");
196+
}
197+
message["content"] = obj.dump(2);
198+
message.erase("name");
199+
}
200+
201+
if (!message["content"].is_null() && !supports_system_role_) {
82202
std::string content = message.at("content");
83203
if (role == "system") {
84204
if (!pending_system.empty()) pending_system += "\n";
@@ -95,24 +215,18 @@ class chat_template {
95215
}
96216
}
97217
}
98-
if (_requires_object_arguments && message.contains("tool_calls")) {
99-
for (auto & tool_call : message.at("tool_calls")) {
100-
if (tool_call["type"] == "function") {
101-
auto & function = tool_call.at("function");
102-
std::string arguments = function.at("arguments");
103-
function["arguments"] = json::parse(arguments);
104-
}
105-
}
106-
}
218+
actual_messages.push_back(message);
107219
}
108220
flush_sys();
221+
} else {
222+
actual_messages = messages;
109223
}
110224

111225
auto context = minja::Context::make(json({
112226
{"messages", actual_messages},
113227
{"add_generation_prompt", add_generation_prompt},
114-
{"bos_token", _bos_token},
115-
{"eos_token", _eos_token},
228+
{"bos_token", bos_token_},
229+
{"eos_token", eos_token_},
116230
}));
117231

118232
if (!tools.is_null()) {
@@ -126,8 +240,8 @@ class chat_template {
126240
}
127241
}
128242

129-
return _template_root->render(context);
243+
return template_root_->render(context);
130244
}
131245
};
132246

133-
} // namespace minja
247+
} // namespace minja

‎include/minja/minja.hpp

+3-3
Original file line numberDiff line numberDiff line change
@@ -1009,7 +1009,7 @@ class FilterNode : public TemplateNode {
10091009
throw std::runtime_error("Filter must be a callable: " + filter_value.dump());
10101010
}
10111011
std::string rendered_body = body->render(context);
1012-
1012+
10131013
ArgumentsValue filter_args = {{Value(rendered_body)}, {}};
10141014
auto result = filter_value.call(context, filter_args);
10151015
out << result.to_str();
@@ -1181,7 +1181,7 @@ class UnaryOpExpr : public Expression {
11811181
case Op::Expansion:
11821182
case Op::ExpansionDict:
11831183
throw std::runtime_error("Expansion operator is only supported in function calls and collections");
1184-
1184+
11851185
}
11861186
throw std::runtime_error("Unknown unary operator");
11871187
}
@@ -2557,7 +2557,7 @@ inline std::shared_ptr<Context> Context::builtins() {
25572557
return (int64_t) items.size();
25582558
}));
25592559
globals.set("safe", simple_function("safe", { "value" }, [](const std::shared_ptr<Context> &, Value & args) -> Value {
2560-
return args.at("value");
2560+
return args.at("value").to_str();
25612561
}));
25622562
globals.set("string", simple_function("string", { "value" }, [](const std::shared_ptr<Context> &, Value & args) -> Value {
25632563
return args.at("value").to_str();

0 commit comments

Comments
 (0)