79 lines
2.0 KiB
Go
79 lines
2.0 KiB
Go
package controller
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"time"
|
|
|
|
"rag-local/kb/model/domain"
|
|
"rag-local/kb/model/dto"
|
|
"rag-local/kb/service"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
)
|
|
|
|
type message struct{}
|
|
|
|
var Message = &message{}
|
|
|
|
func (c *message) List(ctx context.Context, req *dto.ListMessageReq) (*dto.ListMessageRes, error) {
|
|
list, err := service.MessageService.List(ctx, req.ConversationId)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.ListMessageRes{List: list}, nil
|
|
}
|
|
|
|
// Chat RAG 问答 SSE 流式：citations → delta* → done / error
|
|
func (c *message) Chat(ctx context.Context, req *dto.ChatReq) (*dto.ChatRes, error) {
|
|
r := g.RequestFromCtx(ctx)
|
|
r.Response.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
|
|
r.Response.Header().Set("Cache-Control", "no-cache")
|
|
r.Response.Header().Set("X-Accel-Buffering", "no")
|
|
r.Response.Header().Set("Connection", "keep-alive")
|
|
|
|
send := func(event string, data any) bool {
|
|
buf, err := json.Marshal(data)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
if event != "" {
|
|
r.Response.Write("event: " + event + "\n")
|
|
}
|
|
r.Response.Write("data: " + string(buf) + "\n\n")
|
|
r.Response.Flush()
|
|
return true
|
|
}
|
|
|
|
// 心跳：长回答期间保持连接，避免网关断流
|
|
stopHeartbeat := make(chan struct{})
|
|
defer close(stopHeartbeat)
|
|
go func() {
|
|
ticker := time.NewTicker(15 * time.Second)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-stopHeartbeat:
|
|
return
|
|
case <-ticker.C:
|
|
r.Response.Write(": ping\n\n")
|
|
r.Response.Flush()
|
|
}
|
|
}
|
|
}()
|
|
|
|
_, _, _, err := service.MessageService.Chat(ctx, req.ConversationId, req.DatasetId, req.Question,
|
|
func(citations []domain.Citation, conversationId int64) {
|
|
send("citations", map[string]any{"conversation_id": conversationId, "citations": citations})
|
|
},
|
|
func(delta string) {
|
|
send("delta", map[string]string{"content": delta})
|
|
})
|
|
if err != nil {
|
|
send("error", map[string]string{"message": err.Error()})
|
|
return nil, nil
|
|
}
|
|
send("done", map[string]string{"status": "ok"})
|
|
return nil, nil
|
|
}
|