sync: minja https://github.com/google/minja/pull/46

common/chat-template.hpp

@@ -33,6 +33,29 @@ struct chat_template_caps {
     bool requires_typed_content = false;
 };
 
+struct chat_template_inputs {
+    nlohmann::ordered_json messages;
+    nlohmann::ordered_json tools;
+    bool add_generation_prompt = true;
+    nlohmann::ordered_json extra_context;
+    std::chrono::system_clock::time_point now = std::chrono::system_clock::now();
+};
+
+struct chat_template_options {
+    bool apply_polyfills = true;
+    bool use_bos_token = true;
+    bool use_eos_token = true;
+    bool define_strftime_now = true;
+
+    bool polyfill_tools = true;
+    bool polyfill_tool_call_examples = true;
+    bool polyfill_tool_calls = true;
+    bool polyfill_tool_responses = true;
+    bool polyfill_system_role = true;
+    bool polyfill_object_arguments = true;
+    bool polyfill_typed_content = true;
+};
+
 class chat_template {
 
   private:
@@ -50,7 +73,18 @@ class chat_template {
         const nlohmann::ordered_json & extra_context = nlohmann::ordered_json()) const
     {
         try {
-            auto prompt = apply(messages, tools, add_generation_prompt, extra_context, /* adjust_inputs= */ false);
+            chat_template_inputs inputs;
+            inputs.messages = messages;
+            inputs.tools = tools;
+            inputs.add_generation_prompt = add_generation_prompt;
+            inputs.extra_context = extra_context;
+            // Use fixed date for tests
+            inputs.now = std::chrono::system_clock::from_time_t(0);
+
+            chat_template_options opts;
+            opts.apply_polyfills = false;
+
+            auto prompt = apply(inputs, opts);
             // fprintf(stderr, "try_raw_render: %s\n", prompt.c_str());
             return prompt;
         } catch (const std::exception & e) {
@@ -178,11 +212,15 @@ class chat_template {
             caps_.supports_tool_call_id = contains(out, "call_911_");
         }
 
+        try {
             if (!caps_.supports_tools) {
                 const json user_msg {
                     {"role", "user"},
                     {"content", "Hey"},
                 };
+                const json args {
+                    {"arg1", "some_value"},
+                };
                 const json tool_call_msg {
                     {"role", "assistant"},
                     {"content", nullptr},
@@ -193,30 +231,38 @@ class chat_template {
                             {"type", "function"},
                             {"function", {
                                 {"name", "tool_name"},
-                            {"arguments", (json {
+                                {"arguments", (caps_.requires_object_arguments ? args : json(minja::Value(args).dump(-1, /* to_json= */ true)))},
-                                {"arg1", "some_value"},
-                            }).dump()},
                             }},
                         },
                     })},
                 };
-            const json tools;
+                std::string prefix, full;
-            auto prefix = apply(json::array({user_msg}), tools, /* add_generation_prompt= */ true);
+                {
-            auto full = apply(json::array({user_msg, tool_call_msg}), tools, /* add_generation_prompt= */ false);
+                    chat_template_inputs inputs;
-            if (full.find(prefix) != 0 && prefix.length() > 0 && prefix[prefix.length() - 1] == '\n') {
+                    inputs.messages = json::array({user_msg});
-                prefix = prefix.substr(0, prefix.length() - 1);
+                    inputs.add_generation_prompt = true;
+                    prefix = apply(inputs);
                 }
+                {
+                    chat_template_inputs inputs;
+                    inputs.messages = json::array({user_msg, tool_call_msg});
+                    inputs.add_generation_prompt = false;
+                    full = apply(inputs);
+                }
+
                 if (full.find(prefix) != 0) {
                     if (prefix.rfind(eos_token_) == prefix.size() - eos_token_.size()) {
                         prefix = prefix.substr(0, prefix.size() - eos_token_.size());
-                } else {
-                    throw std::runtime_error("prefix not found at start of full: " + prefix + " vs " + full);
                     }
-            } else {
+                }
-
+                if (full.find(prefix) != 0) {
+                    fprintf(stderr, "Failed to infer a tool call example (possible template bug)\n");
                 }
                 tool_call_example_ = full.substr(prefix.size());
             }
+        } catch (const std::exception & e) {
+            fprintf(stderr, "Failed to generate tool call example: %s\n", e.what());
+        }
     }
 
     const std::string & source() const { return source_; }
