84 lines
2.4 KiB
Go
84 lines
2.4 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
|
|
"rag-local/kb/dao"
|
|
"rag-local/kb/model/domain"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/cloudwego/eino/schema"
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
)
|
|
|
|
var MessageService = &messageService{}
|
|
|
|
type messageService struct{}
|
|
|
|
func (s *messageService) List(ctx context.Context, conversationId int64) ([]*entity.Message, error) {
|
|
return dao.Message.List(ctx, conversationId)
|
|
}
|
|
|
|
// Chat RAG 问答：会话解析 → 用户消息落库 → 工作流流式生成 → 助手消息+引用落库。
|
|
// onCitations 在检索完成后回调（先于流式输出）；onDelta 接收模型增量文本。
|
|
func (s *messageService) Chat(ctx context.Context, conversationId, datasetId int64, question string, onCitations func([]domain.Citation, int64), onDelta func(string)) (string, []domain.Citation, int64, error) {
|
|
if conversationId <= 0 {
|
|
title := question
|
|
if r := []rune(title); len(r) > 20 {
|
|
title = string(r[:20])
|
|
}
|
|
id, err := dao.Conversation.Insert(ctx, datasetId, title)
|
|
if err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
conversationId = id
|
|
}
|
|
conv, err := dao.Conversation.GetOne(ctx, conversationId)
|
|
if err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
if conv == nil {
|
|
return "", nil, 0, gerror.New("会话不存在")
|
|
}
|
|
if conv.DatasetId == 0 && datasetId > 0 {
|
|
_ = dao.Conversation.UpdateDataset(ctx, conversationId, datasetId)
|
|
conv.DatasetId = datasetId
|
|
}
|
|
if conv.DatasetId > 0 {
|
|
datasetId = conv.DatasetId
|
|
}
|
|
if datasetId <= 0 {
|
|
return "", nil, 0, gerror.New("请选择知识库数据集")
|
|
}
|
|
|
|
if _, err := dao.Message.Insert(ctx, conversationId, "user", question, ""); err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
|
|
history, err := dao.Message.List(ctx, conversationId)
|
|
if err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
messages := make([]*schema.Message, 0, len(history))
|
|
for _, m := range history {
|
|
messages = append(messages, &schema.Message{Role: schema.RoleType(m.Role), Content: m.Content})
|
|
}
|
|
|
|
answer, citations, err := ChatService.Ask(ctx, datasetId, question, messages,
|
|
func(c []domain.Citation) {
|
|
if onCitations != nil {
|
|
onCitations(c, conversationId)
|
|
}
|
|
}, onDelta)
|
|
if err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
|
|
citJson, _ := json.Marshal(citations)
|
|
if _, err := dao.Message.Insert(ctx, conversationId, "assistant", answer, string(citJson)); err != nil {
|
|
return "", nil, 0, err
|
|
}
|
|
return answer, citations, conversationId, nil
|
|
}
|