Skip to content

Commit

Permalink
Merge pull request #16 from google/tool-call
Browse files Browse the repository at this point in the history 
chat_template.hpp improvements
  • Loading branch information
@ochafik
ochafik authored Dec 30, 2024
2 parents 202aa2f + bbc02d6 commit 58f0ca6
Show file tree
Hide file tree
Showing 5 changed files with 212 additions and 55 deletions.
182 changes: 148 additions & 34 deletions include/minja/chat-template.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,47 +21,114 @@ class chat_template {
public:

private:
bool _supports_tools = true;
bool supports_tools_ = true;
// Meta-Llama-3.1-8B-Instruct's template expects arguments to be an object.
// Most other templates (and OpenAI's API) expect the arguments object to be stringified.
bool _requires_object_arguments = false;
bool _supports_system_role = true;
std::string _source;
std::string _bos_token;
std::string _eos_token;
std::shared_ptr<minja::TemplateNode> _template_root;
bool requires_object_arguments_ = false;
bool supports_system_role_ = true;
bool supports_parallel_tool_calls_ = false;
std::string source_;
std::string bos_token_;
std::string eos_token_;
std::shared_ptr<minja::TemplateNode> template_root_;

std::string try_render(
const nlohmann::ordered_json & messages,
const nlohmann::ordered_json & tools,
bool add_generation_prompt,
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
{
try {
auto prompt = apply(messages, tools, add_generation_prompt, extra_context);
// fprintf(stderr, "Prompt: %s\n", prompt.c_str());
return prompt;
} catch (const std::exception & e) {
// fprintf(stderr, "Error: %s\n", e.what());
return "";
}
}

public:
chat_template(const std::string & source, const std::string & bos_token, const std::string & eos_token)
: _source(source), _bos_token(bos_token), _eos_token(eos_token)
: source_(source), bos_token_(bos_token), eos_token_(eos_token)
{
_supports_tools = source.find("tools") != std::string::npos;
_requires_object_arguments =
source.find("tool_call.arguments | items") != std::string::npos
|| source.find("tool_call.arguments | tojson") != std::string::npos;
_supports_system_role = source.find("System role not supported") == std::string::npos;

_template_root = minja::Parser::parse(_source, {
template_root_ = minja::Parser::parse(source_, {
/* .trim_blocks = */ true,
/* .lstrip_blocks = */ true,
/* .keep_trailing_newline = */ false,
});
supports_tools_ = source.find("tools") != std::string::npos;

auto renders_string_arguments =
try_render({
{
{"role", "user"},
{"content", "Hey"}
},
{
{"role", "assistant"},
{"tool_calls", json::array({
{
{"id", "call_1___"},
{"type", "function"},
{"function", {
{"arguments", "{\"code\": \"print('Hello, World!')\"}"},
{"name", "ipython"},
}},
},
})},
}
}, {}, false).find("{\"code\": \"print") != std::string::npos;
if (!renders_string_arguments) {
auto renders_object_arguments =
try_render({
{
{"role", "user"},
{"content", "Hey"}
},
{
{"role", "assistant"},
{"tool_calls", json::array({
{
{"id", "call_1___"},
{"type", "function"},
{"function", {
{"arguments", {
{"code", "print('Hello, World!')"},
}},
{"name", "ipython"},
}},
},
})},
}
}, {}, false).find("{\"code\": \"print") != std::string::npos;
requires_object_arguments_ = renders_object_arguments;
}
supports_parallel_tool_calls_ = source.find("tool_call_id") != std::string::npos;

supports_system_role_ = try_render({
{{"role", "system"}, {"content", "<System Needle>"}},
{{"role", "user"}, {"content", "Hey"}}
}, {}, false).find("<System Needle>") != std::string::npos;
}

const std::string & source() const { return _source; }
bool supports_tools() const { return _supports_tools; }
const std::string & source() const { return source_; }
bool supports_tools() const { return supports_tools_; }
bool supports_parallel_tool_calls() const { return supports_parallel_tool_calls_; }

std::string apply(
const nlohmann::ordered_json & messages,
const nlohmann::ordered_json & tools,
bool add_generation_prompt,
const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
{
auto actual_messages = messages;
json actual_messages;

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

if (_requires_object_arguments || !_supports_system_role) {
if (requires_object_arguments_ || !supports_system_role_ || !supports_tools_) {
actual_messages = json::array();

std::string pending_system;
auto flush_sys = [&]() {
if (!pending_system.empty()) {
Expand All @@ -72,13 +139,66 @@ class chat_template {
pending_system.clear();
}
};
for (auto & message : actual_messages) {
for (const auto & message_ : messages) {
auto message = message_;
if (!message.contains("role") || !message.contains("content")) {
throw std::runtime_error("message must have 'role' and 'content' fields: " + message.dump());
}
std::string role = message.at("role");

if (!message["content"].is_null() && !_supports_system_role) {
if (message.contains("tool_calls")) {
if (requires_object_arguments_ || !supports_tools_) {
for (auto & tool_call : message.at("tool_calls")) {
if (tool_call["type"] == "function") {
auto & function = tool_call.at("function");
std::string arguments = function.at("arguments");
function["arguments"] = json::parse(arguments);
}
}
}
if (!supports_tools_) {
auto content = message.at("content");
auto tool_calls = json::array();
for (const auto & tool_call : message.at("tool_calls")) {
if (tool_call.at("type") != "function") {
continue;
}
const auto & function = tool_call.at("function");
auto tc = json {
{"name", function.at("name")},
{"arguments", function.at("arguments")},
};
if (tool_call.contains("id")) {
tc["id"] = tool_call["id"];
}
tool_calls.push_back(tc);
}
auto obj = json {
{"tool_calls", tool_calls},
};
if (!content.is_null() && content != "") {
obj["content"] = content;
}
message["content"] = obj.dump(2);
message.erase("tool_calls");
}
}
if (!supports_tools_ && role == "tool") {
message["role"] = "user";
auto obj = json {
{"tool_response", {
{"tool", message.at("name")},
{"content", message.at("content")},
}},
};
if (message.contains("tool_call_id")) {
obj["tool_response"]["tool_call_id"] = message.at("tool_call_id");
}
message["content"] = obj.dump(2);
message.erase("name");
}

if (!message["content"].is_null() && !supports_system_role_) {
std::string content = message.at("content");
if (role == "system") {
if (!pending_system.empty()) pending_system += "\n";
Expand All @@ -95,24 +215,18 @@ class chat_template {
}
}
}
if (_requires_object_arguments && message.contains("tool_calls")) {
for (auto & tool_call : message.at("tool_calls")) {
if (tool_call["type"] == "function") {
auto & function = tool_call.at("function");
std::string arguments = function.at("arguments");
function["arguments"] = json::parse(arguments);
}
}
}
actual_messages.push_back(message);
}
flush_sys();
} else {
actual_messages = messages;
}

auto context = minja::Context::make(json({
{"messages", actual_messages},
{"add_generation_prompt", add_generation_prompt},
{"bos_token", _bos_token},
{"eos_token", _eos_token},
{"bos_token", bos_token_},
{"eos_token", eos_token_},
}));

if (!tools.is_null()) {
Expand All @@ -126,8 +240,8 @@ class chat_template {
}
}

return _template_root->render(context);
return template_root_->render(context);
}
};

} // namespace minja
} // namespace minja
6 changes: 3 additions & 3 deletions include/minja/minja.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1009,7 +1009,7 @@ class FilterNode : public TemplateNode {
throw std::runtime_error("Filter must be a callable: " + filter_value.dump());
}
std::string rendered_body = body->render(context);

ArgumentsValue filter_args = {{Value(rendered_body)}, {}};
auto result = filter_value.call(context, filter_args);
out << result.to_str();
Expand Down Expand Up @@ -1181,7 +1181,7 @@ class UnaryOpExpr : public Expression {
case Op::Expansion:
case Op::ExpansionDict:
throw std::runtime_error("Expansion operator is only supported in function calls and collections");

}
throw std::runtime_error("Unknown unary operator");
}
Expand Down Expand Up @@ -2557,7 +2557,7 @@ inline std::shared_ptr<Context> Context::builtins() {
return (int64_t) items.size();
}));
globals.set("safe", simple_function("safe", { "value" }, [](const std::shared_ptr<Context> &, Value & args) -> Value {
return args.at("value");
return args.at("value").to_str();
}));
globals.set("string", simple_function("string", { "value" }, [](const std::shared_ptr<Context> &, Value & args) -> Value {
return args.at("value").to_str();
Expand Down
Loading

0 comments on commit 58f0ca6

Please sign in to comment.