@@ -225,27 +271,49 @@ class chat_template {
     const chat_template_caps & original_caps() const { return caps_; }
 
     std::string apply(
-        const nlohmann::ordered_json & messages,
+        const chat_template_inputs & inputs,
-        const nlohmann::ordered_json & tools,
+        const chat_template_options & opts = chat_template_options()) const
-        bool add_generation_prompt,
-        const nlohmann::ordered_json & extra_context = nlohmann::ordered_json(),
-        bool adjust_inputs = true) const
     {
         json actual_messages;
 
-        auto needs_adjustments = adjust_inputs && (false
+        auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
-            || !caps_.supports_system_role
+        auto has_tool_calls = false;
-            || !caps_.supports_tools
+        auto has_tool_responses = false;
-            || !caps_.supports_tool_responses
+        auto has_string_content = false;
-            || !caps_.supports_tool_calls
+        for (const auto & message : inputs.messages) {
-            || caps_.requires_object_arguments
+            if (!message["tool_calls"].is_null()) {
-            || caps_.requires_typed_content
+                has_tool_calls = true;
+            }
+            if (message["role"] == "tool") {
+                has_tool_responses = true;
+            }
+            if (message["content"].is_string()) {
+                has_string_content = true;
+            }
+        }
+
+        auto polyfill_system_role = opts.polyfill_system_role && !caps_.supports_system_role;
+        auto polyfill_tools = opts.polyfill_tools && has_tools && !caps_.supports_tools;
+        auto polyfill_tool_call_example = polyfill_tools && opts.polyfill_tool_call_examples;
+        auto polyfill_tool_calls = opts.polyfill_tool_calls && has_tool_calls && !caps_.supports_tool_calls;
+        auto polyfill_tool_responses = opts.polyfill_tool_responses && has_tool_responses && !caps_.supports_tool_responses;
+        auto polyfill_object_arguments = opts.polyfill_object_arguments && has_tool_calls && caps_.requires_object_arguments;
+        auto polyfill_typed_content = opts.polyfill_typed_content && has_string_content && caps_.requires_typed_content;
+
+        auto needs_polyfills = opts.apply_polyfills && (false
+            || polyfill_system_role
+            || polyfill_tools
+            || polyfill_tool_calls
+            || polyfill_tool_responses
+            || polyfill_object_arguments
+            || polyfill_typed_content
         );
-        if (needs_adjustments) {
+
+        if (needs_polyfills) {
             actual_messages = json::array();
 
             auto add_message = [&](const json & msg) {
-                if (caps_.requires_typed_content && msg.contains("content") && !msg.at("content").is_null() && msg.at("content").is_string()) {
+                if (polyfill_typed_content && msg.contains("content") && !msg.at("content").is_null() && msg.at("content").is_string()) {
                     actual_messages.push_back({
                         {"role", msg.at("role")},
                         {"content", {{
@@ -268,16 +336,14 @@ class chat_template {
                     pending_system.clear();
                 }
             };
-            auto needs_tools_in_system = !tools.is_null() && tools.size() > 0 && !caps_.supports_tools;
 
             json adjusted_messages;
-            if (needs_tools_in_system) {
+            if (polyfill_tools) {
-                adjusted_messages = add_system(messages,
+                adjusted_messages = add_system(inputs.messages,
-                    "\n\n"
+                    "You can call any of the following tools to satisfy the user's requests: " + minja::Value(inputs.tools).dump(2, /* to_json= */ true) +
-                    "You can call any of the following tools to satisfy the user's requests: " + tools.dump(2) + "\n\n"
+                    (!polyfill_tool_call_example || tool_call_example_.empty() ? "" : "\n\nExample tool call syntax:\n\n" + tool_call_example_));
-                    "Example tool call syntax:\n\n" + tool_call_example_ + "\n\n");
             } else {
-                adjusted_messages = messages;
+                adjusted_messages = inputs.messages;
             }
 
             for (const auto & message_ : adjusted_messages) {
@@ -288,7 +354,7 @@ class chat_template {
                 std::string role = message.at("role");
 
                 if (message.contains("tool_calls")) {
-                    if (caps_.requires_object_arguments || !caps_.supports_tool_calls) {
+                    if (polyfill_object_arguments || polyfill_tool_calls) {
                         for (auto & tool_call : message.at("tool_calls")) {
                             if (tool_call["type"] == "function") {
                                 auto & function = tool_call.at("function");
@@ -303,7 +369,7 @@ class chat_template {
                             }
                         }
                     }
-                    if (!caps_.supports_tool_calls) {
+                    if (polyfill_tool_calls) {
                         auto content = message.at("content");
                         auto tool_calls = json::array();
                         for (const auto & tool_call : message.at("tool_calls")) {
@@ -330,7 +396,7 @@ class chat_template {
                         message.erase("tool_calls");
                     }
                 }
-                if (!caps_.supports_tool_responses && role == "tool") {
+                if (polyfill_tool_responses && role == "tool") {
                     message["role"] = "user";
                     auto obj = json {
                         {"tool_response", {
@@ -347,7 +413,7 @@ class chat_template {
                     message.erase("name");
                 }
 
-                if (!message["content"].is_null() && !caps_.supports_system_role) {
+                if (!message["content"].is_null() && polyfill_system_role) {
                     std::string content = message.at("content");
                     if (role == "system") {
                         if (!pending_system.empty()) pending_system += "\n";
@@ -366,28 +432,40 @@ class chat_template {
                 }
                 add_message(message);
             }
-            if (!caps_.supports_system_role) {
             flush_sys();
-            }
         } else {
-            actual_messages = messages;
+            actual_messages = inputs.messages;
         }
 
         auto context = minja::Context::make(json({
             {"messages", actual_messages},
-            {"add_generation_prompt", add_generation_prompt},
+            {"add_generation_prompt", inputs.add_generation_prompt},
-            {"bos_token", bos_token_},
-            {"eos_token", eos_token_},
         }));
-
+        if (opts.use_bos_token) {
-        if (!tools.is_null()) {
+            context->set("bos_token", bos_token_);
-            auto tools_val = minja::Value(tools);
-            context->set("tools", tools_val);
         }
-        if (!extra_context.is_null()) {
+        if (opts.use_eos_token) {
-            for (auto & kv : extra_context.items()) {
+            context->set("eos_token", eos_token_);
-                minja::Value val(kv.value());
+        }
-                context->set(kv.key(), val);
+        if (opts.define_strftime_now) {
+            auto now = inputs.now;
+            context->set("strftime_now", Value::callable([now](const std::shared_ptr<minja::Context> &, minja::ArgumentsValue & args) {
+                args.expectArgs("strftime_now", {1, 1}, {0, 0});
+                auto format = args.args[0].get<std::string>();
+
+                auto time = std::chrono::system_clock::to_time_t(now);
+                auto local_time = *std::localtime(&time);
+                std::ostringstream ss;
+                ss << std::put_time(&local_time, format.c_str());
+                return ss.str();
+            }));
+        }
+        if (!inputs.tools.is_null()) {
+            context->set("tools", minja::Value(inputs.tools));
+        }
+        if (!inputs.extra_context.is_null()) {
+            for (auto & kv : inputs.extra_context.items()) {
+                context->set(kv.key(), minja::Value(kv.value()));
             }
         }
 
@@ -404,7 +482,7 @@ class chat_template {
             std::string existing_system = messages_with_system.at(0).at("content");
             messages_with_system[0] = json {
                 {"role", "system"},
-                {"content", existing_system + "\n" + system_prompt},
+                {"content", existing_system + "\n\n" + system_prompt},
             };
         } else {
             messages_with_system.insert(messages_with_system.begin(), json {

common/minja.hpp

@@ -2194,7 +2194,7 @@ private:
     }
 
     TemplateTokenVector tokenize() {
-      static std::regex comment_tok(R"(\{#([-~]?)(.*?)([-~]?)#\})");
+      static std::regex comment_tok(R"(\{#([-~]?)([\s\S\r\n]*?)([-~]?)#\})");
       static std::regex expr_open_regex(R"(\{\{([-~])?)");
       static std::regex block_open_regex(R"(^\{%([-~])?[\s\n\r]*)");
       static std::regex block_keyword_tok(R"((if|else|elif|endif|for|endfor|generation|endgeneration|set|endset|block|endblock|macro|endmacro|filter|endfilter|break|continue)\b)");
@@ -2615,6 +2615,7 @@ inline std::shared_ptr<Context> Context::builtins() {
   }));
   globals.set("join", simple_function("join", { "items", "d" }, [](const std::shared_ptr<Context> &, Value & args) {
     auto do_join = [](Value & items, const std::string & sep) {
+      if (!items.is_array()) throw std::runtime_error("object is not iterable: " + items.dump());
       std::ostringstream oss;
       auto first = true;
       for (size_t i = 0, n = items.size(); i < n; ++i) {
@@ -2695,6 +2696,10 @@ inline std::shared_ptr<Context> Context::builtins() {
     return Value::callable([=](const std::shared_ptr<Context> & context, ArgumentsValue & args) {
       args.expectArgs(is_select ? "select" : "reject", {2, (std::numeric_limits<size_t>::max)()}, {0, 0});
       auto & items = args.args[0];
+      if (items.is_null())
+        return Value::array();
+      if (!items.is_array()) throw std::runtime_error("object is not iterable: " + items.dump());
+
       auto filter_fn = context->get(args.args[1]);
       if (filter_fn.is_null()) throw std::runtime_error("Undefined filter: " + args.args[1].dump());
 
@@ -2772,6 +2777,7 @@ inline std::shared_ptr<Context> Context::builtins() {
       auto & items = args.args[0];
       if (items.is_null())
         return Value::array();
+      if (!items.is_array()) throw std::runtime_error("object is not iterable: " + items.dump());
       auto attr_name = args.args[1].get<std::string>();
 
       bool has_test = false;