feat: refactor retrieve logic to enhance trace handling and knowledge base preparation

internal/ai/rag/retrieve_flow.go

@@ -0,0 +1,38 @@
+package rag
+
+import (
+	"fmt"
+	"log/slog"
+
+	"cs-agent/internal/models"
+)
+
+func newRetrieveTrace() *RetrieveTrace {
+	return &RetrieveTrace{}
+}
+
+func (s *retrieve) prepareRetrievableKnowledgeBases(req RetrieveRequest, trace *RetrieveTrace) ([]models.KnowledgeBase, []int64, bool) {
+	if req.Query == "" {
+		return nil, nil, false
+	}
+	knowledgeBaseIDs := normalizeKnowledgeBaseIDs(req.KnowledgeBaseIDs)
+	if len(knowledgeBaseIDs) == 0 {
+		return nil, nil, false
+	}
+
+	retrievableKnowledgeBases := s.loadRetrievableKnowledgeBases(knowledgeBaseIDs)
+	if len(retrievableKnowledgeBases) == 0 {
+		slog.Info("Skip retrieve for non-enabled knowledge bases",
+			"knowledge_base_ids", fmt.Sprint(knowledgeBaseIDs))
+		return nil, knowledgeBaseIDs, false
+	}
+	return retrievableKnowledgeBases, knowledgeBaseIDs, true
+}
+
+func applySearchTrace(target *RetrieveTrace, source *RetrieveTrace) {
+	if target == nil || source == nil {
+		return
+	}
+	target.EmbeddingMs = source.EmbeddingMs
+	target.VectorSearchMs = source.VectorSearchMs
+}