llama-server : make chat logprobs response OpenAI-compatible

tools/server/server-common.cpp

@@ -1157,7 +1157,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) {
@@ -1401,11 +1400,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

@@ -283,6 +283,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;
@@ -386,6 +394,7 @@ struct server_slot {
         }
         generated_tokens.clear();
         generated_token_probs.clear();
+        held_tokens.clear();
         json_schema = json();
 
         task_prev = std::move(task);
@@ -1844,6 +1853,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());
@@ -1867,15 +1878,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) {
@@ -1958,9 +1973,33 @@ 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) {
+            // generation ended with text still held: non-stream content includes it, while
+            // stream never sends it
+            release_held_tokens(slot, slot.task->params.stream ? 0 : std::string::npos);
+        }
+
         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) {
+            if (held.out.txt.size() > n_visible) {
+                held.out.txt.resize(n_visible);
+            }
+            n_visible -= held.out.txt.size();
+
+            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;
 
@@ -2043,7 +2082,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;
@@ -2063,7 +2102,9 @@ struct server_context_impl {
             res->tokens  = { 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;
@@ -3876,6 +3917,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) {
@@ -4002,6 +4044,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

@@ -282,12 +282,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)
@@ -483,6 +483,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({
@@ -499,6 +501,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
@@ -992,6 +1003,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;
@@ -1146,9 +1166,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

@@ -101,6 +101,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
@@ -111,6 +113,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;
@@ -300,6 +306,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;
@@ -337,6 +348,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;
@@ -366,6 +378,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;
@@ -405,6 +421,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;
 

tools/server/tests/unit/test_chat_completion.py

@@ -440,6 +440,110 @@ def test_logprobs_stream():
     assert aggregated_text == output_text
 
 
+@pytest.mark.parametrize("stream", [False, True])
+def test_logprobs_partial_stop_prefix(stream: bool):
+    # a token matching only a prefix of a stop string is held back, then emitted merged with
+    # the next token once generation diverges - its logprobs bytes must not be lost
+    global server
+    server.start()
+    body = {
+        "max_tokens": 8,
+        "temperature": 0.0,
+        "messages": [
+            {"role": "system", "content": "Book"},
+            {"role": "user", "content": "What is the best book"},
+        ],
+        "logprobs": True,
+        "top_logprobs": 1,
+    }
+    res = server.make_request("POST", "/chat/completions", data=body)
+    assert res.status_code == 200
+    entries = res.body["choices"][0]["logprobs"]["content"]
+    # stop on a mid-generation token followed by a char the model never emits, so that token
+    # only matches a prefix of the stop string and the next token diverges from it
+    prefix = next(e["token"] for e in entries[1:-1] if e["token"].strip())
+    body["stop"] = [prefix + "\x01"]
+
+    all_bytes = bytearray()
+    if stream:
+        content = ""
+        for data in server.make_stream_request("POST", "/chat/completions", data={**body, "stream": True}):
+            if not data["choices"]:
+                continue
+            choice = data["choices"][0]
+            content += choice["delta"].get("content") or ""
+            if choice.get("logprobs"):
+                for entry in choice["logprobs"]["content"]:
+                    all_bytes.extend(entry["bytes"])
+    else:
+        res = server.make_request("POST", "/chat/completions", data=body)
+        assert res.status_code == 200
+        content = res.body["choices"][0]["message"]["content"]
+        for entry in res.body["choices"][0]["logprobs"]["content"]:
+            all_bytes.extend(entry["bytes"])
+
+    assert prefix in content
+    assert bytes(all_bytes).decode("utf-8") == content
+
+
+@pytest.mark.parametrize("stream", [False, True])
+def test_logprobs_matched_stop(stream: bool):
+    # a stop string spanning two tokens: the first is held back as a partial match, then the
+    # second completes it - the logprobs bytes of both must stay out, just like the content
+    global server
+    server.start()
+    body = {
+        "max_tokens": 16,
+        "temperature": 0.0,
+        "messages": [
+            {"role": "system", "content": "Book"},
+            {"role": "user", "content": "What is the best book"},
+        ],
+        "logprobs": True,
+        "top_logprobs": 1,
+    }
+    res = server.make_request("POST", "/chat/completions", data=body)
+    assert res.status_code == 200
+    tokens = [e["token"] for e in res.body["choices"][0]["logprobs"]["content"]]
+    stop = next(
+        tokens[i] + tokens[i + 1] for i in range(1, len(tokens) - 2)
+        if tokens[i].strip() and tokens[i + 1].strip()
+        and tokens[i] + tokens[i + 1] not in "".join(tokens[:i + 1])
+    )
+    body["stop"] = [stop]
+
+    all_bytes = bytearray()
+    if stream:
+        content = ""
+        n_entries = 0
+        for data in server.make_stream_request("POST", "/chat/completions", data={**body, "stream": True, "stream_options": {"include_usage": True}}):
+            if data.get("usage"):
+                completion_tokens = data["usage"]["completion_tokens"]
+            if not data["choices"]:
+                continue
+            choice = data["choices"][0]
+            content += choice["delta"].get("content") or ""
+            if choice.get("finish_reason"):
+                finish_reason = choice["finish_reason"]
+            if choice.get("logprobs"):
+                for entry in choice["logprobs"]["content"]:
+                    all_bytes.extend(entry["bytes"])
+                    n_entries += 1
+        # the stop string's tokens keep their entries, just with no bytes
+        assert n_entries == completion_tokens
+    else:
+        res = server.make_request("POST", "/chat/completions", data=body)
+        assert res.status_code == 200
+        content = res.body["choices"][0]["message"]["content"]
+        finish_reason = res.body["choices"][0]["finish_reason"]
+        for entry in res.body["choices"][0]["logprobs"]["content"]:
+            all_bytes.extend(entry["bytes"])
+
+    assert finish_reason == "stop"
+    assert content and stop not in content
+    assert bytes(all_bytes).decode("utf-8") == content
+
+
 def test_logit_bias():
     global server
     server.start()