feat(go-api): add RAG retrieval to chat completions (#15739)

## Summary - Add knowledge-base retrieval support to Go chat completions. ## What changed - Routes KB-backed chat sessions through the Go retrieval service instead of falling back to solo chat. - Resolves embedding and rerank models, validates accessible knowledge bases, and preserves tenant-aware retrieval. - Rejects mixed embedding models across selected knowledge bases before retrieval to avoid incompatible vector dimensions. - Threads the HTTP request context into streaming retrieval so cancelled requests can stop downstream retrieval work. - Applies metadata filters and message-level `doc_ids` before retrieval. - Expands parent/child chunks before building references and prompt context. - Injects retrieved knowledge through a copied dialog prompt config so the caller's original dialog is not mutated. - Honors configured empty responses when no chunks are found. - Names the metadata no-match sentinel and reuses it across retrieval/handler paths. - Adds a defensive content cast while appending streamed answers. - Adds focused unit coverage for retrieval, metadata filtering, authorization, multimodal messages, references, empty-response behavior, prompt immutability, and mixed embedding models. --------- Co-authored-by: Yingfeng <yingfeng.zhang@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-09-07 17:58:55 +08:00 · 2026-06-09 20:07:45 -07:00
parent 7c1bd9a5a5
commit bbc1f2ecec
7 changed files with 1690 additions and 62 deletions
--- a/internal/service/chat_session.go
+++ b/internal/service/chat_session.go
@@ -17,10 +17,13 @@
 package service

 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"ragflow/internal/common"
+	"ragflow/internal/engine"
+	"ragflow/internal/service/nlp"
 	"strings"
 	"time"

@@ -31,21 +34,61 @@ import (
 	modelModule "ragflow/internal/entity/models"
 )

+type chatKnowledgebaseStore interface {
+	Accessible(kbID, userID string) bool
+	GetByIDs(ids []string) ([]*entity.Knowledgebase, error)
+}
+
+type chatModelProvider interface {
+	GetChatModel(tenantID, compositeModelName string) (*modelModule.ChatModel, error)
+	GetEmbeddingModel(tenantID, compositeModelName string) (*modelModule.EmbeddingModel, error)
+	GetRerankModel(tenantID, compositeModelName string) (*modelModule.RerankModel, error)
+	GetModelConfigFromProviderInstance(tenantID string, modelType entity.ModelType, modelName string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error)
+	GetTenantDefaultModelByType(tenantID string, modelType entity.ModelType) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error)
+}
+
+type chatMetadataService interface {
+	LabelQuestion(question string, kbs []*entity.Knowledgebase) map[string]float64
+	GetFlattedMetaByKBs(kbIDs []string) (common.MetaData, error)
+}
+
+type chatRetrievalService interface {
+	Retrieval(ctx context.Context, req *nlp.RetrievalRequest) (*nlp.RetrievalResult, error)
+}
+
 // ChatSessionService chat session (conversation) service
 type ChatSessionService struct {
 	chatSessionDAO   *dao.ChatSessionDAO
 	chatDAO          *dao.ChatDAO
 	userTenantDAO    *dao.UserTenantDAO
-	modelProviderSvc *ModelProviderService
+	kbDAO            chatKnowledgebaseStore
+	docEngine        engine.DocEngine
+	modelProviderSvc chatModelProvider
+	metadataSvc      chatMetadataService
+	retrievalSvc     chatRetrievalService
 }

 // NewChatSessionService create chat session service
 func NewChatSessionService() *ChatSessionService {
+	docEngine := engine.Get()
+	return newChatSessionServiceWithRetrieval(docEngine, nlp.NewRetrievalService(docEngine, dao.NewDocumentDAO()))
+}
+
+// NewChatSessionServiceWithRetrieval creates a chat session service with a retrieval service.
+func NewChatSessionServiceWithRetrieval(retrievalSvc chatRetrievalService) *ChatSessionService {
+	return newChatSessionServiceWithRetrieval(engine.Get(), retrievalSvc)
+}
+
+func newChatSessionServiceWithRetrieval(docEngine engine.DocEngine, retrievalSvc chatRetrievalService) *ChatSessionService {
 	return &ChatSessionService{
 		chatSessionDAO:   dao.NewChatSessionDAO(),
 		chatDAO:          dao.NewChatDAO(),
 		userTenantDAO:    dao.NewUserTenantDAO(),
+		kbDAO:            dao.NewKnowledgebaseDAO(),
+		docEngine:        docEngine,
 		modelProviderSvc: NewModelProviderService(),
+		metadataSvc:      NewMetadataService(),
+		retrievalSvc:     retrievalSvc,
 	}
 }

