108 lines
2.5 KiB
Go
108 lines
2.5 KiB
Go
package eino
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/elastic/go-elasticsearch/v8"
|
||
|
||
"github.com/cloudwego/eino-ext/components/indexer/es8"
|
||
)
|
||
|
||
const (
|
||
indexName = "eino_example"
|
||
fieldContent = "content"
|
||
fieldContentVector = "content_vector"
|
||
fieldExtraLocation = "location"
|
||
docExtraLocation = "location"
|
||
)
|
||
|
||
func TestIndexer() {
|
||
ctx := context.Background()
|
||
|
||
// 1. 创建 ES 客户端
|
||
client, err := elasticsearch.NewClient(elasticsearch.Config{
|
||
Addresses: []string{"http://localhost:9200"},
|
||
})
|
||
|
||
if err != nil {
|
||
fmt.Printf("create client error: %v\n", err)
|
||
return
|
||
}
|
||
|
||
// 2. 定义 Index Spec（选填：如果索引不存在，将自动创建）
|
||
indexSpec := &es8.IndexSpec{
|
||
Settings: map[string]any{
|
||
"number_of_shards": 1,
|
||
"number_of_replicas": 0,
|
||
},
|
||
Mappings: map[string]any{
|
||
"properties": map[string]any{
|
||
fieldContentVector: map[string]any{
|
||
"type": "dense_vector",
|
||
"dims": 1024,
|
||
"index": true,
|
||
"similarity": "l2_norm",
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
// 4. 准备文档
|
||
// 文档通常包含 ID 和 Content
|
||
// 也可以包含额外的 Metadata 用于过滤或其他用途
|
||
docs := []*schema.Document{
|
||
{
|
||
ID: "1",
|
||
Content: "Eiffel Tower: Located in Paris, France.",
|
||
MetaData: map[string]any{
|
||
docExtraLocation: "France",
|
||
},
|
||
},
|
||
{
|
||
ID: "2",
|
||
Content: "The Great Wall: Located in China.",
|
||
MetaData: map[string]any{
|
||
docExtraLocation: "China",
|
||
},
|
||
},
|
||
}
|
||
|
||
// 5. 创建 ES 索引器组件
|
||
indexer, err := es8.NewIndexer(ctx, &es8.IndexerConfig{
|
||
Client: client,
|
||
Index: indexName,
|
||
IndexSpec: indexSpec, // 添加此项以启用自动索引创建
|
||
BatchSize: 10,
|
||
// DocumentToFields 指定如何将文档字段映射到 ES 字段
|
||
DocumentToFields: func(ctx context.Context, doc *schema.Document) (field2Value map[string]es8.FieldValue, err error) {
|
||
return map[string]es8.FieldValue{
|
||
fieldContent: {
|
||
Value: doc.Content,
|
||
EmbedKey: fieldContentVector, // 对文档内容进行向量化并保存到 "content_vector" 字段
|
||
},
|
||
fieldExtraLocation: {
|
||
// 额外的 metadata 字段
|
||
Value: doc.MetaData[docExtraLocation],
|
||
},
|
||
}, nil
|
||
},
|
||
// 提供 embedding 组件用于向量化
|
||
Embedding: EmbedderDashscope,
|
||
})
|
||
|
||
if err != nil {
|
||
fmt.Printf("create indexer error: %v\n", err)
|
||
return
|
||
}
|
||
|
||
// 6. 索引文档
|
||
ids, err := indexer.Store(ctx, docs)
|
||
if err != nil {
|
||
fmt.Printf("index error: %v\n", err)
|
||
return
|
||
}
|
||
fmt.Println("indexed ids:", ids)
|
||
}
|