catface_backend_go/app/model_es/knowledge.go

package model_es

import (
	"bytes"
	"catface/app/global/variable"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"strings"
	"time"

	"github.com/elastic/go-elasticsearch/v8/esapi"
)

func CreateKnowledgeESFactory() *Knowledge {
	return &Knowledge{}
}

type Knowledge struct {
	Id      int32    `json:"id"` // TIP int64 会炸 ES 的 ‘integer’
	Dirs    []string `json:"dirs"`
	Title   string   `json:"title"`
	Content string   `json:"content"`
}

func (k *Knowledge) IndexName() string {
	return "catface_knowledges"
}

func (k *Knowledge) InsertDocument() error {
	ctx := context.Background()

	k.Id = int32(time.Now().UnixNano() / 1e6) // 将纳秒级时间戳转换为毫秒级 // INFO 自动补充时间戳为 ID

	// 将结构体转换为 JSON 字符串
	data, err := json.Marshal(k)
	if err != nil {
		return err
	}

	// 创建请求
	req := esapi.IndexRequest{
		Index: k.IndexName(),
		// DocumentID: fmt.Sprintf("%d", k.Id),
		Body:    bytes.NewReader(data),
		Refresh: "true",
	}

	// 发送请求
	res, err := req.Do(ctx, variable.ElasticClient)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.IsError() {
		var k map[string]interface{}
		if err := json.NewDecoder(res.Body).Decode(&k); err != nil {
			return fmt.Errorf("error parsing the response body: %s", err)
		} else {
			return fmt.Errorf("[%s] %s: %s",
				res.Status(),
				k["error"].(map[string]interface{})["type"],
				k["error"].(map[string]interface{})["reason"],
			)
		}
	}

	return nil
}

// RandomDocuments 随机查询 num 个文档
func (k *Knowledge) RandomDocuments(num int) ([]*Knowledge, error) {
	ctx := context.Background()

	// 创建本地随机数生成器  // TIP rand.Seed() 在 Go1.20 之后弃用了。
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))

	// 创建请求
	req := esapi.SearchRequest{
		Index: []string{k.IndexName()},
		Body: strings.NewReader(fmt.Sprintf(`{
			"size": %d,
			"query": {
				"function_score": {
					"query": { "match_all": {} },
					"random_score": {
						"seed": %d
					}
				}
			}
		}`, num, rng.Int63())),
	}

	// 发送请求
	res, err := req.Do(ctx, variable.ElasticClient)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.IsError() {
		var k map[string]interface{}
		if err := json.NewDecoder(res.Body).Decode(&k); err != nil {
			return nil, fmt.Errorf("error parsing the response body: %s", err)
		} else {
			return nil, fmt.Errorf("[%s] %s: %s",
				res.Status(),
				k["error"].(map[string]interface{})["type"],
				k["error"].(map[string]interface{})["reason"],
			)
		}
	}

	// 解析响应
	var result map[string]interface{}
	if err := json.NewDecoder(res.Body).Decode(&result); err != nil {
		return nil, err
	}

	// 提取文档
	hits := result["hits"].(map[string]interface{})["hits"].([]interface{})

	documents := make([]*Knowledge, len(hits))
	for i, hit := range hits {
		source := hit.(map[string]interface{})["_source"].(map[string]interface{})

		// 将 []interface{} 转换为 []string
		dirs := make([]string, len(source["dirs"].([]interface{})))
		for j, dir := range source["dirs"].([]interface{}) {
			dirs[j] = dir.(string)
		}

		doc := &Knowledge{
			Dirs:    dirs,
			Title:   source["title"].(string),
			Content: source["content"].(string),
		}
		documents[i] = doc
	}

	return documents, nil
}