diff --git a/internal/service/model_chat.go b/internal/service/model_chat.go index 9319fb6105..d5edb89469 100644 --- a/internal/service/model_chat.go +++ b/internal/service/model_chat.go @@ -18,6 +18,7 @@ package service import ( "context" + "errors" "fmt" "go.uber.org/zap" @@ -26,6 +27,13 @@ import ( modelModule "ragflow/internal/entity/models" ) +// streamDoneSentinel is the OpenAI-style end-of-stream marker model drivers +// emit as their final sender call. It is a transport signal, not answer text. +const streamDoneSentinel = "[DONE]" + +// errStreamDone aborts the driver loop once the terminal sentinel arrives. +var errStreamDone = errors.New("chat stream done") + func (m *ModelProviderService) Chat(tenantID, modelID string, messages []modelModule.Message, config *modelModule.ChatConfig) (*modelModule.ChatResponse, error) { chatModel, err := m.GetChatModel(tenantID, modelID) if err != nil { @@ -51,6 +59,9 @@ func chatStreamWithContext(ctx context.Context, chatModel *modelModule.ChatModel if delta == nil { return nil } + if *delta == streamDoneSentinel { + return errStreamDone + } select { case ch <- *delta: return nil @@ -58,7 +69,7 @@ func chatStreamWithContext(ctx context.Context, chatModel *modelModule.ChatModel return ctx.Err() } }); err != nil { - if err == context.Canceled || err == context.DeadlineExceeded { + if errors.Is(err, errStreamDone) || err == context.Canceled || err == context.DeadlineExceeded { return } common.Warn("ChatStreamlyWithSender returned error", zap.Error(err)) diff --git a/internal/service/model_chat_test.go b/internal/service/model_chat_test.go new file mode 100644 index 0000000000..215299b86f --- /dev/null +++ b/internal/service/model_chat_test.go @@ -0,0 +1,65 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "context" + "strings" + "testing" + + "ragflow/internal/common" + modelModule "ragflow/internal/entity/models" +) + +// doneEmittingDriver streams the given deltas followed by the terminal +// "[DONE]" sentinel, mirroring real OpenAI-compatible drivers. +type doneEmittingDriver struct { + *modelModule.DummyModel + deltas []string +} + +func (d *doneEmittingDriver) ChatStreamlyWithSender(modelName string, messages []modelModule.Message, apiConfig *modelModule.APIConfig, modelConfig *modelModule.ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error { + for _, delta := range d.deltas { + if err := sender(&delta, nil); err != nil { + return err + } + } + done := streamDoneSentinel + return sender(&done, nil) +} + +func TestChatStreamWithContextStripsDoneSentinel(t *testing.T) { + modelName := "fake-model" + chatModel := &modelModule.ChatModel{ + ModelDriver: &doneEmittingDriver{ + DummyModel: modelModule.NewDummyModel(nil, modelModule.URLSuffix{}), + deltas: []string{"Hello", " world"}, + }, + ModelName: &modelName, + APIConfig: &modelModule.APIConfig{}, + } + + ch := chatStreamWithContext(context.Background(), chatModel, nil, &modelModule.ChatConfig{}) + var sb strings.Builder + for delta := range ch { + sb.WriteString(delta) + } + + if got := sb.String(); got != "Hello world" { + t.Fatalf("expected %q, got %q", "Hello world", got) + } +}