refactor: 重构文档处理流程和任务管理

common/eino/retriever.go

@@ -3,6 +3,8 @@ package eino
 import (
 	"context"
 	"errors"
+	"rag/dao"
+	"sort"
 
 	"github.com/cloudwego/eino/callbacks"
 	"github.com/cloudwego/eino/components/embedding"
@@ -16,12 +18,14 @@ type PGVectorRetrieverConfig struct {
 	Embedder     embedding.Embedder
 	DefaultTopK  int
 	DefaultIndex string
+	DSLInfo      map[string]any
 }
 
 type PGVectorRetriever struct {
 	embedder embedding.Embedder
 	topK     int
 	index    string
+	dslInfo  map[string]any
 }
 
 func NewPGVectorRetriever(config *PGVectorRetrieverConfig) (*PGVectorRetriever, error) {
@@ -36,43 +40,62 @@ func NewPGVectorRetriever(config *PGVectorRetrieverConfig) (*PGVectorRetriever,
 		embedder: config.Embedder,
 		topK:     config.DefaultTopK,
 		index:    config.DefaultIndex,
+		dslInfo:  config.DSLInfo,
 	}, nil
 }
 
 func (r *PGVectorRetriever) Retrieve(ctx context.Context, query string, opts ...retriever.Option) ([]*schema.Document, error) {
-
-	// 1. 处理公共 Option（官方标准写法）
 	options := &retriever.Options{
 		Index:     &r.index,
 		TopK:      &r.topK,
+		DSLInfo:   r.dslInfo,
 		Embedding: r.embedder,
 	}
 	options = retriever.GetCommonOptions(options, opts...)
 
-	// 2. 回调（官方标准）
 	ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{
 		Query: query,
 		TopK:  *options.TopK,
 	})
 
-	// 3. 执行检索
-	docs, err := r.doRetrieve(ctx, query, options)
+	// ==========================================
+	// 🔥 双路检索：向量 + 全文
+	// ==========================================
+	docsVector, err := r.doRetrieveVector(ctx, query, options)
 	if err != nil {
 		callbacks.OnError(ctx, err)
 		return nil, err
 	}
 
-	// 4. 完成回调
-	callbacks.OnEnd(ctx, &retriever.CallbackOutput{
-		Docs: docs,
+	docsFulltext, err := r.doRetrieveMeilisearch(ctx, query, options)
+	if err != nil {
+		callbacks.OnError(ctx, err)
+		return nil, err
+	}
+
+	// 合并 + 去重
+	docs := mergeAndDeduplicate(docsVector, docsFulltext)
+
+	// 排序（distance 越小越靠前）
+	sort.Slice(docs, func(i, j int) bool {
+		d1 := gconv.Float64(docs[i].MetaData["distance"])
+		d2 := gconv.Float64(docs[j].MetaData["distance"])
+		return d1 < d2
 	})
 
+	// 最多保留 topK
+	if len(docs) > *options.TopK {
+		docs = docs[:*options.TopK]
+	}
+
+	callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: docs})
 	return docs, nil
 }
 
-func (r *PGVectorRetriever) doRetrieve(ctx context.Context, query string, opts *retriever.Options) ([]*schema.Document, error) {
-
-	// 1. 生成向量
+// ==========================================
+// 1. 向量检索（PG）
+// ==========================================
+func (r *PGVectorRetriever) doRetrieveVector(ctx context.Context, query string, opts *retriever.Options) ([]*schema.Document, error) {
 	vectors, err := opts.Embedding.EmbedStrings(ctx, []string{query})
 	if err != nil {
 		return nil, err
@@ -81,37 +104,76 @@ func (r *PGVectorRetriever) doRetrieve(ctx context.Context, query string, opts *
 		return nil, errors.New("empty query vector")
 	}
 
-	queryVec := pgvector.NewVector(vectors[0])
+	queryVec := pgvector.NewVector(gconv.Float32s(vectors[0]))
 	topK := *opts.TopK
+	datasetIds := gconv.Int64s(opts.DSLInfo["dataset_ids"])
 
-	// 2. PG 向量相似度检索 SQL
-	sql := `
-	SELECT id, content, dataset_id, document_id,
-	       vector <-> ? AS distance
-	FROM document_chunk
-	ORDER BY distance ASC
-	LIMIT ?
-`
-
-	// 3. 查询
-	rows, err := dao.DocumentChunk.GetDB().GetAll(ctx, sql, queryVec, topK)
+	rows, err := dao.DocumentChunk.GetAllByVector(ctx, datasetIds, queryVec, topK)
 	if err != nil {
 		return nil, err
 	}
 
-	// 4. 转为 Eino Document
 	docs := make([]*schema.Document, 0, len(rows))
 	for _, row := range rows {
 		docs = append(docs, &schema.Document{
 			ID:      gconv.String(row["id"]),
 			Content: gconv.String(row["content"]),
-			Metadata: map[string]any{
-				"dataset_id":  row["dataset_id"],
-				"document_id": row["document_id"],
-				"distance":    row["distance"],
+			MetaData: map[string]any{
+				"dataset_id":  gconv.Int64(row["dataset_id"]),
+				"document_id": gconv.Int64(row["document_id"]),
+				"distance":    gconv.Float64(row["distance"]),
+				"retrieve_by": "vector",
 			},
 		})
 	}
-
 	return docs, nil
 }
+
+// ==========================================
+// 2. 全文检索（Meilisearch）🔥 新增
+// ==========================================
+func (r *PGVectorRetriever) doRetrieveMeilisearch(ctx context.Context, query string, opts *retriever.Options) ([]*schema.Document, error) {
+	topK := *opts.TopK
+	datasetIds := gconv.Int64s(opts.DSLInfo["dataset_ids"])
+
+	// 调用你已有的 Meilisearch DAO
+	rows, err := dao.DocumentChunk.SearchByKeywords(ctx, query, datasetIds, topK)
+	if err != nil {
+		return nil, err
+	}
+
+	docs := make([]*schema.Document, 0, len(rows))
+	for _, row := range rows {
+		docs = append(docs, &schema.Document{
+			ID:      gconv.String(row["id"]),
+			Content: gconv.String(row["content"]),
+			MetaData: map[string]any{
+				"dataset_id":  gconv.Int64(row["dataset_id"]),
+				"document_id": gconv.Int64(row["document_id"]),
+				"distance":    0.1, // 全文结果给高分
+				"retrieve_by": "fulltext",
+			},
+		})
+	}
+	return docs, nil
+}
+
+// ==========================================
+// 合并去重
+// ==========================================
+func mergeAndDeduplicate(vecDocs, fullDocs []*schema.Document) []*schema.Document {
+	idMap := make(map[string]*schema.Document)
+	for _, d := range vecDocs {
+		idMap[d.ID] = d
+	}
+	for _, d := range fullDocs {
+		if _, exists := idMap[d.ID]; !exists {
+			idMap[d.ID] = d
+		}
+	}
+	merged := make([]*schema.Document, 0, len(idMap))
+	for _, d := range idMap {
+		merged = append(merged, d)
+	}
+	return merged
+}