jimmyyem 2 달 전
부모
커밋
bb9036a8e0
1개의 변경된 파일12개의 추가작업 그리고 28개의 파일을 삭제
  1. 12 28
      internal/logic/chatrecords/gpts_submit_api_chat_logic.go

+ 12 - 28
internal/logic/chatrecords/gpts_submit_api_chat_logic.go

@@ -80,9 +80,8 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 	// 先查用户的积分余额,如果不足 则查租户的积分余额
 	var switcher bool
 	var balance float32
-	var balanceOwner string
 	//switcher = true	//TODO 打开这里就是开始计费,并且会记录使用积分和token的信息
-	creditBalance, err := l.svcCtx.DB.CreditBalance.Query().Where(creditbalance.UserID(userId)).Only(l.ctx)
+	creditBalance, err := l.svcCtx.DB.CreditBalance.Query().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).Only(l.ctx)
 	if switcher {
 		if err != nil {
 			if !ent.IsNotFound(err) {
@@ -92,20 +91,7 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 		} else {
 			if creditBalance.Balance > 0 {
 				balance = creditBalance.Balance
-				balanceOwner = "user"
-			} else {
-				creditBalance, err = l.svcCtx.DB.CreditBalance.Query().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).Only(l.ctx)
-				if err != nil {
-					if !ent.IsNotFound(err) {
-						l.Logger.Errorf("query user balance failed:%v\n", err)
-						return
-					}
-				} else {
-					if creditBalance.Balance > 0 {
-						balance = creditBalance.Balance
-						balanceOwner = "organization"
-					}
-				}
+
 			}
 		}
 
@@ -211,6 +197,7 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 
 			if finish {
 				if switcher {
+					tx, err := l.svcCtx.DB.Tx(context.Background())
 					// 构造 original_data
 					originalData := custom_types.OriginalData{}
 					originalData.Request = chatReq
@@ -218,7 +205,7 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 
 					agentId := strconv.Itoa(int(*req.AgentId))
 					// 记录Token使用信息
-					usageDetailItem, err := l.svcCtx.DB.UsageDetail.Create().
+					usageDetailItem, err := tx.UsageDetail.Create().
 						SetType(3).            //1-微信 2-名片 3-智能体
 						SetBotID(agentId).     //智能体ID
 						SetReceiverID(userId). //接收者userID
@@ -233,6 +220,7 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 						Save(l.ctx)
 
 					if err != nil {
+						_ = tx.Rollback()
 						l.Logger.Errorf("save data to usage_detail error:%v\n", err)
 					}
 
@@ -240,7 +228,7 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 					// 根据1:1000 根据Token换算积分使用量
 					change := float64(chatData.Metadata.Usage.TotalTokens) / float64(1000)
 					number := float32(math.Round(change*1000) / 1000)
-					_, err = l.svcCtx.DB.CreditUsage.Create().
+					_, err = tx.CreditUsage.Create().
 						SetUserID(userId).
 						SetNumber(number).
 						SetNtype(1).
@@ -249,21 +237,17 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 						SetOrganizationID(*userInfo.DepartmentId).
 						Save(l.ctx)
 					if err != nil {
+						_ = tx.Rollback()
 						l.Logger.Errorf("save data to credit_usage error:%v\n", err)
 					}
 
 					// 减积分
-					if balanceOwner == "user" {
-						_, err := l.svcCtx.DB.CreditBalance.Update().Where(creditbalance.UserID(userId)).SetBalance(creditBalance.Balance - number).Save(l.ctx)
-						if err != nil {
-							l.Logger.Errorf("update user:%v balance error:%v\n", userId, err)
-						}
-					} else {
-						_, err := l.svcCtx.DB.CreditBalance.Update().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).SetBalance(creditBalance.Balance - number).Save(l.ctx)
-						if err != nil {
-							l.Logger.Errorf("update organization:%v balance error:%v\n", *userInfo.DepartmentId, err)
-						}
+					_, err = tx.CreditBalance.Update().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).SetBalance(creditBalance.Balance - number).Save(l.ctx)
+					if err != nil {
+						_ = tx.Rollback()
+						l.Logger.Errorf("update organization:%v balance error:%v\n", *userInfo.DepartmentId, err)
 					}
+					_ = tx.Commit()
 				}
 				break
 			}