写入与查询自动Embedding

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"log"
	"net/http"
	"strings"
	"time"
)

var client *elasticsearch.Client

func init() {
	cfg := elasticsearch.Config{
		Addresses: []string{
			"http://ld-t4n5668xk31ui****-proxy-search-pub.lindorm.aliyuncs.com:30070",
		},
		Username: "<Username>",
		Password: "<Password>",
		Transport: &http.Transport{
			MaxIdleConnsPerHost: 20,
		},
	}
	var err error
	client, err = elasticsearch.NewClient(cfg)
	if err != nil {
		fmt.Println("Init client error %s", err)
	}
}

func handlerCommonError(res *esapi.Response) {
	var errorResponse map[string]interface{}
	if err := json.NewDecoder(res.Body).Decode(&errorResponse); err != nil {
		fmt.Printf("Error parsing the error response: %s\n", err)
		return
	}
	fmt.Printf("ERROR %v\n", errorResponse)
}

type Source struct {
	Field1 int      `json:"field1"`
	Field2 string   `json:"field2"`
	Brand  string   `json:"brand"`
	Merit  string   `json:"merit"`
	Tag    []string `json:"tag"`
}

type KnnResponse struct {
	Hits struct {
		Hits []struct {
			ID     string  `json:"_id"`
			Score  float64 `json:"_score"`
			Source Source  `json:"_source"`
		} `json:"hits"`
	} `json:"hits"`
}

