Havoc412 91073fdf7f refactor(app): 重构 WebSocket 聊天流程并优化文档查询功能
- 调整了 ES TopK 查询逻辑,增加了文档类型筛选
- 优化了 WebSocket 的关闭流程,增加了文档和 token 信息的发送
- 新增了 Doc 模型的 ShowById 和 ShowByIds 方法,用于查询特定文档
- 更新了 prompts.yml 文件中的提示语,将"上下文"改为"知识库"
2024-11-19 13:06:39 +08:00

45 lines
852 B
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package curd
import (
"catface/app/model"
"catface/app/model_es"
"catface/app/model_res"
)
func CreateDocCurdFactory() *DocCurd {
return &DocCurd{
doc: model.CreateDocFactory(""),
doc_es: model_es.CreateDocESFactory()}
}
type DocCurd struct { // 组合方法的使用
doc *model.Doc
doc_es *model_es.Doc
}
func (d *DocCurd) TopK(embedding []float64, k int) (temp []model_res.DocResult, err error) {
// ESTopK
docs_es, err := d.doc_es.TopK(embedding, k)
if err != nil {
return
}
// MySQL补充基本信息
var ids []int64
for _, doc := range docs_es {
ids = append(ids, doc.Id)
}
docs := d.doc.ShowByIds(ids, "id", "name")
// 装载
for _, doc := range docs {
for _, doc_es := range docs_es {
if doc.Id == doc_es.Id {
temp = append(temp, *model_res.NewDocResult(&doc, &doc_es))
}
}
}
return
}