diff --git a/include/ortx_tokenizer.h b/include/ortx_tokenizer.h index 015b90a33..53d3b47c1 100644 --- a/include/ortx_tokenizer.h +++ b/include/ortx_tokenizer.h @@ -276,6 +276,30 @@ extError_t ORTX_API_CALL OrtxApplyChatTemplate(const OrtxTokenizer* tokenizer, c const char* input, const char* tools, OrtxTensorResult** output, bool add_generation_prompt, bool tokenize); +/** + * @brief Applies a chat template with additional template context values. + * + * Behaves like OrtxApplyChatTemplate, while also adding the properties from + * template_kwargs to the chat template context. template_kwargs must be a + * null-terminated JSON object or null. Core context properties such as messages, + * tools, and add_generation_prompt cannot be overridden. + * + * @param tokenizer Pointer to an OrtxTokenizer used for template processing. + * @param template_str Null-terminated string representing the chat template; can be null if tokenizer.json has one. + * @param input Null-terminated string containing the input to be processed. + * @param tools Null-terminated string containing the function tools. + * @param template_kwargs Null-terminated JSON object containing additional template context values; can be null. + * @param output Pointer to an OrtxTensorResult that will be populated with the output strings, + * if tokenize is true, the ids will be in the output as indexed 1. + * @param add_generation_prompt Indicates whether to add a generation prompt to the output. + * @param tokenize Indicates whether to tokenize the templated text to IDs. + * @return extError_t Returns an error code indicating success or the type of failure. + */ +extError_t ORTX_API_CALL OrtxApplyChatTemplateWithOptions(const OrtxTokenizer* tokenizer, const char* template_str, + const char* input, const char* tools, + const char* template_kwargs, OrtxTensorResult** output, + bool add_generation_prompt, bool tokenize); + #ifdef __cplusplus } #endif diff --git a/shared/api/c_api_tokenizer.cc b/shared/api/c_api_tokenizer.cc index 7584d7052..052f00b62 100644 --- a/shared/api/c_api_tokenizer.cc +++ b/shared/api/c_api_tokenizer.cc @@ -467,8 +467,16 @@ extError_t ORTX_API_CALL OrtxApplyChatTemplate(const OrtxTokenizer* tokenizer, c const char* input, const char* tools, OrtxTensorResult** output, bool add_generation_prompt, bool tokenize) { - if (tokenizer == nullptr && template_str == nullptr) { - ReturnableStatus::last_error_message_ = "both tokenizer and template_str are null, no template to apply"; + return OrtxApplyChatTemplateWithOptions(tokenizer, template_str, input, tools, nullptr, output, + add_generation_prompt, tokenize); +} + +extError_t ORTX_API_CALL OrtxApplyChatTemplateWithOptions(const OrtxTokenizer* tokenizer, const char* template_str, + const char* input, const char* tools, + const char* template_kwargs, OrtxTensorResult** output, + bool add_generation_prompt, bool tokenize) { + if (tokenizer == nullptr) { + ReturnableStatus::last_error_message_ = "tokenizer is null"; return kOrtxErrorInvalidArgument; } @@ -485,7 +493,8 @@ extError_t ORTX_API_CALL OrtxApplyChatTemplate(const OrtxTokenizer* tokenizer, c std::string text; std::vector ids_vec; - status = token_ptr->ApplyChatTemplate(template_str, input, tools, text, ids_vec, add_generation_prompt, tokenize); + status = token_ptr->ApplyChatTemplate(template_str, input, tools, template_kwargs, text, ids_vec, + add_generation_prompt, tokenize); if (status.IsOk()) { auto result = std::make_unique(); std::vector> tensors; diff --git a/shared/api/chat_template.cc b/shared/api/chat_template.cc index 1b454a60e..23bd8aca9 100644 --- a/shared/api/chat_template.cc +++ b/shared/api/chat_template.cc @@ -376,8 +376,9 @@ std::string normalize_tool_quotes(const std::string& input) { } OrtxStatus TokenizerImpl::ApplyChatTemplate(const char* template_str, const char* message, const char* tools, - std::string& output, std::vector& ids_vec, - bool add_generation_prompt, bool tokenize) const { + const char* template_kwargs, std::string& output, + std::vector& ids_vec, bool add_generation_prompt, + bool tokenize) const { OrtxStatus status; std::string input_str = minja::normalize_newlines(message); @@ -408,7 +409,21 @@ OrtxStatus TokenizerImpl::ApplyChatTemplate(const char* template_str, const char throw std::runtime_error("Invalid or unsupported chat template."); } - std::shared_ptr context; + json context_values = json::object(); + if (template_kwargs) { + if (*template_kwargs == '\0') { + throw std::runtime_error("template_kwargs must be a JSON object or null."); + } + auto parsed_kwargs = json::parse(minja::normalize_newlines(template_kwargs), nullptr, + /*allow_exceptions=*/false); + if (parsed_kwargs.is_discarded()) { + throw std::runtime_error("Invalid template_kwargs JSON."); + } + if (!parsed_kwargs.is_object()) { + throw std::runtime_error("template_kwargs must be a JSON object."); + } + context_values = std::move(parsed_kwargs); + } // Check Phi-4-mini tool call case for quote normalization bool phi_4_mini = false; @@ -462,20 +477,16 @@ OrtxStatus TokenizerImpl::ApplyChatTemplate(const char* template_str, const char tools_json = NormalizeTools(tools_str.c_str()); } - // Add tools to the context - context = minja::Context::make(json({ - {"messages", actual_messages}, - {"tools", tools_json}, - {"add_generation_prompt", add_generation_prompt}, - })); + context_values["tools"] = std::move(tools_json); } else { - // No tools input, just use the messages - context = minja::Context::make(json({ - {"messages", actual_messages}, - {"add_generation_prompt", add_generation_prompt}, - })); + context_values.erase("tools"); } + // Core request values take precedence over additional template kwargs. + context_values["messages"] = std::move(actual_messages); + context_values["add_generation_prompt"] = add_generation_prompt; + auto context = minja::Context::make(std::move(context_values)); + // Set required context values context->set("strftime_now", minja::Value::callable(strftime_function)); context->set("bos_token", tok_config_->bos_token_); diff --git a/shared/api/tokenizer_impl.h b/shared/api/tokenizer_impl.h index c919c8bff..7dfd53952 100644 --- a/shared/api/tokenizer_impl.h +++ b/shared/api/tokenizer_impl.h @@ -89,7 +89,8 @@ class TokenizerImpl : public OrtxObjectImpl { OrtxStatus Id2Token(extTokenId_t id, std::string& token, TokenizerDecodingState** state, bool skip_special_tokens) const; OrtxStatus GetDecoderPromptIds(size_t batch_size, const char* lang, const char* task, int no_timestamps, std::vector>& t_ids) const; - OrtxStatus ApplyChatTemplate(const char* template_str, const char* message, const char* tools, std::string& output, + OrtxStatus ApplyChatTemplate(const char* template_str, const char* message, const char* tools, + const char* template_kwargs, std::string& output, std::vector& ids_vec, bool add_generation_prompt, bool tokenize) const; private: diff --git a/test/pp_api_test/test_tokenizer_chat.cc b/test/pp_api_test/test_tokenizer_chat.cc index 90e5f3f03..211a55975 100644 --- a/test/pp_api_test/test_tokenizer_chat.cc +++ b/test/pp_api_test/test_tokenizer_chat.cc @@ -2255,4 +2255,118 @@ TEST(OrtxTokenizerTest, ChatTemplateDivisionByZero) { messages_json.c_str(), nullptr, result.ToBeAssigned(), false, false); EXPECT_NE(err, kOrtxOK) << "Expected modulo by zero to return an error."; } +} + +TEST(OrtxTokenizerTest, ChatTemplateAcceptsTypedTemplateKwargs) { + OrtxObjectPtr tokenizer(OrtxCreateTokenizer, "data/phi-4-base"); + ASSERT_EQ(tokenizer.Code(), kOrtxOK) << OrtxGetLastErrorMessage(); + + const std::string template_str = + R"({% if enable_thinking is defined and not enable_thinking %}NO_THINK{% else %}THINK{% endif %}|{{ reasoning_effort }}|{{ level }})"; + const std::string messages_json = R"([{"role":"user","content":"Hello"}])"; + const std::string template_kwargs = + R"({"enable_thinking":false,"reasoning_effort":"low","level":2})"; + OrtxObjectPtr result; + + auto err = OrtxApplyChatTemplateWithOptions( + tokenizer.get(), template_str.c_str(), messages_json.c_str(), nullptr, + template_kwargs.c_str(), result.ToBeAssigned(), true, false); + ASSERT_EQ(err, kOrtxOK) << OrtxGetLastErrorMessage(); + + OrtxObjectPtr tensor; + ASSERT_EQ(OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()), kOrtxOK); + const char* text = nullptr; + ASSERT_EQ(OrtxGetTensorData(tensor.get(), reinterpret_cast(&text), nullptr, nullptr), kOrtxOK); + EXPECT_STREQ(text, "NO_THINK|low|2"); +} + +TEST(OrtxTokenizerTest, ChatTemplateKwargsCannotOverrideCoreContext) { + OrtxObjectPtr tokenizer(OrtxCreateTokenizer, "data/phi-4-base"); + ASSERT_EQ(tokenizer.Code(), kOrtxOK) << OrtxGetLastErrorMessage(); + + const std::string template_str = + R"({{ messages[0].content }}|{% if add_generation_prompt %}GEN{% else %}NO_GEN{% endif %}|{% if tools is defined %}TOOLS{% else %}NO_TOOLS{% endif %})"; + const std::string messages_json = R"([{"role":"user","content":"Hello"}])"; + const std::string template_kwargs = + R"({"messages":[{"role":"user","content":"Override"}],"add_generation_prompt":false,"tools":[{"name":"override"}]})"; + OrtxObjectPtr result; + + auto err = OrtxApplyChatTemplateWithOptions( + tokenizer.get(), template_str.c_str(), messages_json.c_str(), nullptr, + template_kwargs.c_str(), result.ToBeAssigned(), true, false); + ASSERT_EQ(err, kOrtxOK) << OrtxGetLastErrorMessage(); + + OrtxObjectPtr tensor; + ASSERT_EQ(OrtxTensorResultGetAt(result.get(), 0, tensor.ToBeAssigned()), kOrtxOK); + const char* text = nullptr; + ASSERT_EQ(OrtxGetTensorData(tensor.get(), reinterpret_cast(&text), nullptr, nullptr), kOrtxOK); + EXPECT_STREQ(text, "Hello|GEN|NO_TOOLS"); +} + +TEST(OrtxTokenizerTest, ChatTemplateRejectsInvalidTemplateKwargs) { + OrtxObjectPtr tokenizer(OrtxCreateTokenizer, "data/phi-4-base"); + ASSERT_EQ(tokenizer.Code(), kOrtxOK) << OrtxGetLastErrorMessage(); + + const std::string messages_json = R"([{"role":"user","content":"Hello"}])"; + OrtxObjectPtr result; + + auto empty_string = OrtxApplyChatTemplateWithOptions( + tokenizer.get(), "{{ messages[0].content }}", messages_json.c_str(), nullptr, + "", result.ToBeAssigned(), true, false); + EXPECT_EQ(empty_string, kOrtxErrorInvalidArgument); + EXPECT_STREQ(OrtxGetLastErrorMessage(), "template_kwargs must be a JSON object or null."); + + auto invalid_json = OrtxApplyChatTemplateWithOptions( + tokenizer.get(), "{{ messages[0].content }}", messages_json.c_str(), nullptr, + "{", result.ToBeAssigned(), true, false); + EXPECT_EQ(invalid_json, kOrtxErrorInvalidArgument); + EXPECT_STREQ(OrtxGetLastErrorMessage(), "Invalid template_kwargs JSON."); + + auto non_object = OrtxApplyChatTemplateWithOptions( + tokenizer.get(), "{{ messages[0].content }}", messages_json.c_str(), nullptr, + "[]", result.ToBeAssigned(), true, false); + EXPECT_EQ(non_object, kOrtxErrorInvalidArgument); + EXPECT_STREQ(OrtxGetLastErrorMessage(), "template_kwargs must be a JSON object."); +} + +TEST(OrtxTokenizerTest, ChatTemplateRejectsNullTokenizerWithExplicitTemplate) { + const std::string messages_json = R"([{"role":"user","content":"Hello"}])"; + OrtxObjectPtr result; + + auto err = OrtxApplyChatTemplateWithOptions( + nullptr, "{{ messages[0].content }}", messages_json.c_str(), nullptr, + nullptr, result.ToBeAssigned(), true, false); + EXPECT_EQ(err, kOrtxErrorInvalidArgument); + EXPECT_STREQ(OrtxGetLastErrorMessage(), "tokenizer is null"); +} + +TEST(OrtxTokenizerTest, LegacyChatTemplateApiMatchesNullTemplateKwargs) { + OrtxObjectPtr tokenizer(OrtxCreateTokenizer, "data/phi-4-base"); + ASSERT_EQ(tokenizer.Code(), kOrtxOK) << OrtxGetLastErrorMessage(); + + const std::string template_str = R"({{ messages[0].content }}|{{ add_generation_prompt }})"; + const std::string messages_json = R"([{"role":"user","content":"Hello"}])"; + OrtxObjectPtr legacy_result; + OrtxObjectPtr options_result; + + ASSERT_EQ(OrtxApplyChatTemplate( + tokenizer.get(), template_str.c_str(), messages_json.c_str(), nullptr, + legacy_result.ToBeAssigned(), true, false), + kOrtxOK); + ASSERT_EQ(OrtxApplyChatTemplateWithOptions( + tokenizer.get(), template_str.c_str(), messages_json.c_str(), nullptr, nullptr, + options_result.ToBeAssigned(), true, false), + kOrtxOK); + + OrtxObjectPtr legacy_tensor; + OrtxObjectPtr options_tensor; + ASSERT_EQ(OrtxTensorResultGetAt(legacy_result.get(), 0, legacy_tensor.ToBeAssigned()), kOrtxOK); + ASSERT_EQ(OrtxTensorResultGetAt(options_result.get(), 0, options_tensor.ToBeAssigned()), kOrtxOK); + const char* legacy_text = nullptr; + const char* options_text = nullptr; + ASSERT_EQ(OrtxGetTensorData(legacy_tensor.get(), reinterpret_cast(&legacy_text), nullptr, nullptr), + kOrtxOK); + ASSERT_EQ(OrtxGetTensorData(options_tensor.get(), reinterpret_cast(&options_text), nullptr, nullptr), + kOrtxOK); + EXPECT_STREQ(legacy_text, options_text); } \ No newline at end of file