func createWritePipeline() {
        password := "test****" //AI引擎的密码
	wirtePipelineID := "write_embedding_pipeline"
	url := "http://ld-2zeyhrpnv9n7a****-proxy-ai-vpc.lindorm.aliyuncs.com:9002" //AI引擎的专有网络连接地址
	writePipeline := map[string]interface{}{
		"description": "demo_chunking pipeline",
		"processors": []interface{}{
			map[string]interface{}{
				"text-embedding": map[string]interface{}{
					"inputFields":  []interface{}{"text_field"},
					"outputFields": []interface{}{"text_field_embedding"},
					"userName":     "root", //AI引擎的用户名
					"password":     "test****", //AI引擎的密码
					"url":          "http://ld-2zeyhrpnv9n7a****-proxy-ai-vpc.lindorm.aliyuncs.com:9002", // AI引擎的专有网络连接地址
					"modeName":     "bge_m3_model",
				},
			},
		},
	}

	content, err := json.Marshal(writePipeline)
	if err != nil {
		fmt.Printf("Error marshalling JSON: %s\n", err)
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	request := esapi.IngestPutPipelineRequest{
		PipelineID: wirtePipelineID,
		Body:       bytes.NewReader(content),
	}

	response, err := request.Do(ctx, client)
	if err != nil {
		fmt.Println("Knn err %s", err)
		return
	}
	defer response.Body.Close()

	if response.IsError() {
		handlerCommonError(response)
		return
	}
	fmt.Println("Response: ", response.String())
}

func createSearchPipeline() {
        password := "test****" //AI引擎的密码
	searchPipelineID := "knnsearch_pipeline"
	url := "http://ld-2zeyhrpnv9n7a****-proxy-ai-vpc.lindorm.aliyuncs.com:9002" //AI引擎的专有网络连接地址
	searchPipeline := map[string]interface{}{
		"request_processors": []interface{}{
			map[string]interface{}{
				"text-embedding": map[string]interface{}{
					"tag":         "auto-query-embedding",
					"description": "Auto query embedding",
					"model_config": map[string]interface{}{
						"inputFields":  []interface{}{"text_field"},
						"outputFields": []interface{}{"text_field_embedding"},
						"userName":     "root", //AI引擎的用户名
						"password":     "test****", //AI引擎的密码
						"url":          "http://ld-2zeyhrpnv9n7a****-proxy-ai-vpc.lindorm.aliyuncs.com:9002", //AI引擎的专有网络连接地址 
						"modeName":     "bge_m3_model",
					},
				},
			},
		},
	}

	searchPielineUrl := fmt.Sprintf("_search/pipeline/%s", searchPipelineID)

	buf := new(bytes.Buffer)
	if err := json.NewEncoder(buf).Encode(searchPipeline); err != nil {
		fmt.Println("Encode error ", err)
		return
	}
	request, err := http.NewRequest("PUT", searchPielineUrl, buf)
	if err != nil {
		fmt.Println("Err ", err)
		return
	}
	request.Header.Set("Accept", "application/json")
	request.Header.Set("Content-Type", "application/json")
	res, err := client.Transport.Perform(request)
	if err != nil {
		fmt.Println("Err ", err)
		return
	}
	defer res.Body.Close()
	response := esapi.Response{
		StatusCode: res.StatusCode,
		Body:       res.Body,
		Header:     res.Header,
	}

	fmt.Println("Response:", response.String())
}

// 创建索引，指定写入与查询pipeline
func createAutoEmbeddingPipelineIndex() {
	indexName := "search_vector_test"
	writePipelineID := "write_embedding_pipeline"
	searchPipelineID := "knnsearch_pipeline"
	createMapping := map[string]interface{}{
		"settings": map[string]interface{}{
			"index": map[string]interface{}{
				"number_of_shards":        2,
				"knn":                     true,
				"default_pipeline":        writePipelineID,
				"search.default_pipeline": searchPipelineID,
			},
		},
		"mappings": map[string]interface{}{
			"_source": map[string]interface{}{
				"excludes": []interface{}{"text_field_embedding"},
			},
			"properties": map[string]interface{}{
				"text_field": map[string]interface{}{
					"type":     "text",
					"analyzer": "ik_max_word",
				},
				"text_field_embedding": map[string]interface{}{
					"type":      "knn_vector",
					"dimension": 1024,
					"data_type": "float",
					"method": map[string]interface{}{
						"engine":     "lvector",
						"name":       "hnsw",
						"space_type": "cosinesimil",
						"parameters": map[string]interface{}{
							"m":               24,
							"ef_construction": 500,
						},
					},
				},
				"tag": map[string]interface{}{
					"type": "keyword",
				},
				"brand": map[string]interface{}{
					"type": "keyword",
				},
				"merit": map[string]interface{}{
					"type":     "text",
					"analyzer": "ik_max_word",
				},
			},
		},
	}

	body, err := json.Marshal(createMapping)
	if err != nil {
		log.Fatalf("Error marshaling index body: %s", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	req := esapi.IndicesCreateRequest{
		Index: indexName,
		Body:  bytes.NewReader(body),
	}

	res, err := req.Do(ctx, client)
	defer res.Body.Close()
	if err != nil {
		fmt.Println("Create index error: ", err)
		return
	}
	if res.IsError() {
		handlerCommonError(res)
		return
	}
	fmt.Println("Create response: ", res.String())
}

func updateIndexAddAutoEmbeddingPipeline() {
	// 假设 search_vector_test2 没有指定 default_pipeline 与 search.default_pipeline
	indexName := "search_vector_test2"
	putSetting := map[string]interface{}{
		"index": map[string]interface{}{
			"default_pipeline":        "write_embedding_pipeline",
			"search.default_pipeline": "knnsearch_pipeline",
		},
	}

	body, err := json.Marshal(putSetting)
	if err != nil {
		log.Fatalf("Error marshaling index body: %s", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()

	req := esapi.IndicesPutSettingsRequest{
		Index: []string{indexName},
		Body:  bytes.NewReader(body),
	}

	res, err := req.Do(ctx, client)
	defer res.Body.Close()
	if err != nil {
		fmt.Println("Create index error: ", err)
		return
	}

	if res.IsError() {
		handlerCommonError(res)
		return
	}
	fmt.Println("Put setting: ", res.String())
}