[REFACTOR] Refactor JSONFFI Conv template (#2331)

This PR refactors JSONFFI conv template to use immutable processing.
This helps to prevent bugs from multiple requests and concurrent
access to the conversation data structure.

It also reduces the need to deep copy the struct.

Author: tqchen

cpp/json_ffi/conv_template.cc

@@ -131,7 +131,7 @@ ModelConfig ModelConfig::FromJSON(const picojson::object& json_obj) {
 
 /****************** Conversation template ******************/
 
-std::map<MessagePlaceholders, std::string> PLACEHOLDERS = {
+std::unordered_map<MessagePlaceholders, std::string> PLACEHOLDERS = {
     {MessagePlaceholders::SYSTEM, "{system_message}"},
     {MessagePlaceholders::USER, "{user_message}"},
     {MessagePlaceholders::ASSISTANT, "{assistant_message}"},
@@ -153,120 +153,213 @@ Conversation::Conversation()
                       {"assistant", PLACEHOLDERS[MessagePlaceholders::ASSISTANT]},
                       {"tool", PLACEHOLDERS[MessagePlaceholders::TOOL]}}) {}
 
-Result<std::vector<Data>> Conversation::AsPrompt(ModelConfig config, DLDevice device) {
-  using TResult = Result<std::vector<Data>>;
-  // Get the system message
-  std::string system_msg = system_template;
-  size_t pos = system_msg.find(PLACEHOLDERS[MessagePlaceholders::SYSTEM]);
+std::string Conversation::GetSystemText(const std::string& system_msg) const {
+  std::string system_text = this->system_template;
+  static std::string system_placeholder = PLACEHOLDERS[MessagePlaceholders::SYSTEM];
+  size_t pos = system_text.find(system_placeholder);
   if (pos != std::string::npos) {
-    system_msg.replace(pos, PLACEHOLDERS[MessagePlaceholders::SYSTEM].length(),
-                       this->system_message);
+    system_text.replace(pos, system_placeholder.length(), system_msg);
   }
+  return system_text;
+}
 
-  // Get the message strings
-  std::vector<Data> message_list;
-  std::vector<std::string> separators = seps;
-  if (separators.size() == 1) {
-    separators.push_back(separators[0]);
+std::string Conversation::GetRoleText(const std::string& role, const std::string& content,
+                                      const std::optional<std::string>& fn_call_string) const {
+  std::string role_text = this->role_templates.at(role);
+  std::string placeholder = PLACEHOLDERS[MessagePlaceholderFromString(role)];
+  size_t pos = role_text.find(placeholder);
+  if (pos != std::string::npos) {
+    role_text.replace(pos, placeholder.length(), content);
+  }
+  if (fn_call_string) {
+    // replace placeholder[FUNCTION] with function_string
+    // this assumes function calling is used for a single request scenario only
+    pos = role_text.find(PLACEHOLDERS[MessagePlaceholders::FUNCTION]);
+    if (pos != std::string::npos) {
+      role_text.replace(pos, PLACEHOLDERS[MessagePlaceholders::FUNCTION].length(),
+                        fn_call_string.value());
+    }
   }
+  return role_text;
+}
 
-  if (!system_msg.empty()) {
-    system_msg += separators[0];
-    message_list.push_back(TextData(system_message));
+/// Try to detect if function calling is needed, if so, return the function calling string
+Result<std::optional<std::string>> TryGetFunctionCallingString(
+    const Conversation& conv, const ChatCompletionRequest& request) {
+  using TResult = Result<std::optional<std::string>>;
+  if (!request.tools.has_value() ||
+      (request.tool_choice.has_value() && request.tool_choice.value() == "none")) {
+    return TResult::Ok(std::nullopt);
+  }
+  std::vector<ChatTool> tools_ = request.tools.value();
+  std::string tool_choice_ = request.tool_choice.value();
+
+  // TODO: support with tool choice as dict
+  for (const auto& tool : tools_) {
+    if (tool.function.name == tool_choice_) {
+      picojson::value function_str(tool.function.AsJSON());
+      return TResult::Ok(function_str.serialize());
+    }
   }
 
-  for (int i = 0; i < messages.size(); i++) {
-    std::string role = messages[i].role;
-    // Todo(mlc-team): support content to be a single string.
-    std::optional<std::vector<std::unordered_map<std::string, std::string>>> content =
-        messages[i].content;
-    if (roles.find(role) == roles.end()) {
-      return TResult::Error("Role \"" + role + "\" is not supported");
-    }
+  if (tool_choice_ != "auto") {
+    return TResult::Error("Invalid tool_choice value in the request: " + tool_choice_);
+  }
+
+  picojson::array function_list;
+  for (const auto& tool : tools_) {
+    function_list.push_back(picojson::value(tool.function.AsJSON()));
+  }
 
-    std::string separator = separators[role == "assistant"];  // check assistant role
+  picojson::value function_list_json(function_list);
+  return TResult::Ok(function_list_json.serialize());
+};
 
-    // If content is empty, add the role and separator
-    // assistant's turn to generate text
-    if (!content.has_value()) {
-      message_list.push_back(TextData(roles[role] + role_empty_sep));
-      continue;
-    }
+Result<std::vector<Data>> CreatePrompt(const Conversation& conv,
+                                       const ChatCompletionRequest& request,
+                                       const ModelConfig& config, DLDevice device) {
+  using TResult = Result<std::vector<Data>>;
+
+  Result<std::optional<std::string>> fn_call_str_tmp = TryGetFunctionCallingString(conv, request);
+  if (fn_call_str_tmp.IsErr()) {
+    return TResult::Error(fn_call_str_tmp.UnwrapErr());
+  }
+  std::optional<std::string> fn_call_string = fn_call_str_tmp.Unwrap();
+
+  // Handle system message
+  // concz
+  bool has_custom_system = false;
+  std::string custom_system_inputs;
 
-    std::string message = "";
-    std::string role_prefix = "";
-    // Do not append role prefix if this is the first message and there
-    // is already a system message
-    if (add_role_after_system_message || system_msg.empty() || i != 0) {
-      role_prefix = roles[role] + role_content_sep;
+  auto f_populate_system_message = [&](const std::vector<ChatCompletionMessage>& msg_vec) {
+    for (ChatCompletionMessage msg : msg_vec) {
+      if (msg.role == "system") {
+        ICHECK(msg.content.IsText()) << "System message must be text";
+        custom_system_inputs += msg.content.Text();
+        has_custom_system = true;
+      }
     }
+  };
+  // go through messages in template and passed in.
+  f_populate_system_message(conv.messages);
+  f_populate_system_message(request.messages);
 
-    message += role_prefix;
+  // pending text records the text to be put into data
+  // we lazily accumulate the pending text
+  // to reduce amount of segments in the Data vector
+  std::string pending_text =
+      conv.GetSystemText(has_custom_system ? custom_system_inputs : conv.system_message);
 
-    for (const auto& item : content.value()) {
-      auto it_type = item.find("type");
-      if (it_type == item.end()) {
-        return TResult::Error("The content of a message does not have \"type\" field");
+  // the seperator after system message.
+  if (!pending_text.empty()) {
+    pending_text += conv.seps[0];
+  }
+
+  // Get the message strings
+  std::vector<Data> message_list;
+  size_t non_system_msg_count = 0;
+
+  // returns error if error happens
+  auto f_process_messages =
+      [&](const std::vector<ChatCompletionMessage>& msg_vec) -> std::optional<TResult> {
+    for (size_t i = 0; i < msg_vec.size(); ++i) {
+      const ChatCompletionMessage& msg = msg_vec[i];
+      auto role_it = conv.roles.find(msg.role);
+      if (role_it == conv.roles.end()) {
+        return TResult::Error("Role \"" + msg.role + "\" is not supported");
       }
-      if (it_type->second == "text") {
-        auto it_text = item.find("text");
-        if (it_text == item.end()) {
-          return TResult::Error("The text type content of a message does not have \"text\" field");
-        }
-        // replace placeholder[ROLE] with input message from role
-        std::string role_text = role_templates[role];
-        std::string placeholder = PLACEHOLDERS[MessagePlaceholderFromString(role)];
-        size_t pos = role_text.find(placeholder);
-        if (pos != std::string::npos) {
-          role_text.replace(pos, placeholder.length(), it_text->second);
-        }
-        if (use_function_calling) {
-          // replace placeholder[FUNCTION] with function_string
-          // this assumes function calling is used for a single request scenario only
-          if (!function_string.has_value()) {
-            return TResult::Error(
-                "The function string in conversation template is not defined for function "
-                "calling.");
+      const std::string& role_name = role_it->second;
+      // skip system message as it is already processed
+      if (msg.role == "system") continue;
+      // skip when content is empty
+      if (msg.content.IsNull()) {
+        pending_text += role_name + conv.role_empty_sep;
+        continue;
+      }
+      ++non_system_msg_count;
+      // assistant uses conv.seps[1] if there are two seps
+      int sep_offset = msg.role == "assistant" ? 1 : 0;
+      const std::string& seperator = conv.seps[sep_offset % conv.seps.size()];
+      // setup role prefix
+      std::string role_prefix = "";
+      // Do not append role prefix if this is the first message and there is already a system
+      // message
+      if (conv.add_role_after_system_message || pending_text.empty() || non_system_msg_count != 1) {
+        role_prefix = role_name + conv.role_content_sep;
+      }
+      pending_text += role_prefix;
+
+      if (msg.content.IsParts()) {
+        for (const auto& item : msg.content.Parts()) {
+          auto it_type = item.find("type");
+          if (it_type == item.end()) {
+            return TResult::Error("The content of a message does not have \"type\" field");
           }
-          pos = role_text.find(PLACEHOLDERS[MessagePlaceholders::FUNCTION]);
-          if (pos != std::string::npos) {
-            role_text.replace(pos, PLACEHOLDERS[MessagePlaceholders::FUNCTION].length(),
-                              function_string.value());
+          if (it_type->second == "text") {
+            auto it_text = item.find("text");
+            if (it_text == item.end()) {
+              return TResult::Error(
+                  "The text type content of a message does not have \"text\" field");
+            }
+            // replace placeholder[ROLE] with input message from role
+            pending_text += conv.GetRoleText(msg.role, it_text->second, fn_call_string);
+          } else if (it_type->second == "image_url") {
+            if (item.find("image_url") == item.end()) {
+              return TResult::Error("Content should have an image_url field");
+            }
+            std::string image_url =
+                item.at("image_url");  // TODO(mlc-team): According to OpenAI API reference this
+                                       // should be a map, with a "url" key containing the URL, but
+                                       // we are just assuming this as the URL for now
+            std::string base64_image = image_url.substr(image_url.find(",") + 1);
+            Result<NDArray> image_data_res = LoadImageFromBase64(base64_image);
+            if (image_data_res.IsErr()) {
+              return TResult::Error(image_data_res.UnwrapErr());
+            }
+            if (!config.vision_config.has_value()) {
+              return TResult::Error("Vision config is required for image input");
+            }
+            int image_size = config.vision_config.value().image_size;
+            int patch_size = config.vision_config.value().patch_size;
+
+            int embed_size = (image_size * image_size) / (patch_size * patch_size);
+
+            auto image_ndarray = ClipPreprocessor(image_data_res.Unwrap(), image_size, device);
+            // lazily commit text data
+            if (pending_text.length() != 0) {
+              message_list.push_back(TextData(pending_text));
+              pending_text = "";
+            }
+            message_list.push_back(ImageData(image_ndarray, embed_size));
+          } else {
+            return TResult::Error("Unsupported content type: " + it_type->second);
           }
         }
-        message += role_text;
-      } else if (it_type->second == "image_url") {
-        if (item.find("image_url") == item.end()) {
-          return TResult::Error("Content should have an image_url field");
-        }
-        std::string image_url =
-            item.at("image_url");  // TODO(mlc-team): According to OpenAI API reference this
-                                   // should be a map, with a "url" key containing the URL, but
-                                   // we are just assuming this as the URL for now
-        std::string base64_image = image_url.substr(image_url.find(",") + 1);
-        Result<NDArray> image_data_res = LoadImageFromBase64(base64_image);
-        if (image_data_res.IsErr()) {
-          return TResult::Error(image_data_res.UnwrapErr());
-        }
-        if (!config.vision_config.has_value()) {
-          return TResult::Error("Vision config is required for image input");
-        }
-        int image_size = config.vision_config.value().image_size;
-        int patch_size = config.vision_config.value().patch_size;
-
-        int embed_size = (image_size * image_size) / (patch_size * patch_size);
-
-        auto image_ndarray = ClipPreprocessor(image_data_res.Unwrap(), image_size, device);
-        message_list.push_back(ImageData(image_ndarray, embed_size));
       } else {
-        return TResult::Error("Unsupported content type: " + it_type->second);
+        ICHECK(msg.content.IsText());
+        pending_text += conv.GetRoleText(msg.role, msg.content.Text(), fn_call_string);
       }
+      pending_text += seperator;
     }
+    return std::nullopt;
+  };
 
-    message += separator;
-    message_list.push_back(TextData(message));
+  if (auto err = f_process_messages(conv.messages)) {
+    return err.value();
+  }
+  if (auto err = f_process_messages(request.messages)) {
+    return err.value();
+  }
+  // append last assistant begin message
+  ChatCompletionMessage last_assistant_begin;
+  last_assistant_begin.role = "assistant";
+  last_assistant_begin.content = std::nullopt;
+  if (auto err = f_process_messages({last_assistant_begin})) {
+    return err.value();
+  }
+  if (pending_text.length() != 0) {
+    message_list.push_back(TextData(pending_text));
   }
-
   return TResult::Ok(message_list);
 }
 
@@ -383,7 +476,10 @@ Result<Conversation> Conversation::FromJSON(const picojson::object& json_obj) {
         content.push_back(std::move(item_map));
       }
     }
-    conv.messages.push_back({role_res.Unwrap(), content});
+    ChatCompletionMessage msg;
+    msg.role = role_res.Unwrap();
+    msg.content = content;
+    conv.messages.push_back(msg);
   }
 
   Result<picojson::array> seps_arr_res =
@@ -438,21 +534,6 @@ Result<Conversation> Conversation::FromJSON(const picojson::object& json_obj) {
     }
     conv.stop_token_ids.push_back(stop.get<int64_t>());
   }
-
-  Result<std::optional<std::string>> function_string_res =
-      json::LookupOptionalWithResultReturn<std::string>(json_obj, "function_string");
-  if (function_string_res.IsErr()) {
-    return TResult::Error(function_string_res.UnwrapErr());
-  }
-  conv.function_string = function_string_res.Unwrap();
-
-  Result<bool> use_function_calling_res = json::LookupOrDefaultWithResultReturn<bool>(
-      json_obj, "use_function_calling", conv.use_function_calling);
-  if (use_function_calling_res.IsErr()) {
-    return TResult::Error(use_function_calling_res.UnwrapErr());
-  }
-  conv.use_function_calling = use_function_calling_res.Unwrap();
-
   return TResult::Ok(conv);
 }