move states to server_response_reader

Author: ngxson

tools/server/server-context.cpp

@@ -2571,12 +2571,12 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
     auto completion_id = gen_chatcmplid();
     auto & rd = res->rd;
 
-    // tracking generation state and partial tool calls
-    std::vector<task_result_state> states;
-
     try {
         std::vector<server_task> tasks;
 
+        // tracking generation state and partial tool calls
+        std::vector<task_result_state> states;
+
         const auto & prompt = data.at("prompt");
         // TODO: this log can become very long, put it behind a flag or think about a more compact format
         //SRV_DBG("Prompt: %s\n", prompt.is_string() ? prompt.get<std::string>().c_str() : prompt.dump(2).c_str());
@@ -2615,6 +2615,7 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
             tasks.push_back(std::move(task));
         }
 
+        rd.set_states(std::move(states));
         rd.post_tasks(std::move(tasks));
     } catch (const std::exception & e) {
         res->error(format_error_response(e.what(), ERROR_TYPE_INVALID_REQUEST));
@@ -2635,7 +2636,6 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
             json arr = json::array();
             for (auto & res : all_results.results) {
                 GGML_ASSERT(dynamic_cast<server_task_result_cmpl_final*>(res.get()) != nullptr);
-                res->update(states[res->get_index()]); // update generation state
                 arr.push_back(res->to_json());
             }
             // if single request, return single object instead of array
@@ -2656,7 +2656,6 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
                 dynamic_cast<server_task_result_cmpl_partial*>(first_result.get()) != nullptr
                 || dynamic_cast<server_task_result_cmpl_final*>(first_result.get()) != nullptr
             );
-            first_result->update(states[first_result->get_index()]); // update generation state
         }
 
         // next responses are streamed
@@ -2669,7 +2668,7 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
         }
         res->status = 200;
         res->content_type = "text/event-stream";
-        res->next = [res_this = res.get(), res_type, &should_stop, states = std::move(states)](std::string & output) mutable -> bool {
+        res->next = [res_this = res.get(), res_type, &should_stop](std::string & output) mutable -> bool {
             static auto format_error = [](task_response_type res_type, const json & res_json) {
                 if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) {
                     return format_anthropic_sse({
@@ -2728,7 +2727,6 @@ static std::unique_ptr<server_res_generator> handle_completions_impl(
                         dynamic_cast<server_task_result_cmpl_partial*>(result.get()) != nullptr
                         || dynamic_cast<server_task_result_cmpl_final*>(result.get()) != nullptr
                     );
-                    result->update(states[result->get_index()]); // update generation state
                     json res_json = result->to_json();
                     if (res_type == TASK_RESPONSE_TYPE_ANTHROPIC) {
                         output = format_anthropic_sse(res_json);

tools/server/server-queue.cpp

@@ -271,6 +271,10 @@ void server_response::terminate() {
 // server_response_reader
 //
 
+void server_response_reader::set_states(std::vector<task_result_state> && states) {
+    this->states = std::move(states);
+}
+
 void server_response_reader::post_tasks(std::vector<server_task> && tasks) {
     id_tasks = server_task::get_list_id(tasks);
     queue_results.add_waiting_tasks(tasks);
@@ -298,6 +302,12 @@ server_task_result_ptr server_response_reader::next(const std::function<bool()>
                 SRV_DBG("%s", "received error result, stopping further processing\n");
                 return result;
             }
+            if (!states.empty()) {
+                // update the generation state if needed
+                auto idx = result->get_index();
+                GGML_ASSERT(idx >= 0 && idx < states.size());
+                result->update(states[idx]);
+            }
             if (result->is_stop()) {
                 received_count++;
             }

tools/server/server-queue.h

@@ -7,6 +7,8 @@
 #include <mutex>
 #include <unordered_set>
 
+// struct for managing server tasks
+// in most cases, use server_response_reader to post new tasks and retrieve results
 struct server_queue {
 private:
     int id = 0;
@@ -67,6 +69,8 @@ struct server_queue {
     void cleanup_pending_task(int id_target);
 };
 
+// struct for managing server responses
+// in most cases, use server_response_reader to retrieve results
 struct server_response {
 private:
     bool running = true;
@@ -120,13 +124,18 @@ struct server_response_reader {
     bool cancelled = false;
     int polling_interval_seconds;
 
+    // tracking generation state and partial tool calls
+    // only used by streaming completions
+    std::vector<task_result_state> states;
+
     // should_stop function will be called each polling_interval_seconds
     server_response_reader(std::pair<server_queue &, server_response &> server_queues, int polling_interval_seconds)
         : queue_tasks(server_queues.first), queue_results(server_queues.second), polling_interval_seconds(polling_interval_seconds) {}
     ~server_response_reader() {
         stop();
     }
 
+    void set_states(std::vector<task_result_state> && states);
     void post_tasks(std::vector<server_task> && tasks);
     bool has_next() const;