go get github.com/qdrant/go-client/qdrant
go get github.com/cloudwego/eino-ext/components/indexer/qdrant
go get github.com/cloudwego/eino-ext/components/retriever/qdrant

var qdrantClient *qdrant.Client
var so sync.Once

func NewQdrantClient() *qdrant.Client {
	so.Do(func() {
		var err error
		qdrantClient, err = qdrant.NewClient(&qdrant.Config{
			Host: "127.0.0.1",
			Port: 6334,
		})
		if err != nil {
			log.Fatalf("NewClient failed, err=%v", err)
		}
	})
	return qdrantClient
}

func NewQdrantIndexer(ctx context.Context, client *qdrant.Client) *qdrant_indexer.Indexer {
	CollectName := "test_collection"
	var indexer *qdrant_indexer.Indexer
	
	if true, err := client.CollectionExists(ctx, CollectName); err == nil {
		// 如果集合已存在，先删除
		if err := client.DeleteCollection(ctx, CollectName); err != nil {
			log.Fatalf("DeleteCollection failed, err=%v", err)
		}
		
		// 创建新的索引器
		indexer, err = qdrant_indexer.NewIndexer(ctx, &qdrant_indexer.Config{
			Client: client,
			Embedding: &NormalizedEmbedder{
				OriginalEmbedder: Embeddder,
			},
			VectorDim:  768,                 // 向量维度
			Distance:   qdrant.Distance_Dot, // 使用内积距离
			Collection: CollectName,
		})
		if err != nil {
			log.Fatalf("NewIndexer failed, err=%v", err)
		}
	}
	return indexer
}

type NormalizedEmbedder struct {
	OriginalEmbedder *ollama.Embedder
}

func (ne *NormalizedEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...embedding.Option) ([][]float64, error) {
	// 1. 调用原始Embedding模型获取向量
	vectors, err := ne.OriginalEmbedder.EmbedStrings(ctx, texts, opts...)
	if err != nil {
		return nil, err
	}

	// 2. 对向量进行归一化
	for i, v := range vectors {
		vectors[i] = NormalizeVector(v)
	}

	// 3. 返回模长为1的向量
	return vectors, nil
}

func StoreEmbeddings(ctx context.Context, docs []*schema.Document, indexer *qdrant_indexer.Indexer) error {
	ids, err := indexer.Store(ctx, docs)
	if err != nil {
		return err
	}
	fmt.Printf("Store %d documents, ids: %v\n", len(ids), ids)
	return nil
}

func GetEmbeddings(ctx context.Context, query string, CollectName string, client *qdrant.Client, indexer *qdrant_indexer.Indexer) ([]*schema.Document, error) {
	scoreThresh := 0.3
	retriever, _ := qdrant_retriever.NewRetriever(ctx, &qdrant_retriever.Config{
		Client:        client,
		Embedding:     &NormalizedEmbedder{OriginalEmbedder: Embeddder},
		TopK:          20,            // 返回最相似的20个文档
		ScoreThreshold: &scoreThresh, // 相似度阈值
		Collection:    CollectName,
	})
	
	// 执行检索，并添加过滤条件
	neighbors, _ := retriever.Retrieve(ctx, query,
		qdrant_retriever.WithFilter(&qdrant.Filter{
			Must: []*qdrant.Condition{
				qdrant.NewMatch("metadata.source", "sohu"), // 只返回source为sohu的文档
			},
		}),
	)
	return neighbors, nil
}

func TestGetEmbeddings(t *testing.T) {
	ctx := context.Background()
	
	// 准备测试文档
	docs := []*schema.Document{
		{
			ID:      uuid.NewString(),
			Content: "美国9个国家公共假日，元旦：1月1日；马丁路德金日：一月的第三个星期一...",
			MetaData: map[string]any{
				"source": "sohu",
			},
		},
		{
			ID:      uuid.NewString(),
			Content: "美国平均退休年龄：66.33岁",
			MetaData: map[string]any{
				"source": "sina",
			},
		},
		{
			ID:      uuid.NewString(),
			Content: "美国平均退休年龄：66.33岁",
			MetaData: map[string]any{
				"source": "sohu",
			},
		},
	}
	
	// 初始化
	NewEmbedder()
	client := NewQdrantClient()
	indexer := NewQdrantIndexer(ctx, client)
	
	// 存储文档
	err := StoreEmbeddings(ctx, docs, indexer)
	if err != nil {
		t.Fatal(err)
	}
	
	// 语义检索
	docreturn, err := GetEmbeddings(ctx, "美国多少岁可以退休？", "test_collection", client, indexer)
	if err != nil {
		t.Fatalf("GetEmbeddings failed: %v", err)
	}
	
	// 输出结果
	for i, doc := range docreturn {
		fmt.Printf("%d  %.4f  %s\n", i, doc.Score(), doc.Content)
	}
}

go test -v -run TestGetEmbeddings ./rag/

// 归一化后，内积 = 余弦相似度
Distance: qdrant.Distance_Dot

TopK: 20,  // 根据实际需求调整，不是越大越好

scoreThresh := 0.3
ScoreThreshold: &scoreThresh,

qdrant_retriever.WithFilter(&qdrant.Filter{
    Must: []*qdrant.Condition{
        qdrant.NewMatch("metadata.source", "sohu"),
    },
})