diff --git a/service/AIStreamService.go b/service/AIStreamService.go index 58d2755..77582d8 100644 --- a/service/AIStreamService.go +++ b/service/AIStreamService.go @@ -6,7 +6,7 @@ import ( "log" ) -func CreateAIStreamMsg(userID, modelID, sessionID int, msg interface{}, t string) { +func CreateAIStreamMsg(userID, modelID, sessionID int, msg interface{}, t string) (uint, error) { msgStr, err := json.Marshal(msg) if err != nil { log.Println("ai stream msg Error parsing JSON:", err) @@ -15,4 +15,5 @@ func CreateAIStreamMsg(userID, modelID, sessionID int, msg interface{}, t string if id == 0 { log.Println("ai stream msg insert error") } + return id, nil } diff --git a/service/doubao.go b/service/doubao.go index dd2e145..3df3cf4 100644 --- a/service/doubao.go +++ b/service/doubao.go @@ -34,6 +34,9 @@ func doubao(modelParam proto.ModelParam, question, channel string, SessionID, Us req := model.CreateChatCompletionRequest{ Model: modelParam.Model, Messages: messages, + StreamOptions: &model.StreamOptions{ + IncludeUsage: true, + }, } stream, err := doubaoClient.CreateChatCompletionStream(ctx, req) @@ -122,8 +125,15 @@ func DouBaoV2(modelParam proto.ModelParam, imCtx *proto.IMParamContext) { fmt.Printf("doubao Stream chat error: %v\n", err) return } - go CreateAIStreamMsg(imCtx.UserID, imCtx.ModelID, imCtx.SessionID, recv, modelParam.Model) + aiStream_id, err3 := CreateAIStreamMsg(imCtx.UserID, imCtx.ModelID, imCtx.SessionID, recv, modelParam.Model) + if err3 != nil { + log.Println("create ai stream message error:", err3) + } doubaoToGeneralMassageAndSendMsgQueue(&recv, imCtx.Channel, imCtx.SessionID, imCtx.UserID) + if recv.Choices == nil { + log.Println("doubao stream recv choices is nil:", recv, "\t aiStream_id:", aiStream_id) + continue + } choices := recv.Choices[0] if choices.FinishReason == proto.FinishReasonStop { answer += choices.Delta.Content