From 4d974e4eabef93b90078b41e2cd8d95a26d55e55 Mon Sep 17 00:00:00 2001
From: junleea <354425203@qq.com>
Date: Tue, 1 Apr 2025 14:21:14 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0doubao=E6=97=A5=E5=BF=97?=
 =?UTF-8?q?=E8=BE=93=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 service/AIStreamService.go |  3 ++-
 service/doubao.go          | 12 +++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

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