145 lines
3.2 KiB
Go
145 lines
3.2 KiB
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
|
||
}
|