server: realtime API: send complete after committed audio drains (#30994)

* server: realtime API: send complete after committed audio drains

* server: realtime API: complete a commit at its end row

Completion is triggered when the stream runs dry, which is a timing
proxy: audio appended while the tail drains leaks into the committed
item, a final commit before the drain merges both items, and back to
back commits collapse into one.

mtmd_audio_stream_flush() records the row count at the end of the
flushed audio (mtmd_audio_stream_get_n_rows_flushed), next to the rows
produced so far (mtmd_audio_stream_get_n_rows). The slot keeps a FIFO
of these end rows and completes each commit once every row up to it is
decoded.

---------

Co-authored-by: Pascal <admin@serveurperso.com>

tools/mtmd/mtmd.cpp

@@ -3096,7 +3096,8 @@ struct mtmd_audio_stream {
     std::vector<float> out;     // rows of the last step
     std::vector<float> tok_embd;
 
-    int64_t n_rows = 0; // rows produced so far
+    int64_t n_rows = 0;         // rows produced so far
+    int64_t n_rows_flushed = 0; // rows produced once the audio up to the last flush is consumed
     bool pending_flush = false;
     bool ended = false;
 
@@ -3209,6 +3210,8 @@ int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) {
     // align to a whole row, then add the silence the model expects after the audio
     const int64_t n_align = (s->samples_per_tok - s->n_samples % s->samples_per_tok) % s->samples_per_tok;
     mtmd_audio_stream_add_silence(s, n_align + (int64_t) s->n_right_pad * s->samples_per_tok);
+    // one row per samples_per_tok, the left padding included
+    s->n_rows_flushed = s->n_samples / s->samples_per_tok;
     s->pending_flush = true;
     s->ended = end;
     return 0;
@@ -3358,3 +3361,11 @@ int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) {
 const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * s) {
     return s->out.data();
 }
+
+int64_t mtmd_audio_stream_get_n_rows(const mtmd_audio_stream * s) {
+    return s->n_rows;
+}
+
+int64_t mtmd_audio_stream_get_n_rows_flushed(const mtmd_audio_stream * s) {
+    return s->n_rows_flushed;
+}

tools/mtmd/mtmd.h

@@ -412,6 +412,13 @@ MTMD_API int32_t mtmd_audio_stream_step(mtmd_audio_stream * stream, llama_token
 // rows of the last step, the reading size (in bytes) is n_rows * llama_model_n_embd_inp(model) * sizeof(float)
 MTMD_API const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * stream);
 
+// rows produced so far, the prefix rows included
+MTMD_API int64_t mtmd_audio_stream_get_n_rows(const mtmd_audio_stream * stream);
+
+// rows the stream produces up to the end of the last flush
+// once get_n_rows() reaches it, every row of the flushed audio has been produced
+MTMD_API int64_t mtmd_audio_stream_get_n_rows_flushed(const mtmd_audio_stream * stream);
+
 /////////////////////////////////////////
 // EXPERIMENTAL API for audio generation, subjected to breaking changes
 

tools/server/server-context.cpp

@@ -25,6 +25,7 @@
 #include <algorithm>
 #include <cstddef>
 #include <cinttypes>
+#include <deque>
 #include <exception>
 #include <memory>
 #include <filesystem>
@@ -399,6 +400,7 @@ struct server_slot {
     // realtime transcription, see SERVER_TASK_TYPE_TRANSCRIBE_RT
     mtmd::audio_stream_ptr audio_stream;
     bool audio_wait = false; // the stream needs more audio before it can produce rows
+    std::deque<int64_t> audio_commit_rows; // end row of each non-final commit not yet completed
 
     server_slot_stats stats;
 
@@ -456,6 +458,7 @@ struct server_slot {
 
         audio_stream.reset();
         audio_wait = false;
+        audio_commit_rows.clear();
         inp_embd.clear();
     }
 
@@ -3071,7 +3074,11 @@ private:
                             if (task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_APPEND) {
                                 mtmd_audio_stream_push(stream, task.params.control_audio.data(), task.params.control_audio.size());
                             } else {
-                                mtmd_audio_stream_flush(stream, task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_END);
+                                const int32_t ret = mtmd_audio_stream_flush(stream, task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_END);
+                                // a non-final commit completes its item once the rows of its audio are decoded
+                                if (ret == 0 && task.params.control_action == SERVER_CONTROL_ACTION_AUDIO_COMMIT) {
+                                    slot->audio_commit_rows.push_back(mtmd_audio_stream_get_n_rows_flushed(stream));
+                                }
                             }
                             slot->audio_wait = false;
                             res->success = true;
@@ -3402,6 +3409,22 @@ private:
             return true; // rows from a previous step, not decoded yet
         }
 
+        // every produced row is decoded and its delta queued, the commits ending there are complete
+        const int64_t n_rows = mtmd_audio_stream_get_n_rows(slot.audio_stream.get());
+        while (!slot.audio_commit_rows.empty() && slot.audio_commit_rows.front() <= n_rows) {
+            slot.audio_commit_rows.pop_front();
+            auto res = std::make_unique<server_task_result_cmpl_partial>();
+            res->id                   = slot.task->id;
+            res->index                = slot.task->index;
+            res->n_decoded            = slot.stats.n_gen;
+            res->n_prompt_tokens      = slot.task->n_tokens();
+            res->res_type             = slot.task->params.res_type;
+            res->oaicompat_model      = slot.task->params.oaicompat_model;
+            res->oaicompat_cmpl_id    = slot.task->params.oaicompat_cmpl_id;
+            res->audio_commit_drained = true;
+            queue_results.send(std::move(res));
+        }
+
         const int32_t n = mtmd_audio_stream_step(slot.audio_stream.get(), slot.sampled);
         if (n > 0) {
             const float * rows = mtmd_audio_stream_get_output_embd(slot.audio_stream.get());
@@ -6111,16 +6134,13 @@ void server_routes::init_routes() {
                     if (!post_control(SERVER_CONTROL_ACTION_AUDIO_COMMIT, {})) {
                         break;
                     }
-                    // the current item is the committed one, its transcript completes it and the next item opens
+                    // completed is sent once the rows of the committed audio are decoded (audio_commit_drained below)
                     std::lock_guard<std::mutex> lock(mutex_item);
                     send_event(json {
                         {"type",             "input_audio_buffer.committed"},
                         {"previous_item_id", item_prev.empty() ? json(nullptr) : json(item_prev)},
                         {"item_id",          item_id},
                     });
-                    send_completed();
-                    item_prev = item_id;
-                    item_id   = "item_" + random_string();
                 } else if (type == "input_audio_buffer.clear") {
                     send_event(json {{"type", "input_audio_buffer.cleared"}});
                 } else {
@@ -6144,7 +6164,13 @@ void server_routes::init_routes() {
                 break;
             }
             if (auto * partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get())) {
-                if (!partial->content.text.empty()) {
+                if (partial->audio_commit_drained) {
+                    // complete the committed item and open the next one
+                    std::lock_guard<std::mutex> lock(mutex_item);
+                    send_completed();
+                    item_prev = item_id;
+                    item_id   = "item_" + random_string();
+                } else if (!partial->content.text.empty()) {
                     std::lock_guard<std::mutex> lock(mutex_item);
                     item_text += partial->content.text;
                     send_event(json {

tools/server/server-task.h

@@ -470,6 +470,7 @@ struct server_task_result_cmpl_partial : server_task_result {
     bool is_progress = false;
     bool is_begin = false; // whether to send 200 status to HTTP client (begin of SSE stream)
                            // ref: https://github.com/ggml-org/llama.cpp/pull/23884
+    bool audio_commit_drained = false; // realtime: non-final commit has fully decoded
     completion_token_output prob_output;
     server_slot_stats stats;
     result_prompt_progress progress;