diff --git a/internal/router/router.go b/internal/router/router.go index dbe8548335..08f722d08a 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -17,9 +17,10 @@ package router import ( - "github.com/gin-gonic/gin" "ragflow/internal/common" "ragflow/internal/handler" + + "github.com/gin-gonic/gin" ) type Router struct { diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index a5b5fc5aa5..5a65067a05 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -1609,7 +1609,15 @@ func (s *ChatSessionService) ChatCompletions( } sendOrCancel(fmt.Sprintf("data:%s\n\n", sseMarshalChunk(sanitizeJSONFloats(ans).(map[string]interface{}), chatID))) } else { - ans := s.structureAnswer(session, "", messageID, sessionID, reference) + ans := s.structureAnswer(session, result.Answer, messageID, sessionID, reference) + if result.Reference != nil { + ans["reference"] = result.Reference + } + ans["audio_binary"] = result.AudioBinary + ans["prompt"] = result.Prompt + if result.CreatedAt != 0 { + ans["created_at"] = result.CreatedAt + } ans["final"] = true if chatID != "" { ans["chat_id"] = chatID @@ -1779,7 +1787,6 @@ func (s *ChatSessionService) buildDefaultCompletionDialog(tenantID string) *enti } } -// createSessionForCompletion mirrors Python _create_session_for_completion. func (s *ChatSessionService) createSessionForCompletion(chatID string, dialog *entity.Chat, userID string) (*entity.ChatSession, error) { newID := utility.GenerateUUID() name := "New session" diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index 1fad7a3e6c..c7060461d3 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -903,6 +903,123 @@ func TestChatCompletionsPassesRequestUserIDToPipeline(t *testing.T) { } } +func TestChatCompletionsStreamFinalCarriesDecoratedReference(t *testing.T) { + store := newFakeSessionStore() + store.sessions["session-1"] = &entity.ChatSession{ + ID: "session-1", + DialogID: "dialog-1", + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), + Reference: json.RawMessage(`[]`), + } + store.dialogs["dialog-1"] = &entity.Chat{ + ID: "dialog-1", + TenantID: "tenant-owner", + LLMID: "chat@factory", + LLMSetting: entity.JSONMap{}, + PromptConfig: entity.JSONMap{ + "parameters": []interface{}{}, + }, + } + store.dialogExists["tenant-owner|dialog-1"] = true + + finalReference := map[string]interface{}{ + "chunks": []map[string]interface{}{ + { + "id": "chunk-1", + "content": "Marigold is a depth-estimation model.", + "document_id": "doc-1", + "document_name": "paper.pdf", + }, + }, + "doc_aggs": []interface{}{ + map[string]interface{}{"doc_id": "doc-1", "doc_name": "paper.pdf", "count": 1}, + }, + "total": 1, + } + pipeline := &fakePipeline{ + resultChan: makeResultChan( + AsyncChatResult{Answer: "Marigold", Reference: map[string]interface{}{"chunks": []interface{}{}}, Final: false}, + AsyncChatResult{ + Answer: "Marigold is a depth-estimation model. [ID:0]", + Reference: finalReference, + Prompt: "### Query: what is marigold", + CreatedAt: 123, + Final: true, + }, + ), + } + + svc := &ChatSessionService{ + chatSessionDAO: store, + userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-owner"}}, + pipeline: pipeline, + } + + streamChan := make(chan string, 8) + _, err := svc.ChatCompletions( + context.Background(), + "user-1", + "dialog-1", + "session-1", + []map[string]interface{}{{"role": "user", "content": "what is marigold"}}, + "", + nil, + "", + nil, + nil, + false, + false, + true, + streamChan, + ) + if err != nil { + t.Fatalf("ChatCompletions failed: %v", err) + } + + var finalData map[string]interface{} + eventCount := len(streamChan) + for i := 0; i < eventCount; i++ { + event := <-streamChan + payload := strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(event), "data:")) + var wrapper map[string]interface{} + if err := json.Unmarshal([]byte(payload), &wrapper); err != nil { + t.Fatalf("failed to parse SSE payload %q: %v", payload, err) + } + data, ok := wrapper["data"].(map[string]interface{}) + if !ok { + continue + } + if final, _ := data["final"].(bool); final { + finalData = data + break + } + } + if finalData == nil { + t.Fatal("missing final SSE data") + } + if got := finalData["answer"]; got != "Marigold is a depth-estimation model. [ID:0]" { + t.Fatalf("final answer = %v", got) + } + if got := finalData["prompt"]; got != "### Query: what is marigold" { + t.Fatalf("final prompt = %v", got) + } + ref, ok := finalData["reference"].(map[string]interface{}) + if !ok { + t.Fatalf("final reference = %#v", finalData["reference"]) + } + chunks, ok := ref["chunks"].([]interface{}) + if !ok || len(chunks) != 1 { + t.Fatalf("final reference chunks = %#v", ref["chunks"]) + } + docAggs, ok := ref["doc_aggs"].([]interface{}) + if !ok || len(docAggs) != 1 { + t.Fatalf("final reference doc_aggs = %#v", ref["doc_aggs"]) + } + if got := ref["total"]; got != float64(1) { + t.Fatalf("final reference total = %v", got) + } +} + func TestCompletion_EmptyMessages(t *testing.T) { svc := &ChatSessionService{ chatSessionDAO: &fakeSessionStore{},