fix: loss ending done data

Signed-off-by: thxCode <thxcode0824@gmail.com>
gpustack · Jul 3, 2024 · 76c20fd · 76c20fd
1 parent b4406f7
commit 76c20fd
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/llama-box/main.cpp b/llama-box/main.cpp
@@ -3131,6 +3131,12 @@ int main(int argc, char **argv) {
                     continue;
                 }
 
+                const std::string done = "data: [DONE] \n\n";
+                if (!sink.write(done.c_str(), done.size())) {
+                    sink.done();
+                    return false;
+                }
+
                 sink.done_with_trailer({{"X-Response-Tokens-Per-Second",
                                          std::to_string(json_value(result.data.at("timings"),
                                                                    "predicted_per_second", tps))}});