Fix stop words (#22)

* Updated README
* Updated chat to reflect llama.cpp approach (stop conditions)
* Update pkg/llamacpp/chat.go

Author: djthorpe

.gitignore

@@ -26,10 +26,11 @@ go.work.sum
 
 # Editor/IDE
 .vscode/
+.claude
 
 # Build artifacts
 build/
 
 # Other
-/gollama
+/go-llama
 .DS*

pkg/llamacpp/chat.go

@@ -281,25 +281,23 @@ var defaultStopSequences = []string{
 
 type stopMarkerFilter struct {
 	stops   []string
-	maxLen  int
 	buffer  string
 	stopped bool
 }
 
 func newStopMarkerFilter(stops []string) *stopMarkerFilter {
-	maxLen := 0
-	for _, s := range stops {
-		if len(s) > maxLen {
-			maxLen = len(s)
-		}
-	}
-	return &stopMarkerFilter{stops: stops, maxLen: maxLen}
+	return &stopMarkerFilter{stops: stops}
 }
 
 func (f *stopMarkerFilter) Stopped() bool {
 	return f.stopped
 }
 
+// Process handles incoming tokens using llama.cpp's two-phase approach:
+// 1. Check for full stop sequences
+// 2. If no full stop, check if text ends with a prefix of any stop word (partial match)
+//
+// Returns the text safe to send and whether generation should stop.
 func (f *stopMarkerFilter) Process(text string) (string, bool) {
 	if f.stopped {
 		return "", true
@@ -309,25 +307,28 @@ func (f *stopMarkerFilter) Process(text string) (string, bool) {
 	}
 
 	combined := f.buffer + text
-	idx, found := indexAnyStop(combined, f.stops)
-	if found {
+
+	// Phase 1: Check for full stop sequence
+	if idx, found := indexAnyStop(combined, f.stops); found {
 		f.stopped = true
+		f.buffer = ""
 		return combined[:idx], true
 	}
 
-	if f.maxLen <= 1 {
-		f.buffer = ""
-		return combined, false
-	}
-	keep := f.maxLen - 1
-	if len(combined) <= keep {
-		f.buffer = combined
-		return "", false
+	// Phase 2: Check if text ends with a partial stop sequence
+	// If so, withhold the partial match in the buffer
+	if partialPos := findPartialStop(combined, f.stops); partialPos != -1 {
+		// Text ends with a prefix of a stop word - withhold from partialPos onward
+		f.buffer = combined[partialPos:]
+		if partialPos == 0 {
+			return "", false
+		}
+		return combined[:partialPos], false
 	}
-	cut := len(combined) - keep
-	out := combined[:cut]
-	f.buffer = combined[cut:]
-	return out, false
+
+	// No full or partial match - safe to send everything
+	f.buffer = ""
+	return combined, false
 }
 
 func (f *stopMarkerFilter) Flush() string {
@@ -339,6 +340,49 @@ func (f *stopMarkerFilter) Flush() string {
 	return out
 }
 
+// findPartialStop checks if s ends with a prefix of any stop word.
+// Returns the position where the partial match starts, or -1 if no partial match.
+// This mirrors llama.cpp's string_find_partial_stop function.
+func findPartialStop(s string, stops []string) int {
+	if len(s) == 0 {
+		return -1
+	}
+
+	lastChar := s[len(s)-1]
+
+	for _, stop := range stops {
+		if len(stop) == 0 {
+			continue
+		}
+
+		// Check each possible prefix of the stop word, starting from longest.
+	maxPrefixLen := 0
+
+	for _, stop := range stops {
+		if len(stop) == 0 {
+			continue
+		}
+
+		// Check each possible prefix of the stop word, starting from longest.
+		// We look for prefixes that end with the last character of s.
+		for prefixLen := len(stop); prefixLen >= 1; prefixLen-- {
+			if stop[prefixLen-1] != lastChar {
+				continue
+			}
+
+			prefix := stop[:prefixLen]
+			if strings.HasSuffix(s, prefix) && prefixLen > maxPrefixLen {
+				maxPrefixLen = prefixLen
+			}
+		}
+	}
+
+	if maxPrefixLen > 0 {
+		return len(s) - maxPrefixLen
+	}
+	return -1
+}
+
 func indexAnyStop(s string, stops []string) (int, bool) {
 	idx := -1
 	for _, stop := range stops {

pkg/llamacpp/chat_test.go

@@ -150,3 +150,197 @@ func TestChatLogsOutput(t *testing.T) {
 	t.Logf("chat response: %q", resp.Message.Content)
 	t.Logf("usage: input=%d output=%d total=%d", resp.Usage.InputTokens, resp.Usage.OutputTokens, resp.Usage.TotalTokens())
 }
+
+func TestFindPartialStop(t *testing.T) {
+	tests := []struct {
+		name     string
+		text     string
+		stops    []string
+		expected int
+	}{
+		{
+			name:     "no partial match",
+			text:     "hello world",
+			stops:    []string{"</s>", "<|end|>"},
+			expected: -1,
+		},
+		{
+			name:     "full stop not detected as partial",
+			text:     "hello</s>",
+			stops:    []string{"</s>"},
+			expected: 5, // partial match starts at position 5
+		},
+		{
+			name:     "partial match single char",
+			text:     "hello<",
+			stops:    []string{"</s>"},
+			expected: 5,
+		},
+		{
+			name:     "partial match two chars",
+			text:     "hello</",
+			stops:    []string{"</s>"},
+			expected: 5,
+		},
+		{
+			name:     "partial match three chars",
+			text:     "hello</s",
+			stops:    []string{"</s>"},
+			expected: 5,
+		},
+		{
+			name:     "partial match with pipe",
+			text:     "test<|",
+			stops:    []string{"<|end|>"},
+			expected: 4,
+		},
+		{
+			name:     "partial match longer prefix",
+			text:     "test<|end",
+			stops:    []string{"<|end|>"},
+			expected: 4,
+		},
+		{
+			name:     "partial match almost complete",
+			text:     "test<|end|",
+			stops:    []string{"<|end|>"},
+			expected: 4,
+		},
+		{
+			name:     "multiple stops first matches",
+			text:     "hello<",
+			stops:    []string{"</s>", "<|end|>"},
+			expected: 5,
+		},
+		{
+			name:     "empty text",
+			text:     "",
+			stops:    []string{"</s>"},
+			expected: -1,
+		},
+		{
+			name:     "empty stops",
+			text:     "hello",
+			stops:    []string{},
+			expected: -1,
+		},
+		{
+			name:     "text is just partial",
+			text:     "</",
+			stops:    []string{"</s>"},
+			expected: 0,
+		},
+		{
+			name:     "no match when char not in stop",
+			text:     "hellox",
+			stops:    []string{"</s>"},
+			expected: -1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := findPartialStop(tt.text, tt.stops)
+			if result != tt.expected {
+				t.Errorf("findPartialStop(%q, %v) = %d, want %d", tt.text, tt.stops, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestStopMarkerFilter(t *testing.T) {
+	tests := []struct {
+		name           string
+		stops          []string
+		tokens         []string
+		expectedChunks []string
+		expectedStop   bool
+	}{
+		{
+			name:           "no stop sequences",
+			stops:          []string{"</s>"},
+			tokens:         []string{"hello", " ", "world"},
+			expectedChunks: []string{"hello", " ", "world"},
+			expectedStop:   false,
+		},
+		{
+			name:           "full stop in single token",
+			stops:          []string{"</s>"},
+			tokens:         []string{"hello", "</s>", "extra"},
+			expectedChunks: []string{"hello", ""},
+			expectedStop:   true,
+		},
+		{
+			name:           "stop split across tokens",
+			stops:          []string{"</s>"},
+			tokens:         []string{"hello<", "/s>more"},
+			expectedChunks: []string{"hello", ""},
+			expectedStop:   true,
+		},
+		{
+			name:           "partial withholding then release",
+			stops:          []string{"</s>"},
+			tokens:         []string{"hello<", "notastop"},
+			expectedChunks: []string{"hello", "<notastop"},
+			expectedStop:   false,
+		},
+		{
+			name:           "partial withholding with longer sequence",
+			stops:          []string{"<|end|>"},
+			tokens:         []string{"test<|", "end|>done"},
+			expectedChunks: []string{"test", ""},
+			expectedStop:   true,
+		},
+		{
+			name:           "buffer accumulates partial",
+			stops:          []string{"</s>"},
+			tokens:         []string{"a<", "/", "s>"},
+			expectedChunks: []string{"a", "", ""},
+			expectedStop:   true,
+		},
+		{
+			name:           "stop at very beginning",
+			stops:          []string{"</s>"},
+			tokens:         []string{"</s>rest"},
+			expectedChunks: []string{""},
+			expectedStop:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			filter := newStopMarkerFilter(tt.stops)
+			var chunks []string
+
+			for _, token := range tt.tokens {
+				chunk, stopped := filter.Process(token)
+				chunks = append(chunks, chunk)
+				if stopped {
+					break
+				}
+			}
+
+			// If not stopped, flush remaining buffer
+			if !filter.Stopped() {
+				if tail := filter.Flush(); tail != "" {
+					chunks[len(chunks)-1] += tail
+				}
+			}
+
+			if len(chunks) != len(tt.expectedChunks) {
+				t.Errorf("got %d chunks %v, want %d chunks %v", len(chunks), chunks, len(tt.expectedChunks), tt.expectedChunks)
+				return
+			}
+
+			for i, chunk := range chunks {
+				if chunk != tt.expectedChunks[i] {
+					t.Errorf("chunk[%d] = %q, want %q", i, chunk, tt.expectedChunks[i])
+				}
+			}
+
+			if filter.Stopped() != tt.expectedStop {
+				t.Errorf("Stopped() = %v, want %v", filter.Stopped(), tt.expectedStop)
+			}
+		})
+	}
+}

pkg/llamacpp/completion.go

@@ -34,13 +34,21 @@ func (l *Llama) Complete(ctx context.Context, req schema.CompletionRequest, onCh
 
 		opts := buildCompletionOptions(ctx, req)
 		var callbackErr error
+		var stopFilter *stopMarkerFilter
 		if onChunk != nil {
+			stopFilter = newStopMarkerFilter(opts.StopWords)
 			opts.OnToken = func(token string) bool {
 				if callbackErr != nil {
 					return false
 				}
-				if err := onChunk(schema.CompletionChunk{Text: token}); err != nil {
-					callbackErr = err
+				filtered, stopped := stopFilter.Process(token)
+				if filtered != "" {
+					if err := onChunk(schema.CompletionChunk{Text: filtered}); err != nil {
+						callbackErr = err
+						return false
+					}
+				}
+				if stopped {
 					return false
 				}
 				return true
@@ -55,6 +63,18 @@ func (l *Llama) Complete(ctx context.Context, req schema.CompletionRequest, onCh
 			return callbackErr
 		}
 
+		// Flush any buffered content if we didn't hit a stop
+		if onChunk != nil && stopFilter != nil && !stopFilter.Stopped() {
+			if tail := stopFilter.Flush(); tail != "" {
+				if err := onChunk(schema.CompletionChunk{Text: tail}); err != nil {
+					return err
+				}
+			}
+		}
+
+		// Trim stop sequences from final text
+		text, _ = trimAtStop(text, opts.StopWords)
+
 		usage, err := completionUsage(task.Model(), req.Prompt, text)
 		if err != nil {
 			return err

pkg/llamacpp/httphandler/model.go

@@ -143,9 +143,8 @@ func modelLoadUnload(w http.ResponseWriter, r *http.Request, llamaInstance *llam
 	var req schema.LoadModelRequest
 	if err := httprequest.Read(r, &req); err != nil {
 		return httpresponse.Error(w, httpresponse.ErrBadRequest.With(err.Error()))
-	} else {
-		req.Name = r.PathValue("id")
 	}
+	req.Name = r.PathValue("id")
 
 	// Check if this is an unload request
 	isUnload := req.Load != nil && !*req.Load

pkg/llamacpp/httphandler/model_test.go

@@ -235,7 +235,8 @@ func TestModelLoad_NonExistentModel(t *testing.T) {
 	router := http.NewServeMux()
 	RegisterModelHandlers(router, "/api", llama, noopMiddleware())
 
-	req := httptest.NewRequest(http.MethodPost, "/api/model/nonexistent", nil)
+	req := httptest.NewRequest(http.MethodPost, "/api/model/nonexistent", strings.NewReader("{}"))
+	req.Header.Set("Content-Type", "application/json")
 	rw := httptest.NewRecorder()
 
 	router.ServeHTTP(rw, req)