Repository: ollama/ollama
Pull request: Flush pending responses before returning (#6707)

I'll cross-port the server.cpp portion to main after this goes in.

Changed files:
filename: llama/runner/runner.go
patch: @@ -218,9 +218,22 @@ func incompleteUnicode(token string) bool {
 	return incomplete
 }
 
+func flushPending(seq *Sequence) bool {
+	for _, p := range seq.pendingResponses {
+		select {
+		case seq.responses <- p:
+		case <-seq.quit:
+			return false
+		}
+	}
+
+	return true
+}
+
 func (s *Server) removeSequence(seqIndex int, reason string) {
 	seq := s.seqs[seqIndex]
 
+	flushPending(seq)
 	seq.doneReason = reason
 	close(seq.responses)
 	close(seq.embedding)
@@ -348,24 +361,10 @@ func (s *Server) processBatch() {
 		seq.pendingResponses = append(seq.pendingResponses, piece)
 		sequence := strings.Join(seq.pendingResponses, "")
 
-		if incompleteUnicode(sequence) {
-			continue
-		}
-
 		if ok, stop := findStop(sequence, seq.stop); ok {
 			slog.Info("hit stop token", "stop", seq.stop)
 
-			truncated := truncateStop(seq.pendingResponses, stop)
-
-		truncatedLoop:
-			for _, p := range truncated {
-				select {
-				case seq.responses <- p:
-				case <-seq.quit:
-					break truncatedLoop
-				}
-			}
-
+			seq.pendingResponses = truncateStop(seq.pendingResponses, stop)
 			s.removeSequence(i, "stop")
 			continue
 		}
@@ -374,16 +373,14 @@ func (s *Server) processBatch() {
 			continue
 		}
 
-	pendingLoop:
-		for _, p := range seq.pendingResponses {
-			select {
-			case seq.responses <- p:
-			case <-seq.quit:
-				s.removeSequence(i, "connection")
-				break pendingLoop
-			}
+		if incompleteUnicode(sequence) {
+			continue
 		}
 
+		if !flushPending(seq) {
+			seq.pendingResponses = []string{}
+			s.removeSequence(i, "connection")
+		}
 		seq.pendingResponses = []string{}
 	}
 }
@@ -515,7 +512,8 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 
 	// Send the stop
 	if err := json.NewEncoder(w).Encode(&CompletionResponse{
-		Stop: true,
+		Stop:         true,
+		StoppedLimit: seq.doneReason == "limit",
 		Timings: Timings{
 			PromptN:     seq.numPromptTokens,
 			PromptMS:    float64(seq.startGenerationTime.Sub(seq.startProcessingTime).Milliseconds()),

filename: llm/ext_server/server.cpp
patch: @@ -913,7 +913,9 @@ struct llama_server_context
         slot.sampled = result.tok;
 
         // search stop word and delete it
-        slot.generated_text += token_str;
+        if (!llama_token_is_eog(model, result.tok))
+            slot.generated_text += token_str;
+
         slot.has_next_token = true;
 
         if (slot.ctx_sampling->params.use_penalty_prompt_tokens && result.tok != -1)
@@ -954,30 +956,36 @@ struct llama_server_context
         if (!incomplete)
         {
             size_t pos = std::min(slot.n_sent_text, slot.generated_text.size());
-            const std::string str_test = slot.generated_text.substr(pos);
-            bool is_stop_full = false;
-            size_t stop_pos = find_stopping_strings(str_test, token_str.size(), STOP_FULL, slot);
-            if (stop_pos != std::string::npos)
-            {
-                is_stop_full = true;
-                slot.generated_text.erase(
-                    slot.generated_text.begin() + pos + stop_pos,
-                    slot.generated_text.end());
-                pos = std::min(slot.n_sent_text, slot.generated_text.size());
-            }
-            else
-            {
-                is_stop_full = false;
-                stop_pos = find_stopping_strings(str_test, token_str.size(), STOP_PARTIAL, slot);
-            }
 
-            // check if there is any token to predict
-            if (stop_pos == std::string::npos || (!slot.has_next_token && !is_stop_full && stop_pos > 0))
-            {
-                // 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
+            if (!llama_token_is_eog(model, result.tok)) {
+                const std::string str_test = slot.generated_text.substr(pos);
+                bool is_stop_full = false;
+                size_t stop_pos = find_stopping_strings(str_test, token_str.size(), STOP_FULL, slot);
+                if (stop_pos != std::string::npos)
+                {
+                    is_stop_full = true;
+                    slot.generated_text.erase(
+                        slot.generated_text.begin() + pos + stop_pos,
+                        slot.generated_text.end());
+                    pos = std::min(slot.n_sent_text, slot.generated_text.size());
+                }
+                else
+                {
+                    is_stop_full = false;
+                    stop_pos = find_stopping_strings(str_test, token_str.size(), STOP_PARTIAL, slot);
+                }
+
+                // check if there is any token to predict
+                if (stop_pos == std::string::npos || (!slot.has_next_token && !is_stop_full && stop_pos > 0))
+                {
+                    // 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
+                }
+            } else {
+                    result.text_to_send = slot.generated_text.substr(pos, std::string::npos);
+                    slot.n_sent_text += result.text_to_send.size();
             }
 
             if (slot.params.stream)
@@ -1117,9 +1125,7 @@ struct llama_server_context
             {"multimodal", multimodal}
         };
 
-        if (!llama_token_is_eog(model, tkn.tok)) {
-            res.result_json["content"] = tkn.text_to_send;
-        }
+        res.result_json["content"] = tkn.text_to_send;
 
         if (slot.sparams.n_probs > 0)
         {