Merge pull request #1402 from xiaomakuaiz/feat-search-simlarity

feat: change search similarity to 0.2

Author: xiaomakuaiz

backend/store/rag/ct/rag.go

@@ -47,7 +47,7 @@ func (s *CTRAG) CreateKnowledgeBase(ctx context.Context) (string, error) {
 	return dataset.ID, nil
 }
 
-func (s *CTRAG) QueryRecords(ctx context.Context, datasetIDs []string, query string, groupIds []int, historyMsgs []*schema.Message) ([]*domain.NodeContentChunk, error) {
+func (s *CTRAG) QueryRecords(ctx context.Context, datasetIDs []string, query string, groupIds []int, similarityThreshold float64, historyMsgs []*schema.Message) ([]*domain.NodeContentChunk, error) {
 	var chatMsgs []rag.ChatMessage
 	for _, msg := range historyMsgs {
 		switch msg.Role {
@@ -66,14 +66,17 @@ func (s *CTRAG) QueryRecords(ctx context.Context, datasetIDs []string, query str
 		}
 	}
 	s.logger.Debug("retrieving by history msgs", log.Any("history_msgs", historyMsgs), log.Any("chat_msgs", chatMsgs))
-	chunks, _, rewriteQuery, err := s.client.RetrieveChunks(ctx, rag.RetrievalRequest{
+	retrieveReq := rag.RetrievalRequest{
 		DatasetIDs:   datasetIDs,
 		Question:     query,
 		TopK:         10,
 		UserGroupIDs: groupIds,
 		ChatMessages: chatMsgs,
-		// SimilarityThreshold: 0.2,
-	})
+	}
+	if similarityThreshold != 0 {
+		retrieveReq.SimilarityThreshold = similarityThreshold
+	}
+	chunks, _, rewriteQuery, err := s.client.RetrieveChunks(ctx, retrieveReq)
 	s.logger.Info("retrieve chunks result", log.Int("chunks count", len(chunks)), log.String("query", rewriteQuery))
 
 	if err != nil {

backend/store/rag/rag.go

@@ -16,7 +16,7 @@ import (
 type RAGService interface {
 	CreateKnowledgeBase(ctx context.Context) (string, error)
 	UpsertRecords(ctx context.Context, datasetID string, nodeRelease *domain.NodeReleaseWithDirPath, authGroupId []int) (string, error)
-	QueryRecords(ctx context.Context, datasetIDs []string, query string, groupIDs []int, historyMsgs []*schema.Message) ([]*domain.NodeContentChunk, error)
+	QueryRecords(ctx context.Context, datasetIDs []string, query string, groupIDs []int, similarityThreshold float64, historyMsgs []*schema.Message) ([]*domain.NodeContentChunk, error)
 	DeleteRecords(ctx context.Context, datasetID string, docIDs []string) error
 	DeleteKnowledgeBase(ctx context.Context, datasetID string) error
 	UpdateDocumentGroupIDs(ctx context.Context, datasetID string, docID string, groupIds []int) error

backend/usecase/chat.go

@@ -372,7 +372,7 @@ func (u *ChatUsecase) Search(ctx context.Context, req *domain.ChatSearchReq) (*d
 	if err != nil {
 		return nil, err
 	}
-	rankedNodes, err := u.llmUsecase.GetRankNodes(ctx, []string{kb.DatasetID}, req.Message, groupIds, nil)
+	rankedNodes, err := u.llmUsecase.GetRankNodes(ctx, []string{kb.DatasetID}, req.Message, groupIds, 0.2, nil)
 	if err != nil {
 		return nil, err
 	}

backend/usecase/llm.go

@@ -100,7 +100,7 @@ func (u *LLMUsecase) FormatConversationMessages(
 			if err != nil {
 				return nil, nil, fmt.Errorf("get kb failed: %w", err)
 			}
-			rankedNodes, err = u.GetRankNodes(ctx, []string{kb.DatasetID}, question, groupIDs, historyMessages[:len(historyMessages)-1])
+			rankedNodes, err = u.GetRankNodes(ctx, []string{kb.DatasetID}, question, groupIDs, 0, historyMessages[:len(historyMessages)-1])
 			if err != nil {
 				return nil, nil, fmt.Errorf("get rank nodes failed: %w", err)
 			}
@@ -297,10 +297,17 @@ func (u *LLMUsecase) SplitByTokenLimit(text string, maxTokens int) ([]string, er
 	return result, nil
 }
 
-func (u *LLMUsecase) GetRankNodes(ctx context.Context, datasetIDs []string, question string, groupIDs []int, historyMessages []*schema.Message) ([]*domain.RankedNodeChunks, error) {
+func (u *LLMUsecase) GetRankNodes(
+	ctx context.Context,
+	datasetIDs []string,
+	question string,
+	groupIDs []int,
+	similarityThreshold float64,
+	historyMessages []*schema.Message,
+) ([]*domain.RankedNodeChunks, error) {
 	var rankedNodes []*domain.RankedNodeChunks
 	// get related documents from raglite
-	records, err := u.rag.QueryRecords(ctx, datasetIDs, question, groupIDs, historyMessages)
+	records, err := u.rag.QueryRecords(ctx, datasetIDs, question, groupIDs, similarityThreshold, historyMessages)
 	if err != nil {
 		return nil, fmt.Errorf("get records from raglite failed: %w", err)
 	}