@@ -294,7 +337,7 @@ func (s *ChatSessionService) Completion(userID string, conversationID string, me
 	}

 	// Perform chat completion with RAG
-	result, err := s.asyncChat(dialog, session, messages, chatModelConfig, messageID, reference, false)
+	result, err := s.asyncChat(userID, dialog, session, messages, chatModelConfig, messageID, reference, false)
 	if err != nil {
 		return nil, err
 	}
@@ -308,7 +351,11 @@ func (s *ChatSessionService) Completion(userID string, conversationID string, me
 }

 // CompletionStream performs streaming chat completion with full RAG support
-func (s *ChatSessionService) CompletionStream(userID string, conversationID string, messages []map[string]interface{}, llmID string, chatModelConfig map[string]interface{}, messageID string, streamChan chan<- string) error {
+func (s *ChatSessionService) CompletionStream(ctx context.Context, userID string, conversationID string, messages []map[string]interface{}, llmID string, chatModelConfig map[string]interface{}, messageID string, streamChan chan<- string) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
 	// Validate the last message is from user
 	if len(messages) == 0 {
 		streamChan <- fmt.Sprintf("data: %s\n\n", `{"code": 500, "message": "messages cannot be empty", "data": {"answer": "**ERROR**: messages cannot be empty", "reference": []}}`)
@@ -356,7 +403,7 @@ func (s *ChatSessionService) CompletionStream(userID string, conversationID stri
 	}

 	// Perform streaming chat completion with RAG
-	resultChan, err := s.asyncChatStream(dialog, session, messages, chatModelConfig, messageID, reference)
+	resultChan, err := s.asyncChatStream(ctx, userID, dialog, session, messages, chatModelConfig, messageID, reference)
 	if err != nil {
 		streamChan <- fmt.Sprintf("data: %s\n\n", fmt.Sprintf(`{"code": 500, "message": "%s", "data": {"answer": "**ERROR**: %s", "reference": []}}`, err.Error(), err.Error()))
 		return err
@@ -450,7 +497,7 @@ func (s *ChatSessionService) updateSessionMessages(session *entity.ChatSession,
 }

 // asyncChat performs chat with RAG support (non-streaming)
-func (s *ChatSessionService) asyncChat(dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}, stream bool) (map[string]interface{}, error) {
+func (s *ChatSessionService) asyncChat(userID string, dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}, stream bool) (map[string]interface{}, error) {
 	// Check if we need RAG (knowledge base or tavily)
 	hasKB := len(dialog.KBIDs) > 0
 	hasTavily := false
@@ -465,21 +512,20 @@ func (s *ChatSessionService) asyncChat(dialog *entity.Chat, session *entity.Chat
 		return s.asyncChatSolo(dialog, session, messages, config, messageID, reference, stream)
 	}

-	// TODO: Full RAG implementation with knowledge base retrieval
-	// This would include:
-	// 1. Get embedding model and rerank model
-	// 2. Extract questions from messages
-	// 3. Retrieve chunks from knowledge bases
-	// 4. Rerank chunks
-	// 5. Build prompt with context
-	// 6. Call LLM
+	if hasKB {
+		return s.asyncChatWithRetrieval(context.Background(), userID, dialog, session, messages, config, messageID, reference, stream)
+	}

-	// For now, fall back to solo chat
+	common.Warn("Tavily-backed chat retrieval is not implemented in Go; falling back to solo chat",
+		zap.String("dialog_id", dialog.ID))
 	return s.asyncChatSolo(dialog, session, messages, config, messageID, reference, stream)
 }

 // asyncChatStream performs streaming chat with RAG support
-func (s *ChatSessionService) asyncChatStream(dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}) (<-chan map[string]interface{}, error) {
+func (s *ChatSessionService) asyncChatStream(ctx context.Context, userID string, dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}) (<-chan map[string]interface{}, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	resultChan := make(chan map[string]interface{})

 	go func() {
@@ -500,14 +546,599 @@ func (s *ChatSessionService) asyncChatStream(dialog *entity.Chat, session *entit
 			return
 		}

-		// TODO: Full RAG streaming implementation
-		// For now, fall back to solo chat
+		if hasKB {
+			ragMessages, ragDialog, emptyResponse, err := s.messagesWithRetrievedKnowledge(ctx, userID, dialog, messages, reference)
+			if err != nil {
+				resultChan <- s.structureAnswer(session, "**ERROR**: "+err.Error(), messageID, session.ID, reference)
+				return
+			}
+			if emptyResponse != nil {
+				resultChan <- s.structureAnswer(session, *emptyResponse, messageID, session.ID, reference)
+				return
+			}
+			s.asyncChatSoloStream(ragDialog, session, ragMessages, config, messageID, reference, resultChan)
+			return
+		}
+
+		common.Warn("Tavily-backed streaming chat retrieval is not implemented in Go; falling back to solo chat",
+			zap.String("dialog_id", dialog.ID))
 		s.asyncChatSoloStream(dialog, session, messages, config, messageID, reference, resultChan)
 	}()

 	return resultChan, nil
 }

+func (s *ChatSessionService) asyncChatWithRetrieval(ctx context.Context, userID string, dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}, stream bool) (map[string]interface{}, error) {
+	ragMessages, ragDialog, emptyResponse, err := s.messagesWithRetrievedKnowledge(ctx, userID, dialog, messages, reference)
+	if err != nil {
+		return nil, err
+	}
+	if emptyResponse != nil {
+		var lastRef interface{}
+		if len(reference) > 0 {
+			lastRef = reference[len(reference)-1]
+		}
+		ans := map[string]interface{}{
+			"answer":    *emptyResponse,
+			"reference": lastRef,
+			"final":     true,
+		}
+		return s.structureAnswerWithConv(session, ans, messageID, session.ID, reference), nil
+	}
+	return s.asyncChatSolo(ragDialog, session, ragMessages, config, messageID, reference, stream)
+}
+
+func (s *ChatSessionService) messagesWithRetrievedKnowledge(ctx context.Context, userID string, dialog *entity.Chat, messages []map[string]interface{}, reference []interface{}) ([]map[string]interface{}, *entity.Chat, *string, error) {
+	kbIDs := stringSliceFromJSON(dialog.KBIDs)
+	if len(kbIDs) == 0 {
+		return messages, dialog, nil, nil
+	}
+	if s.retrievalSvc == nil {
+		return nil, nil, nil, errors.New("retrieval service is not configured")
+	}
+
+	question := latestUserQuestion(messages)
+	if question == "" {
+		return messages, dialog, nil, nil
+	}
+
+	kbs, err := s.kbDAO.GetByIDs(kbIDs)
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("failed to load knowledge bases: %w", err)
+	}
+	kbs, err = s.knowledgebasesForDialog(userID, dialog, kbIDs, kbs)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	embeddingTenantID, embeddingModelName, err := validateKnowledgebaseEmbeddingModels(kbs, dialog.TenantID, resolveEmbeddingModelName)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	embeddingModel, err := s.modelProviderSvc.GetEmbeddingModel(embeddingTenantID, embeddingModelName)
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("failed to get embedding model: %w", err)
+	}
+	rerankModel, err := s.rerankModelForDialog(dialog)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	top := int(dialog.TopK)
+	pageSize := int(dialog.TopN)
+	if pageSize <= 0 {
+		pageSize = 6
+	}
+	similarityThreshold := dialog.SimilarityThreshold
+	vectorSimilarityWeight := dialog.VectorSimilarityWeight
+	var rankFeature map[string]float64
+	if s.metadataSvc != nil {
+		rankFeature = s.metadataSvc.LabelQuestion(question, kbs)
+	}
+	baseDocIDs := docIDsFromMessages(messages)
+	docIDs, err := s.filteredDocIDsForDialog(ctx, dialog, kbIDs, question, baseDocIDs)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	tenantIDs := tenantIDsFromKnowledgebases(kbs, dialog.TenantID)
+
+	retrievalResult, err := s.retrievalSvc.Retrieval(ctx, &nlp.RetrievalRequest{
+		Question:               question,
+		TenantIDs:              tenantIDs,
+		KbIDs:                  kbIDs,
+		DocIDs:                 docIDs,
+		Page:                   1,
+		PageSize:               pageSize,
+		Top:                    &top,
+		SimilarityThreshold:    &similarityThreshold,
+		VectorSimilarityWeight: &vectorSimilarityWeight,
+		RankFeature:            &rankFeature,
+		EmbeddingModel:         embeddingModel,
+		RerankModel:            rerankModel,
+	})
+	if err != nil {
+		return nil, nil, nil, fmt.Errorf("retrieval search failed: %w", err)
+	}
+	if retrievalResult == nil {
+		retrievalResult = &nlp.RetrievalResult{}
+	}
+
+	chunks := retrievalResult.Chunks
+	if s.docEngine != nil {
+		chunks = nlp.RetrievalByChildren(chunks, tenantIDs, s.docEngine, ctx)
+	}
+	setLatestReference(reference, chunks, retrievalResult.DocAggs)
+	knowledge := buildKnowledgeBlock(chunks)
+	if knowledge == "" {
+		return messages, dialog, emptyResponseForDialog(dialog), nil
+	}
+	if ragDialog, ok := dialogWithInjectedKnowledgePrompt(dialog, knowledge); ok {
+		return copyMessages(messages), ragDialog, nil, nil
+	}
+
+	return injectKnowledge(messages, knowledge), dialog, nil, nil
+}
+
+type embeddingModelNameResolver func(tenantID string, kb *entity.Knowledgebase) (string, error)
+
+func validateKnowledgebaseEmbeddingModels(kbs []*entity.Knowledgebase, fallbackTenantID string, resolve embeddingModelNameResolver) (string, string, error) {
+	if len(kbs) == 0 {
+		return fallbackTenantID, "", nil
+	}
+
+	expected := ""
+	expectedKBID := ""
+	expectedTenantID := fallbackTenantID
+	for _, kb := range kbs {
+		if kb == nil {
+			return "", "", errors.New("knowledge base is nil")
+		}
+		tenantID := kb.TenantID
+		if tenantID == "" {
+			tenantID = fallbackTenantID
+		}
+		modelName, err := resolve(tenantID, kb)
+		if err != nil {
+			return "", "", err
+		}
+		modelName = strings.TrimSpace(modelName)
+		if modelName == "" {
+			return "", "", fmt.Errorf("knowledge base %s has no embedding model", kb.ID)
+		}
+		if expected == "" {
+			expected = modelName
+			expectedKBID = kb.ID
+			expectedTenantID = tenantID
+			continue
+		}
+		if modelName != expected {
+			return "", "", fmt.Errorf("knowledge bases must use the same embedding model: %s resolves to %q, expected %q from %s", kb.ID, modelName, expected, expectedKBID)
+		}
+	}
+	return expectedTenantID, expected, nil
+}
+
+func (s *ChatSessionService) rerankModelForDialog(dialog *entity.Chat) (*modelModule.RerankModel, error) {
+	compositeName, err := resolveRerankModelName(dialog)
+	if err != nil {
+		return nil, err
+	}
+	if compositeName == "" {
+		return nil, nil
+	}
+	rerankModel, err := s.modelProviderSvc.GetRerankModel(dialog.TenantID, compositeName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get rerank model: %w", err)
+	}
+	return rerankModel, nil
+}
+
+func (s *ChatSessionService) filteredDocIDsForDialog(ctx context.Context, dialog *entity.Chat, kbIDs []string, question string, baseDocIDs []string) ([]string, error) {
+	if dialog.MetaDataFilter == nil || len(*dialog.MetaDataFilter) == 0 {
+		return baseDocIDs, nil
+	}
+	if s.metadataSvc == nil {
+		return nil, errors.New("metadata service is not configured")
+	}
+
+	filter := make(map[string]interface{}, len(*dialog.MetaDataFilter))
+	for key, value := range *dialog.MetaDataFilter {
+		filter[key] = value
+	}
+
+	metaData, err := s.metadataSvc.GetFlattedMetaByKBs(kbIDs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get flattened metadata for chat retrieval: %w", err)
+	}
+
+	var filterChatModel *modelModule.ChatModel
+	method, _ := filter["method"].(string)
+	if method == "auto" || method == "semi_auto" {
+		filterChatModel, err = s.modelProviderSvc.GetChatModel(dialog.TenantID, dialog.LLMID)
+		if err != nil {
+			common.Warn("Failed to get chat model for chat metadata filter", zap.Error(err))
+		}
+	}
+
+	docIDs, empty := ApplyMetaDataFilter(ctx, filter, metaData, question, filterChatModel, baseDocIDs, kbIDs)
+	if empty {
+		return []string{NoMatchDocIDSentinel}, nil
+	}
+	return docIDs, nil
+}
+
+func resolveEmbeddingModelName(tenantID string, kb *entity.Knowledgebase) (string, error) {
+	if kb.TenantEmbdID != nil && *kb.TenantEmbdID > 0 {
+		_, compositeName, err := dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), *kb.TenantEmbdID)
+		if err != nil {
+			return "", fmt.Errorf("failed to get embedding model by tenant_embd_id: %w", err)
+		}
+		return compositeName, nil
+	}
+	if kb.EmbdID != "" {
+		if strings.Contains(kb.EmbdID, "@") {
+			return kb.EmbdID, nil
+		}
+		_, compositeName, err := dao.LookupTenantLLMByName(dao.NewTenantLLMDAO(), tenantID, kb.EmbdID, entity.ModelTypeEmbedding)
+		if err != nil {
+			return "", fmt.Errorf("failed to get embedding model by embd_id: %w", err)
+		}
+		return compositeName, nil
+	}
+
+	tenantLLM, err := dao.NewTenantLLMDAO().GetByTenantAndType(tenantID, entity.ModelTypeEmbedding)
+	if err != nil {
+		return "", fmt.Errorf("failed to get tenant default embedding model: %w", err)
+	}
+	if tenantLLM == nil || tenantLLM.LLMName == nil || *tenantLLM.LLMName == "" {
+		return "", fmt.Errorf("no default embedding model found for tenant %s", tenantID)
+	}
+	return fmt.Sprintf("%s@%s", *tenantLLM.LLMName, tenantLLM.LLMFactory), nil
+}
+
+func resolveRerankModelName(dialog *entity.Chat) (string, error) {
+	if dialog.TenantRerankID != nil && *dialog.TenantRerankID > 0 {
+		_, compositeName, err := dao.LookupTenantLLMByID(dao.NewTenantLLMDAO(), *dialog.TenantRerankID)
+		if err != nil {
+			return "", fmt.Errorf("failed to get rerank model by tenant_rerank_id: %w", err)
+		}
+		return compositeName, nil
+	}
+	if dialog.RerankID == "" {
+		return "", nil
+	}
+	if strings.Contains(dialog.RerankID, "@") {
+		return dialog.RerankID, nil
+	}
+	_, compositeName, err := dao.LookupTenantLLMByName(dao.NewTenantLLMDAO(), dialog.TenantID, dialog.RerankID, entity.ModelTypeRerank)
+	if err != nil {
+		return "", fmt.Errorf("failed to get rerank model by rerank_id: %w", err)
+	}
+	return compositeName, nil
+}
+
+func stringSliceFromJSON(values entity.JSONSlice) []string {
+	result := make([]string, 0, len(values))
+	seen := make(map[string]struct{}, len(values))
+	for _, value := range values {
+		str, ok := value.(string)
+		if !ok || str == "" {
+			continue
+		}
+		if _, exists := seen[str]; exists {
+			continue
+		}
+		seen[str] = struct{}{}
+		result = append(result, str)
+	}
+	return result
+}
+
+func tenantIDsFromKnowledgebases(kbs []*entity.Knowledgebase, fallback string) []string {
+	seen := make(map[string]struct{}, len(kbs)+1)
+	var tenantIDs []string
+	for _, kb := range kbs {
+		if kb == nil || kb.TenantID == "" {
+			continue
+		}
+		if _, exists := seen[kb.TenantID]; exists {
+			continue
+		}
+		seen[kb.TenantID] = struct{}{}
+		tenantIDs = append(tenantIDs, kb.TenantID)
+	}
+	if len(tenantIDs) == 0 && fallback != "" {
+		tenantIDs = append(tenantIDs, fallback)
+	}
+	return tenantIDs
+}
+
+func (s *ChatSessionService) knowledgebasesForDialog(userID string, dialog *entity.Chat, kbIDs []string, loaded []*entity.Knowledgebase) ([]*entity.Knowledgebase, error) {
+	byID := make(map[string]*entity.Knowledgebase, len(loaded))
+	for _, kb := range loaded {
+		if kb != nil {
+			byID[kb.ID] = kb
+		}
+	}
+
+	kbs := make([]*entity.Knowledgebase, 0, len(kbIDs))
+	for _, kbID := range kbIDs {
+		kb := byID[kbID]
+		if kb == nil {
+			return nil, fmt.Errorf("knowledge base %s not found", kbID)
+		}
+		if userID != "" && !s.kbDAO.Accessible(kbID, userID) {
+			return nil, fmt.Errorf("knowledge base %s is not authorized for user", kbID)
+		}
+		if userID == "" && kb.TenantID != dialog.TenantID {
+			return nil, fmt.Errorf("knowledge base %s is not authorized for dialog tenant", kbID)
+		}
+		kbs = append(kbs, kb)
+	}
+	if len(kbs) == 0 {
+		return nil, errors.New("no valid knowledge bases found")
+	}
+	return kbs, nil
+}
+
+func docIDsFromMessages(messages []map[string]interface{}) []string {
+	for i := len(messages) - 1; i >= 0; i-- {
+		if role, _ := messages[i]["role"].(string); role != "user" {
+			continue
+		}
+		return stringSliceFromValue(messages[i]["doc_ids"])
+	}
+	return nil
+}
+
+func latestUserQuestion(messages []map[string]interface{}) string {
+	for i := len(messages) - 1; i >= 0; i-- {
+		if role, _ := messages[i]["role"].(string); role != "user" {
+			continue
+		}
+		return textFromMessageContent(messages[i]["content"])
+	}
+	return ""
+}
+
+func stringSliceFromValue(value interface{}) []string {
+	switch typed := value.(type) {
+	case nil:
+		return nil
+	case []string:
+		return uniqueNonEmptyStrings(typed)
+	case []interface{}:
+		values := make([]string, 0, len(typed))
+		for _, item := range typed {
+			if str, ok := item.(string); ok {
+				values = append(values, str)
+			}
+		}
+		return uniqueNonEmptyStrings(values)
+	default:
+		return nil
+	}
+}
+
+func uniqueNonEmptyStrings(values []string) []string {
+	result := make([]string, 0, len(values))
+	seen := make(map[string]struct{}, len(values))
+	for _, value := range values {
+		value = strings.TrimSpace(value)
+		if value == "" {
+			continue
+		}
+		if _, exists := seen[value]; exists {
+			continue
+		}
+		seen[value] = struct{}{}
+		result = append(result, value)
+	}
+	if len(result) == 0 {
+		return nil
+	}
+	return result
+}
+
+func emptyResponseForDialog(dialog *entity.Chat) *string {
+	if dialog.PromptConfig == nil {
+		return nil
+	}
+	emptyResponse, ok := dialog.PromptConfig["empty_response"].(string)
+	if !ok || emptyResponse == "" {
+		return nil
+	}
+	return &emptyResponse
+}
+
+func buildKnowledgeBlock(chunks []map[string]interface{}) string {
+	var builder strings.Builder
+	for i, chunk := range chunks {
+		content := chunkText(chunk)
+		if content == "" {
+			continue
+		}
+		if builder.Len() > 0 {
+			builder.WriteString("\n\n")
+		}
+		builder.WriteString(fmt.Sprintf("[%d]", i+1))
+		if docName, ok := chunk["docnm_kwd"].(string); ok && docName != "" {
+			builder.WriteString(" ")
+			builder.WriteString(docName)
+		}
+		builder.WriteString("\n")
+		builder.WriteString(content)
+	}
+	return builder.String()
+}
+
+func chunkText(chunk map[string]interface{}) string {
+	for _, key := range []string{"content_with_weight", "content_ltks", "content"} {
+		if value, ok := chunk[key].(string); ok && strings.TrimSpace(value) != "" {
+			return strings.TrimSpace(value)
+		}
+	}
+	return ""
+}
+
+func injectKnowledge(messages []map[string]interface{}, knowledge string) []map[string]interface{} {
+	copied := copyMessages(messages)
+	if len(copied) == 0 {
+		return copied
+	}
+
+	knowledgePrompt := fmt.Sprintf("Use the following knowledge snippets to answer the user's question. If the snippets do not contain the answer, say that the knowledge base does not provide enough information.\n\n%s", knowledge)
+	for i := len(copied) - 1; i >= 0; i-- {
+		if role, _ := copied[i]["role"].(string); role != "user" {
+			continue
+		}
+		copied[i]["content"] = injectKnowledgeIntoContent(copied[i]["content"], knowledgePrompt)
+		return copied
+	}
+
+	copied = append(copied, map[string]interface{}{
+		"role":    "system",
+		"content": knowledgePrompt,
+	})
+	return copied
+}
+
+func injectKnowledgeIntoContent(content interface{}, knowledgePrompt string) interface{} {
+	switch typed := content.(type) {
+	case []interface{}:
+		injected := make([]interface{}, 0, len(typed)+1)
+		injected = append(injected, knowledgeTextBlock(knowledgePrompt))
+		injected = append(injected, typed...)
+		return injected
+	case []map[string]interface{}:
+		injected := make([]interface{}, 0, len(typed)+1)
+		injected = append(injected, knowledgeTextBlock(knowledgePrompt))
+		for _, block := range typed {
+			injected = append(injected, block)
+		}
+		return injected
+	default:
+		contentText := ""
+		if content != nil {
+			contentText = fmt.Sprint(content)
+		}
+		return strings.TrimSpace(knowledgePrompt + "\n\nQuestion:\n" + contentText)
+	}
+}
+
+func knowledgeTextBlock(knowledgePrompt string) map[string]interface{} {
+	return map[string]interface{}{
+		"type": "text",
+		"text": knowledgePrompt + "\n\nQuestion:",
+	}
+}
+
+func textFromMessageContent(content interface{}) string {
+	switch typed := content.(type) {
+	case string:
+		return strings.TrimSpace(typed)
+	case []interface{}:
+		return strings.TrimSpace(strings.Join(textsFromContentBlocks(typed), "\n"))
+	case []map[string]interface{}:
+		blocks := make([]interface{}, 0, len(typed))
+		for _, block := range typed {
+			blocks = append(blocks, block)
+		}
+		return strings.TrimSpace(strings.Join(textsFromContentBlocks(blocks), "\n"))
+	default:
+		if content == nil {
+			return ""
+		}
+		return strings.TrimSpace(fmt.Sprint(content))
+	}
+}
+
+func textsFromContentBlocks(blocks []interface{}) []string {
+	texts := make([]string, 0, len(blocks))
+	for _, block := range blocks {
+		switch typed := block.(type) {
+		case string:
+			if text := strings.TrimSpace(typed); text != "" {
+				texts = append(texts, text)
+			}
+		case map[string]interface{}:
+			if text, ok := typed["text"].(string); ok && strings.TrimSpace(text) != "" {
+				texts = append(texts, strings.TrimSpace(text))
+			}
+		}
+	}
+	return texts
+}
+
+func dialogWithInjectedKnowledgePrompt(dialog *entity.Chat, knowledge string) (*entity.Chat, bool) {
+	if dialog.PromptConfig == nil {
+		return dialog, false
+	}
+	systemPrompt, ok := dialog.PromptConfig["system"].(string)
+	if !ok || !strings.Contains(systemPrompt, "{knowledge}") {
+		return dialog, false
+	}
+
+	copied := cloneJSONMap(dialog.PromptConfig)
+	copied["system"] = strings.ReplaceAll(systemPrompt, "{knowledge}", knowledge)
+	dialogCopy := *dialog
+	dialogCopy.PromptConfig = copied
+	return &dialogCopy, true
+}
+
+func cloneJSONMap(values entity.JSONMap) entity.JSONMap {
+	copied := make(entity.JSONMap, len(values))
+	for key, value := range values {
+		copied[key] = value
+	}
+	return copied
+}
+
+func copyMessages(messages []map[string]interface{}) []map[string]interface{} {
+	copied := make([]map[string]interface{}, len(messages))
+	for i, msg := range messages {
+		copied[i] = make(map[string]interface{}, len(msg))
+		for key, value := range msg {
+			copied[i][key] = value
+		}
+	}
+	return copied
+}
+
+func setLatestReference(reference []interface{}, chunks []map[string]interface{}, docAggs []map[string]interface{}) {
+	ref := map[string]interface{}{
+		"chunks":   chunksForReference(chunks),
+		"doc_aggs": mapsForReference(docAggs),
+	}
+	if len(reference) == 0 {
+		return
+	}
+	reference[len(reference)-1] = ref
+}
+
+func chunksForReference(chunks []map[string]interface{}) []interface{} {
+	result := make([]interface{}, 0, len(chunks))
+	for _, chunk := range chunks {
+		copied := make(map[string]interface{}, len(chunk))
+		for key, value := range chunk {
+			if key == "vector" {
+				continue
+			}
+			copied[key] = value
+		}
+		result = append(result, copied)
+	}
+	return result
+}
+
+func mapsForReference(values []map[string]interface{}) []interface{} {
+	result := make([]interface{}, 0, len(values))
+	for _, value := range values {
+		result = append(result, value)
+	}
+	return result
+}
+
 // asyncChatSolo performs simple chat without RAG (non-streaming)
 func (s *ChatSessionService) asyncChatSolo(dialog *entity.Chat, session *entity.ChatSession, messages []map[string]interface{}, config map[string]interface{}, messageID string, reference []interface{}, stream bool) (map[string]interface{}, error) {
 	common.Info("asyncChatSolo started",
@@ -765,7 +1396,8 @@ func (s *ChatSessionService) structureAnswerWithConv(session *entity.ChatSession
 			if ans["final"] == true && ans["answer"] != nil {
 				lastMsg["content"] = ans["answer"]
 			} else {
-				lastMsg["content"] = (lastMsg["content"].(string)) + content
+				existing, _ := lastMsg["content"].(string)
+				lastMsg["content"] = existing + content
 			}
 			lastMsg["created_at"] = float64(time.Now().Unix())
 			lastMsg["id"] = messageID