455 lines
12 KiB
Go
455 lines
12 KiB
Go
package provider
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"io"
|
||
|
||
"dubbo.apache.org/dubbo-go/v3/common/constant"
|
||
"github.com/topfans/backend/pkg/logger"
|
||
"github.com/topfans/backend/services/aiChatService/model"
|
||
"github.com/topfans/backend/services/aiChatService/service"
|
||
pb "github.com/topfans/backend/pkg/proto/ai_chat"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// memoryExtractionInterval 记忆提取的触发轮数间隔。
|
||
// 每 N 轮对话触发一次 LLM 记忆提取（第 N, 2N, 3N... 轮）。
|
||
const memoryExtractionInterval = 5
|
||
|
||
// AIChatProvider AI Chat 服务 Provider 实现
|
||
type AIChatProvider struct {
|
||
chatService *service.ChatService
|
||
personaService *service.PersonaService
|
||
memoryService *service.MemoryService
|
||
auditService *service.AuditService
|
||
}
|
||
|
||
// 确保 AIChatProvider 实现了 AIChatServiceHandler 接口
|
||
var _ pb.AIChatServiceHandler = (*AIChatProvider)(nil)
|
||
|
||
// NewAIChatProvider 创建 AIChatProvider 实例
|
||
func NewAIChatProvider(
|
||
chatService *service.ChatService,
|
||
personaService *service.PersonaService,
|
||
memoryService *service.MemoryService,
|
||
auditService *service.AuditService,
|
||
) *AIChatProvider {
|
||
return &AIChatProvider{
|
||
chatService: chatService,
|
||
personaService: personaService,
|
||
memoryService: memoryService,
|
||
auditService: auditService,
|
||
}
|
||
}
|
||
|
||
// InitSession 初始化会话，返回欢迎消息（同时创建当前明星的默认人设）
|
||
func (p *AIChatProvider) InitSession(ctx context.Context, req *pb.InitSessionRequest) (*pb.InitSessionResponse, error) {
|
||
userID, starID, err := extractUserInfoFromDubboAttachments(ctx)
|
||
if err != nil {
|
||
logger.Logger.Error("Failed to extract user info from attachments",
|
||
zap.Error(err),
|
||
)
|
||
return nil, err
|
||
}
|
||
|
||
sessionID := req.SessionId
|
||
if sessionID == "" {
|
||
sessionID = fmt.Sprintf("%d_%d", userID, starID)
|
||
}
|
||
|
||
logger.Logger.Info("Received InitSession request",
|
||
zap.Int64("user_id", userID),
|
||
zap.Int64("star_id", starID),
|
||
zap.String("session_id", sessionID),
|
||
)
|
||
|
||
// 进入聊天即创建当前明星的默认人设，防止出现 A 明星用 B 人设的情况
|
||
if _, err := p.personaService.GetPersonaOrDefault(ctx, userID, starID, ""); err != nil {
|
||
logger.Logger.Warn("Failed to ensure default persona in InitSession",
|
||
zap.Int64("user_id", userID),
|
||
zap.Int64("star_id", starID),
|
||
zap.Error(err),
|
||
)
|
||
// 不阻塞进入聊天，SendMessage 时会再次尝试
|
||
}
|
||
|
||
// 获取欢迎消息
|
||
welcomeMessage := p.chatService.GetWelcomeMessage(sessionID, userID, starID)
|
||
|
||
return &pb.InitSessionResponse{
|
||
WelcomeMessage: welcomeMessage,
|
||
SessionId: sessionID,
|
||
}, nil
|
||
}
|
||
|
||
// SendMessage 发送消息（流式返回）
|
||
func (p *AIChatProvider) SendMessage(ctx context.Context, req *pb.ChatMessageRequest, stream pb.AIChatService_SendMessageServer) error {
|
||
userID, starID, err := extractUserInfoFromDubboAttachments(ctx)
|
||
if err != nil {
|
||
logger.Logger.Error("Failed to extract user info from attachments",
|
||
zap.Error(err),
|
||
)
|
||
return err
|
||
}
|
||
|
||
sessionID := req.SessionId
|
||
if sessionID == "" {
|
||
sessionID = fmt.Sprintf("%d_%d", userID, starID)
|
||
}
|
||
|
||
message := req.Message
|
||
personaID := req.PersonaId
|
||
|
||
logger.Logger.Info("Received SendMessage request",
|
||
zap.Int64("user_id", userID),
|
||
zap.String("session_id", sessionID),
|
||
zap.String("message", message),
|
||
)
|
||
|
||
// 1. 前置审核
|
||
if !p.auditService.AuditText(message) {
|
||
logger.Logger.Info("Message blocked by audit")
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: p.auditService.DefaultSafeResponse(),
|
||
SessionId: sessionID,
|
||
IsEnd: false,
|
||
})
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
})
|
||
return nil
|
||
}
|
||
|
||
// 2. 获取人设（传入 starID 用于首次使用时创建明星专属人设）
|
||
persona, err := p.personaService.GetPersonaOrDefault(ctx, userID, starID, personaID)
|
||
if err != nil {
|
||
logger.Logger.Error("Failed to get persona", zap.Error(err))
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: err.Error(),
|
||
IsEnd: true,
|
||
Error: err.Error(),
|
||
})
|
||
return err
|
||
}
|
||
|
||
// 3. 记忆召回
|
||
memoryText, _ := p.memoryService.RecallMemories(ctx, userID, message, 5)
|
||
|
||
// 4. 获取对话历史
|
||
history, _ := p.memoryService.GetContext(ctx, sessionID)
|
||
|
||
// 5. 构建 Prompt
|
||
tokenizer := &service.Tokenizer{}
|
||
messages, _ := service.BuildPrompt(
|
||
persona.SystemPrompt,
|
||
memoryText,
|
||
history,
|
||
message,
|
||
tokenizer,
|
||
)
|
||
|
||
// 6. 检查是否需要调用大模型
|
||
if service.IsNoNeedLLMCall(message) {
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: "好的，我听到了。",
|
||
SessionId: sessionID,
|
||
IsEnd: false,
|
||
})
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
})
|
||
return nil
|
||
}
|
||
|
||
// 7. 调用大模型（流式）
|
||
streamReader, err := p.chatService.StreamChat(ctx, messages)
|
||
if err != nil {
|
||
logger.Logger.Error("AI call failed", zap.Error(err))
|
||
|
||
// 检查是否是敏感内容错误
|
||
if _, ok := err.(*service.SensitiveContentError); ok {
|
||
logger.Logger.Info("Content blocked by safety filter")
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: p.auditService.DefaultSafeResponse(),
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
})
|
||
return nil
|
||
}
|
||
|
||
// 其他错误
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: "抱歉，服务暂时不可用",
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
Error: err.Error(),
|
||
})
|
||
return err
|
||
}
|
||
defer streamReader.Close()
|
||
|
||
// 8. 流式处理
|
||
var fullResponse string
|
||
var sentEnd = false
|
||
for {
|
||
content, done, err := streamReader.Next()
|
||
if err != nil {
|
||
if err == io.EOF {
|
||
// 流结束，发送 is_end
|
||
if !sentEnd {
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
})
|
||
sentEnd = true
|
||
}
|
||
break
|
||
}
|
||
logger.Logger.Error("Stream read error", zap.Error(err))
|
||
break
|
||
}
|
||
|
||
// 后置审核（逐 token）
|
||
if !p.auditService.AuditResponse(content) {
|
||
logger.Logger.Info("Response blocked by audit")
|
||
streamReader.Close()
|
||
// 发送安全回复作为替代
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
Content: p.auditService.DefaultSafeResponse(),
|
||
SessionId: sessionID,
|
||
IsEnd: false,
|
||
})
|
||
stream.Send(&pb.ChatMessageResponse{
|
||
SessionId: sessionID,
|
||
IsEnd: true,
|
||
})
|
||
sentEnd = true
|
||
return nil
|
||
}
|
||
|
||
fullResponse += content
|
||
|
||
// 发送 token 给客户端
|
||
if err := stream.Send(&pb.ChatMessageResponse{
|
||
Content: content,
|
||
SessionId: sessionID,
|
||
IsEnd: done,
|
||
}); err != nil {
|
||
logger.Logger.Error("Failed to send message to stream", zap.Error(err))
|
||
return err
|
||
}
|
||
if done {
|
||
sentEnd = true
|
||
}
|
||
}
|
||
|
||
// 9. 保存上下文
|
||
newHistory := append(history, model.Message{Role: "user", Content: message})
|
||
newHistory = append(newHistory, model.Message{Role: "assistant", Content: fullResponse})
|
||
p.memoryService.SaveContext(ctx, sessionID, newHistory, personaID)
|
||
|
||
// 10. 触发记忆提取（每5轮触发一次，在第5、10、15...轮提取）
|
||
newTurns := len(newHistory) / 2
|
||
shouldExtract := newTurns >= memoryExtractionInterval && newTurns%memoryExtractionInterval == 0
|
||
logger.Logger.Info("Memory extraction check",
|
||
zap.Int("message_count", len(newHistory)),
|
||
zap.Int("turns", newTurns),
|
||
zap.Bool("should_extract", shouldExtract),
|
||
)
|
||
if shouldExtract {
|
||
logger.Logger.Info("Triggering LLM memory extraction", zap.Int64("user_id", userID))
|
||
// 使用 LLM 提取记忆
|
||
extractedMemories, err := service.ExtractMemoriesWithLLM(ctx, p.chatService.Chat, newHistory)
|
||
if err != nil {
|
||
logger.Logger.Error("LLM memory extraction failed", zap.Error(err))
|
||
} else if len(extractedMemories) > 0 {
|
||
if err := p.memoryService.ExtractMemory(ctx, userID, newHistory, extractedMemories); err != nil {
|
||
logger.Logger.Error("Failed to save extracted memories", zap.Error(err))
|
||
} else {
|
||
logger.Logger.Info("Memories extracted and saved successfully",
|
||
zap.Int64("user_id", userID),
|
||
zap.Int("count", len(extractedMemories)),
|
||
)
|
||
}
|
||
}
|
||
}
|
||
|
||
logger.Logger.Info("SendMessage completed",
|
||
zap.Int64("user_id", userID),
|
||
zap.String("session_id", sessionID),
|
||
zap.Int("response_length", len(fullResponse)),
|
||
)
|
||
|
||
return nil
|
||
}
|
||
|
||
// GetHistory 获取对话历史
|
||
func (p *AIChatProvider) GetHistory(ctx context.Context, req *pb.ChatHistoryRequest) (*pb.ChatHistoryResponse, error) {
|
||
userID, starID, err := extractUserInfoFromDubboAttachments(ctx)
|
||
if err != nil {
|
||
logger.Logger.Error("Failed to extract user info from attachments",
|
||
zap.Error(err),
|
||
)
|
||
return nil, err
|
||
}
|
||
|
||
sessionID := req.SessionId
|
||
if sessionID == "" {
|
||
sessionID = fmt.Sprintf("%d_%d", userID, starID)
|
||
}
|
||
|
||
logger.Logger.Info("Received GetHistory request",
|
||
zap.Int64("user_id", userID),
|
||
zap.String("session_id", sessionID),
|
||
)
|
||
|
||
messages, err := p.memoryService.GetContext(ctx, sessionID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
pbMessages := make([]*pb.Message, len(messages))
|
||
for i, m := range messages {
|
||
pbMessages[i] = &pb.Message{
|
||
Role: m.Role,
|
||
Content: m.Content,
|
||
}
|
||
}
|
||
|
||
return &pb.ChatHistoryResponse{
|
||
History: pbMessages,
|
||
}, nil
|
||
}
|
||
|
||
// GetPersonas 获取用户的所有人设
|
||
func (p *AIChatProvider) GetPersonas(ctx context.Context, req *pb.GetPersonasRequest) (*pb.PersonaListResponse, error) {
|
||
userID, _, err := extractUserInfoFromDubboAttachments(ctx)
|
||
if err != nil {
|
||
logger.Logger.Error("Failed to extract user info from attachments",
|
||
zap.Error(err),
|
||
)
|
||
return nil, err
|
||
}
|
||
|
||
logger.Logger.Info("Received GetPersonas request",
|
||
zap.Int64("user_id", userID),
|
||
)
|
||
|
||
personas, err := p.personaService.GetPersonas(ctx, userID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
pbPersonas := make([]*pb.PersonaInfo, len(personas))
|
||
for i, persona := range personas {
|
||
pbPersonas[i] = &pb.PersonaInfo{
|
||
Id: persona.ID,
|
||
StarId: persona.StarID,
|
||
Name: persona.Name,
|
||
Description: persona.Description,
|
||
AvatarUrl: persona.AvatarURL,
|
||
TalkStyle: persona.TalkStyle,
|
||
IsDefault: persona.IsDefault,
|
||
CreatedAt: persona.CreatedAt,
|
||
UpdatedAt: persona.UpdatedAt,
|
||
}
|
||
}
|
||
|
||
return &pb.PersonaListResponse{
|
||
Personas: pbPersonas,
|
||
}, nil
|
||
}
|
||
|
||
// extractUserInfoFromDubboAttachments 从 Dubbo attachments 中提取用户信息
|
||
func extractUserInfoFromDubboAttachments(ctx context.Context) (int64, int64, error) {
|
||
logger.Logger.Debug("Extracting user info from Dubbo attachments",
|
||
zap.Any("context_type", fmt.Sprintf("%T", ctx)),
|
||
)
|
||
|
||
// Try to get any value from context
|
||
if attachments := ctx.Value(constant.AttachmentKey); attachments != nil {
|
||
logger.Logger.Debug("Found attachments via constant.AttachmentKey",
|
||
zap.Any("attachments", attachments),
|
||
zap.String("type", fmt.Sprintf("%T", attachments)),
|
||
)
|
||
if attMap, ok := attachments.(map[string]interface{}); ok {
|
||
logger.Logger.Debug("Attachments map content",
|
||
zap.Any("map", attMap),
|
||
)
|
||
userID := parseIntValue(attMap["user_id"])
|
||
starID := parseIntValue(attMap["star_id"])
|
||
logger.Logger.Debug("Parsed user info from attachments",
|
||
zap.Any("user_id_raw", attMap["user_id"]),
|
||
zap.Int64("user_id", userID),
|
||
zap.Any("star_id_raw", attMap["star_id"]),
|
||
zap.Int64("star_id", starID),
|
||
)
|
||
if userID > 0 && starID > 0 {
|
||
return userID, starID, nil
|
||
}
|
||
logger.Logger.Warn("Parsed user_id or star_id is zero",
|
||
zap.Int64("user_id", userID),
|
||
zap.Int64("star_id", starID),
|
||
)
|
||
} else {
|
||
logger.Logger.Warn("Attachments is not map[string]interface{}",
|
||
zap.String("actual_type", fmt.Sprintf("%T", attachments)),
|
||
)
|
||
}
|
||
} else {
|
||
logger.Logger.Warn("ctx.Value(constant.AttachmentKey) returned nil",
|
||
zap.String("constant_attachment_key", string(constant.AttachmentKey)),
|
||
)
|
||
}
|
||
|
||
// Debug: list all keys in context
|
||
logger.Logger.Warn("Checking alternative key: 'attachment'")
|
||
if val := ctx.Value("attachment"); val != nil {
|
||
logger.Logger.Debug("Found value with key 'attachment'",
|
||
zap.Any("value", val),
|
||
zap.String("type", fmt.Sprintf("%T", val)),
|
||
)
|
||
}
|
||
|
||
return 0, 0, fmt.Errorf("user info not found in Dubbo attachments")
|
||
}
|
||
|
||
// parseIntValue 解析各种类型的值为 int64
|
||
func parseIntValue(v interface{}) int64 {
|
||
switch val := v.(type) {
|
||
case int64:
|
||
return val
|
||
case int:
|
||
return int64(val)
|
||
case float64:
|
||
return int64(val)
|
||
case string:
|
||
var result int64
|
||
fmt.Sscanf(val, "%d", &result)
|
||
return result
|
||
case []string:
|
||
if len(val) > 0 {
|
||
var result int64
|
||
fmt.Sscanf(val[0], "%d", &result)
|
||
return result
|
||
}
|
||
case []interface{}:
|
||
if len(val) > 0 {
|
||
switch s := val[0].(type) {
|
||
case string:
|
||
var result int64
|
||
fmt.Sscanf(s, "%d", &result)
|
||
return result
|
||
case int:
|
||
return int64(s)
|
||
case int64:
|
||
return s
|
||
}
|
||
}
|
||
}
|
||
return 0
|
||
}
|