llama-server : make chat logprobs response OpenAI-compatible

tools/server/server-common.cpp

@@ -1278,7 +1278,6 @@ json oaicompat_chat_params_parse(
 
     auto tools = json_value(body, "tools", json());
     auto has_tools = tools.is_array() && !tools.empty();
-    auto stream = json_value(body, "stream", false);
     auto tool_choice = json_value(body, "tool_choice", std::string("auto"));
 
     if (!opt.use_jinja) {
@@ -1454,11 +1453,10 @@ json oaicompat_chat_params_parse(
     }
 
     // Handle "logprobs" field
-    // TODO: The response format of this option is not yet OAI-compatible, but seems like no one really using it; We may need to fix it in the future
+    // Response is OpenAI-modern shape (choices[].logprobs.content[]); each entry also carries the
+    // integer token id. With tools, the streaming path emits per-token logprobs for all tokens
+    // (including tool-call tokens), matching the non-streaming path.
     if (json_value(body, "logprobs", false)) {
-        if (has_tools && stream) {
-            throw std::invalid_argument("logprobs is not supported with tools + stream");
-        }
         llama_params["n_probs"] = json_value(body, "top_logprobs", 20);
     } else if (body.contains("top_logprobs") && !body.at("top_logprobs").is_null()) {
         throw std::invalid_argument("top_logprobs requires logprobs to be set to true");

tools/server/server-context.cpp

@@ -292,6 +292,14 @@ struct server_slot {
 
     std::vector<completion_token_output> generated_token_probs;
 
+    // tokens whose text is not sent yet (incomplete UTF-8 or partial stop word) - their
+    // logprobs wait with it, since the text may still turn out to be a stop word
+    struct held_token {
+        completion_token_output out;
+        uint64_t                n_gen; // n_decoded of its partial response
+    };
+    std::vector<held_token> held_tokens;
+
     bool has_next_token = true;
     bool has_new_line   = false;
     bool truncated      = false;
@@ -395,6 +403,7 @@ struct server_slot {
         }
         generated_tokens.clear();
         generated_token_probs.clear();
+        held_tokens.clear();
         json_schema = json();
 
         task_prev = std::move(task);
@@ -1990,6 +1999,8 @@ struct server_context_impl {
         // check if there is incomplete UTF-8 character at the end
         bool incomplete = validate_utf8(slot.generated.text) < slot.generated.text.size();
 
+        bool resolved = false;
+
         // search stop word and delete it
         if (!incomplete) {
             size_t pos = std::min(slot.n_sent_text, slot.generated.text.size());
@@ -2011,15 +2022,19 @@ struct server_context_impl {
                 // no send the stop word in the response
                 result.text_to_send = slot.generated.text.substr(pos, std::string::npos);
                 slot.n_sent_text += result.text_to_send.size();
-                // add the token to slot queue and cache
+                resolved = true;
             } else {
                 result.text_to_send = "";
             }
+        } else {
+            result.text_to_send = "";
+        }
 
-            slot.add_token(result);
-            if (slot.task->params.stream) {
-                send_partial_response(slot, result, false);
-            }
+        // every sampled token gets its own logprobs entry (and, if streaming, its own
+        // partial chunk), released once its text is either sent or cut as a stop word
+        slot.held_tokens.push_back({ result, slot.stats.n_gen });
+        if (resolved) {
+            release_held_tokens(slot, result.text_to_send.size());
         }
 
         if (incomplete) {
@@ -2102,9 +2117,45 @@ struct server_context_impl {
 
         SLT_DBG(slot, "n_gen = %d, n_remaining = %d, next token: %5d '%s'\n", (int) slot.stats.n_gen, slot.n_remaining(), result.tok, token_str.c_str());
 
+        if (!slot.has_next_token && !slot.held_tokens.empty()) {
+            // generation ended with text still held, which non-stream content includes, so the
+            // last chunk must send it too
+            const size_t pos = std::min(slot.n_sent_text, slot.generated.text.size());
+            auto & last = slot.held_tokens.back().out;
+            last.text_to_send = slot.generated.text.substr(pos);
+            slot.n_sent_text += last.text_to_send.size();
+            release_held_tokens(slot, last.text_to_send.size());
+        }
+
         return slot.has_next_token; // continue
     }
 
+    // the held tokens' pieces add up to the text not sent yet, of which only the first
+    // n_visible bytes reach content (the rest is a matched stop word), so cut each
+    // token's logprobs text to its share of those
+    void release_held_tokens(server_slot & slot, size_t n_visible) {
+        for (auto & held : slot.held_tokens) {
+            const bool hidden = n_visible == 0 && !held.out.txt.empty();
+
+            if (held.out.txt.size() > n_visible) {
+                held.out.txt.resize(n_visible);
+            }
+            n_visible -= held.out.txt.size();
+
+            if (hidden) {
+                // its chunk still goes out, as it can carry the content released before the
+                // stop word - empty probs keep the logprobs entry out of it
+                held.out.probs.clear();
+            } else {
+                slot.add_token(held.out);
+            }
+            if (slot.task->params.stream) {
+                send_partial_response(slot, held.out, false, false, held.n_gen);
+            }
+        }
+        slot.held_tokens.clear();
+    }
+
     void populate_token_probs(const server_slot & slot, completion_token_output & result, bool post_sampling, bool special, int idx) const {
         const size_t n_probs_request = slot.task->params.sampling.n_probs;
 
@@ -2187,7 +2238,7 @@ struct server_context_impl {
         queue_results.send(std::move(res));
     }
 
-    void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false) {
+    void send_partial_response(server_slot & slot, const completion_token_output & tkn, bool is_progress, bool is_begin = false, uint64_t n_gen = 0) {
         auto res = std::make_unique<server_task_result_cmpl_partial>();
 
         res->id    = slot.task->id;
@@ -2207,7 +2258,9 @@ struct server_context_impl {
             res->tokens.assign(1, tkn.tok);
         }
 
-        res->n_decoded             = slot.stats.n_gen;
+        // a held token's chunk goes out late, but must keep its own count (e.g. n_decoded == 1
+        // opens the stream)
+        res->n_decoded             = n_gen > 0 ? n_gen : slot.stats.n_gen;
         res->n_prompt_tokens       = slot.task->n_tokens();
         res->n_prompt_tokens_cache = slot.stats.n_prompt_cached;
         res->post_sampling_probs   = slot.task->params.post_sampling_probs;
@@ -2273,19 +2326,11 @@ struct server_context_impl {
         res->oaicompat_cmpl_id = slot.task->params.oaicompat_cmpl_id;
 
         // populate res.probs_output
+        // the tokens of a matched stop word are already left out - see release_held_tokens()
         if (slot.task->params.sampling.n_probs > 0) {
-            if (!slot.task->params.stream && slot.stop == STOP_TYPE_WORD) {
-                const llama_tokens stop_word_toks = common_tokenize(ctx_tgt, slot.stopping_word, false);
-
-                size_t safe_offset = std::min(slot.generated_token_probs.size(), stop_word_toks.size());
-                res->probs_output = std::vector<completion_token_output>(
-                        slot.generated_token_probs.begin(),
-                        slot.generated_token_probs.end() - safe_offset);
-            } else {
-                res->probs_output = std::vector<completion_token_output>(
-                        slot.generated_token_probs.begin(),
-                        slot.generated_token_probs.end());
-            }
+            res->probs_output = std::vector<completion_token_output>(
+                    slot.generated_token_probs.begin(),
+                    slot.generated_token_probs.end());
         }
 
         res->generation_params = slot.task->params; // copy the parameters
@@ -4388,6 +4433,7 @@ struct server_context_impl {
             completion_token_output result;
             result.tok          = id;
             result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok));
+            result.txt          = result.text_to_send;
             result.prob         = 1.0f; // TODO: set it here instead of doing inside populate_token_probs
 
             if (slot.task->params.sampling.n_probs > 0) {
@@ -4527,6 +4573,7 @@ struct server_context_impl {
 
                 result.tok          = ids[i];
                 result.text_to_send = common_token_to_piece(slot.ctx_tgt, result.tok, accept_special_token(slot, result.tok));
+                result.txt          = result.text_to_send;
                 result.prob         = 1.0f; // set later
 
                 // TODO: set result.probs

tools/server/server-task.cpp

@@ -274,12 +274,12 @@ json completion_token_output::to_json(bool post_sampling_probs) const {
 json completion_token_output::probs_vector_to_json(const std::vector<completion_token_output> & probs, bool post_sampling_probs) {
     json out = json::array();
     for (const auto & p : probs) {
-        std::string txt(p.text_to_send);
+        std::string txt(p.txt);
         txt.resize(validate_utf8(txt));
         out.push_back(json {
             {"id",           p.tok},
             {"token",        txt},
-            {"bytes",        str_to_bytes(p.text_to_send)},
+            {"bytes",        str_to_bytes(p.txt)},
             {
                 post_sampling_probs ? "prob" : "logprob",
                 post_sampling_probs ? p.prob : logarithm(p.prob)
@@ -475,6 +475,8 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat_stream() {
             {"object", "chat.completion.chunk"},
         });
     }
+    // index of the last diff chunk, if any, before the finish-reason chunk below is pushed
+    size_t last_diff_idx = deltas.empty() ? SIZE_MAX : deltas.size() - 1;
 
     deltas.push_back({
         {"choices", json::array({
@@ -491,6 +493,15 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat_stream() {
         {"object",             "chat.completion.chunk"},
     });
 
+    // attach any leftover pending_probs to the last diff chunk, or the finish-reason
+    // chunk if there was no diff at all
+    if (!probs_to_emit.empty()) {
+        size_t target_idx = (last_diff_idx != SIZE_MAX) ? last_diff_idx : deltas.size() - 1;
+        deltas[target_idx].at("choices").at(0)["logprobs"] = json {
+            {"content", completion_token_output::probs_vector_to_json(probs_to_emit, post_sampling_probs)},
+        };
+    }
+
     if (include_usage) {
         // OpenAI API spec for chat.completion.chunks specifies an empty `choices` array for the last chunk when including usage
         // https://platform.openai.com/docs/api-reference/chat_streaming/streaming#chat_streaming/streaming-choices
@@ -984,6 +995,15 @@ void server_task_result_cmpl_partial::update(task_result_state & state) {
     }
     state.update_chat_msg(content, true, oaicompat_msg_diffs);
 
+    // queue this token's logprobs, flush once a chunk has visible delta content to attach them to
+    if (prob_output.probs.size() > 0) {
+        state.pending_probs.push_back(prob_output);
+    }
+    if (is_progress || n_decoded == 1 || !oaicompat_msg_diffs.empty()) {
+        probs_to_emit = std::move(state.pending_probs);
+        state.pending_probs.clear();
+    }
+
     // Copy current state for use in to_json_*() (reflects state BEFORE this chunk)
     thinking_block_started = state.thinking_block_started;
     text_block_started     = state.text_block_started;
@@ -1138,9 +1158,9 @@ json server_task_result_cmpl_partial::to_json_oaicompat_chat() {
         auto & last_json = deltas[deltas.size() - 1];
         GGML_ASSERT(last_json.at("choices").size() >= 1);
 
-        if (prob_output.probs.size() > 0) {
+        if (!probs_to_emit.empty()) {
             last_json.at("choices").at(0)["logprobs"] = json {
-                {"content", completion_token_output::probs_vector_to_json({prob_output}, post_sampling_probs)},
+                {"content", completion_token_output::probs_vector_to_json(probs_to_emit, post_sampling_probs)},
             };
         }
 

tools/server/server-task.h

@@ -103,6 +103,8 @@ struct task_params {
     json to_json(bool only_metrics = false) const;
 };
 
+struct completion_token_output; // defined below, used by task_result_state
+
 // struct for tracking the state of a task (e.g., for streaming)
 struct task_result_state {
     // tracking diffs for partial tool calls
@@ -112,6 +114,10 @@ struct task_result_state {
     std::vector<std::string> generated_tool_call_ids;
     std::unordered_set<size_t> sent_tool_call_names;
 
+    // logprobs for tokens whose text hasn't produced a chat-msg diff yet (still
+    // being buffered by the parser); flushed onto the next chunk that does
+    std::vector<completion_token_output> pending_probs;
+
     // for OpenAI Responses and Anthropic streaming API:
     // track output item / content block state across chunks
     bool thinking_block_started = false;
@@ -343,6 +349,11 @@ struct completion_token_output {
     llama_token tok;
     float prob;
     std::string text_to_send;
+    // this token's own decoded piece, for logprobs attribution - independent of
+    // text_to_send, which can be withheld (incomplete UTF-8, partial stop word) or
+    // merged with earlier tokens; cut to the part that reached content if a stop word
+    // matches
+    std::string txt;
     struct prob_info {
         llama_token tok;
         std::string txt;
@@ -380,6 +391,7 @@ struct server_task_result_cmpl_final : server_task_result {
 
     bool post_sampling_probs;
     std::vector<completion_token_output> probs_output;
+    std::vector<completion_token_output> probs_to_emit; // leftover pending_probs flushed by update()
     std::vector<std::string>  response_fields;
 
     task_params generation_params;
@@ -409,6 +421,10 @@ struct server_task_result_cmpl_final : server_task_result {
         is_updated = true;
         oaicompat_msg = state.update_chat_msg(content, false, oaicompat_msg_diffs);
 
+        // last update for the request: drain pending_probs unconditionally
+        probs_to_emit = std::move(state.pending_probs);
+        state.pending_probs.clear();
+
         oai_resp_id = state.oai_resp_id;
         oai_resp_reasoning_id = state.oai_resp_reasoning_id;
         oai_resp_message_id = state.oai_resp_message_id;
@@ -448,6 +464,7 @@ struct server_task_result_cmpl_partial : server_task_result {
     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
     completion_token_output prob_output;
+    std::vector<completion_token_output> probs_to_emit; // pending_probs flushed by update() for this chunk
     server_slot_stats stats;
     result_prompt_progress progress;