75 lines
1.7 KiB
Go
75 lines
1.7 KiB
Go
package model_es
|
|
|
|
import (
|
|
"catface/app/utils/data_bind"
|
|
"catface/app/utils/model_handler"
|
|
"encoding/json"
|
|
"fmt"
|
|
)
|
|
|
|
// INFO @brief 这个文件就是处理 ES 存储文档特征向量的集中处理；
|
|
|
|
func CreateDocESFactory() *Doc {
|
|
return &Doc{}
|
|
}
|
|
|
|
type Doc struct {
|
|
Id int64 `json:"id"` // 对应 MySQL 中存储的文档 ID。
|
|
Content string `json:"content"`
|
|
Embedding []float64 `json:"embedding"`
|
|
}
|
|
|
|
func (d *Doc) IndexName() string {
|
|
return "catface_docs"
|
|
}
|
|
|
|
func (d *Doc) TopK(embedding []float64, k int) ([]Doc, error) {
|
|
// 将 embedding 数组转换为 JSON 格式
|
|
params := map[string]interface{}{
|
|
"query_vector": embedding,
|
|
}
|
|
paramsJSON, err := json.Marshal(params)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 构建请求体
|
|
body := fmt.Sprintf(`{
|
|
"size": %d,
|
|
"query": {
|
|
"script_score": {
|
|
"query": {"match_all": {}},
|
|
"script": {
|
|
"source": "cosineSimilarity(params.query_vector, 'embedding') + 1.0",
|
|
"params": %s
|
|
}
|
|
}
|
|
},
|
|
"_source": ["content"]
|
|
}`, k, string(paramsJSON))
|
|
|
|
hits, err := model_handler.SearchRequest(body, d.IndexName())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var docs []Doc
|
|
for _, hit := range hits {
|
|
hitMap := hit.(map[string]interface{})
|
|
source := hitMap["_source"].(map[string]interface{})
|
|
var doc Doc
|
|
if err := data_bind.ShouldBindFormMapToModel(source, &doc); err != nil {
|
|
continue
|
|
}
|
|
|
|
docs = append(docs, doc)
|
|
}
|
|
|
|
return docs, nil
|
|
}
|
|
|
|
// UPDATE 因为 chunck 还是 Python 来处理会比较方便，所以 Go 这边主要还是处理查询相关的操作。
|
|
// func (d *Doc) InsertDocument() error {
|
|
// return nil
|
|
// }
|