Re-Prompt after toolcall

Author: bandoti

examples/main/main.cpp

@@ -90,39 +90,51 @@ static void sigint_handler(int signo) {
 
 class chat_formatter {
 public:
-    chat_formatter(common_params & params, std::vector<common_chat_msg> & chat_msgs, struct common_chat_templates * chat_templates)
+
+    struct result {
+        std::string formatted;
+        bool tool_was_called;
+    };
+
+    chat_formatter(common_params & params,
+                   std::vector<common_chat_msg> & chat_msgs,
+                   struct common_chat_templates * chat_templates)
+
         : params_(params), chat_msgs_(chat_msgs), chat_templates_(chat_templates) {}
 
 #ifdef LLAMA_USE_TOOLCALL
     chat_formatter(common_params & params,
                    std::vector<common_chat_msg> & chat_msgs,
                    struct common_chat_templates * chat_templates,
                    const llama_vocab * vocab,
-                   toolcall::client::ptr tc_client,
-                   common_chat_format * chat_format)
+                   toolcall::client::ptr tc_client)
 
-        : params_(params), chat_msgs_(chat_msgs), chat_templates_(chat_templates), vocab_(vocab), tc_client_(tc_client), chat_format_(chat_format) {}
+        : params_(params), chat_msgs_(chat_msgs), chat_templates_(chat_templates),
+          vocab_(vocab), tc_client_(tc_client),
+          chat_format_(COMMON_CHAT_FORMAT_CONTENT_ONLY),
+          formatted_() {}
 #endif
 
-    std::string operator () (const std::string & role, const std::string & content, [[maybe_unused]] bool use_toolcalls = false) {
+    chat_formatter::result operator() (const std::string & role, const std::string & content) {
+
+        common_chat_msg new_msg = common_chat_parse(content, chat_format_);
+        new_msg.role = role;
 
         common_chat_templates_inputs cinputs;
         cinputs.use_jinja = params_.use_jinja;
         cinputs.add_generation_prompt = (role == "user");
 #ifdef LLAMA_USE_TOOLCALL
-        if (tc_client_ != nullptr && use_toolcalls) {
+        if (tc_client_ != nullptr) {
             cinputs.tool_choice = common_chat_tool_choice_parse_oaicompat(tc_client_->tool_choice());
             cinputs.tools = common_chat_tools_parse_oaicompat(tc_client_->tool_list());
         }
 #endif
-        for (const auto & msg : chat_msgs_) {
-            cinputs.messages.push_back(common_chat_msg(msg));
-        }
-
-        common_chat_msg new_msg = common_chat_parse(content, *chat_format_);
-        new_msg.role = role;
+        cinputs.messages.assign(chat_msgs_.cbegin(), chat_msgs_.cend());
+        cinputs.messages.push_back(new_msg);
+        chat_msgs_.push_back(new_msg);
 
-        if (! new_msg.tool_calls.empty()) {
+        bool tool_was_called = false;
+        if (! new_msg.tool_calls.empty()) { // Call tool and re-prompt
             nlohmann::json result_array = nlohmann::json::array();
             for (const auto & tc : new_msg.tool_calls) {
                 toolcall::result_set res = tc_client_->call(tc.name, tc.arguments, tc.id);
@@ -132,21 +144,28 @@ class chat_formatter {
                     }
                 }
             }
-            new_msg.content += result_array.dump(-1);
+            common_chat_msg toolcall_msg;
+            toolcall_msg.role = "tool";
+            toolcall_msg.content = result_array.dump(-1);
+
+            cinputs.add_generation_prompt = true;
+            cinputs.messages.push_back(toolcall_msg);
+            chat_msgs_.push_back(toolcall_msg);
+
+            tool_was_called = true;
         }
 
-        cinputs.messages.push_back(new_msg);
         common_chat_params cparams = common_chat_templates_apply(chat_templates_, cinputs);
+        std::string formatted = cparams.prompt.substr(formatted_.size(), cparams.prompt.size());
+        formatted_ = cparams.prompt;
 
-        auto formatted = cparams.prompt;
-        chat_msgs_.push_back(new_msg);
         LOG_DBG("formatted: '%s'\n", formatted.c_str());
 
 #ifdef LLAMA_USE_TOOLCALL
-        if (chat_format_) *chat_format_ = cparams.format;
+        chat_format_ = cparams.format;
         common_chat_grammar_to_sampler(&cparams, vocab_, &params_.sampling);
 #endif
-        return formatted;
+        return chat_formatter::result{std::move(formatted), tool_was_called};
     }
 
 private:
@@ -157,7 +176,8 @@ class chat_formatter {
 #ifdef LLAMA_USE_TOOLCALL
     const llama_vocab * vocab_;
     toolcall::client::ptr tc_client_;
-    common_chat_format * chat_format_;
+    common_chat_format chat_format_;
+    std::string formatted_;
 #endif
 };
 
@@ -355,8 +375,7 @@ int main(int argc, char ** argv) {
     if (tc_client) {
         tc_client->initialize();
     }
-    common_chat_format chat_format = COMMON_CHAT_FORMAT_CONTENT_ONLY;
-    chat_formatter chat_add_and_format(params, chat_msgs, chat_templates.get(), vocab, tc_client, &chat_format);
+    chat_formatter chat_add_and_format(params, chat_msgs, chat_templates.get(), vocab, tc_client);
 #else
     chat_formatter chat_add_and_format(params, chat_msgs, chat_templates.get());
 #endif
@@ -366,12 +385,12 @@ int main(int argc, char ** argv) {
         if (params.conversation_mode && params.enable_chat_template) {
             if (!params.system_prompt.empty()) {
                 // format the system prompt (will use template default if empty)
-                chat_add_and_format("system", params.system_prompt, true);
+                chat_add_and_format("system", params.system_prompt);
             }
 
             if (!params.prompt.empty()) {
                 // format and append the user prompt
-                chat_add_and_format("user", params.prompt, true);
+                chat_add_and_format("user", params.prompt);
             } else {
                 waiting_for_first_input = true;
             }
@@ -905,9 +924,15 @@ int main(int argc, char ** argv) {
                     }
 
                     if (params.enable_chat_template) {
-                        chat_add_and_format("assistant", assistant_ss.str(), true);
-                        is_interacting = true;
-                        LOG("\n");
+                        auto format_res = chat_add_and_format("assistant", assistant_ss.str());
+                        if (format_res.tool_was_called) {
+                            auto format_res_tok = common_tokenize(ctx, format_res.formatted, false, true);
+                            embd_inp.insert(embd_inp.end(), format_res_tok.begin(), format_res_tok.end());
+
+                        } else {
+                            is_interacting = true;
+                            LOG("\n");
+                        }
                     }
                 }
             }
@@ -975,7 +1000,7 @@ int main(int argc, char ** argv) {
 
                     bool format_chat = params.conversation_mode && params.enable_chat_template;
                     std::string user_inp = format_chat
-                        ? chat_add_and_format("user", std::move(buffer))
+                        ? chat_add_and_format("user", std::move(buffer)).formatted
                         : std::move(buffer);
                     // TODO: one inconvenient of current chat template implementation is that we can't distinguish between user input and special tokens (prefix/postfix)
                     const auto line_pfx = common_tokenize(ctx, params.input_prefix, false, true);