catface_backend_go/app/model_es/knowledge.go

package model_es

import (
	"bytes"
	"catface/app/global/consts"
	"catface/app/global/variable"
	"catface/app/utils/data_bind"
	"catface/app/utils/model_handler"
	"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
}

/**
 * @description: 使用 ES 的 match 匹配虽有字段，同时处理 highlight
 * @param {string} query 查询字符串
 * @param {int} num 查询数量
 * @return {*}
 */
func (k *Knowledge) QueryDocumentsMatchAll(query string, num int) ([]Knowledge, error) {
	ctx := context.Background()

	body := fmt.Sprintf(`{
  "size": %d,
  "query": {
    "bool": {
      "should": [
        { "match": {"title": "%s" }},
        { "match": {"content": "%s" }}
      ]
    }
  },
 "highlight": {
    "pre_tags": ["%v"],
    "post_tags": ["%v"],
    "fields": {
      "title": {},
      "content": {
        "fragment_size" : 20
      }
    }
  }
}`, num, query, query, consts.PreTags, consts.PostTags) // TODO dirs 我还没想好如何处理
	// 创建请求
	req := esapi.SearchRequest{
		Index: []string{k.IndexName()},
		Body:  strings.NewReader(body),
	}

	// 发送请求
	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, ok := result["hits"].(map[string]interface{})["hits"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("error extracting hits from response")
	}

	var knowledges []Knowledge
	for _, hit := range hits {
		hitMap := hit.(map[string]interface{})
		source := hitMap["_source"].(map[string]interface{})
		highlight := hitMap["highlight"].(map[string]interface{})

		for k, v := range highlight {
			// INFO Knowledge 暂时不涉及 keywords 类型，就先这样处理。
			source[k] = model_handler.TransStringSliceToString(v.([]interface{}))
		}

		var k Knowledge
		if err := data_bind.ShouldBindFormMapToModel(source, &k); err != nil {
			return nil, err
		}
		knowledges = append(knowledges, k)
	}

	return knowledges, nil
}