fix:add switcher

internal/logic/chatrecords/gpts_submit_api_chat_logic.go

@@ -78,41 +78,45 @@ 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)
-	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 = "user"
+	if switcher {
+		if err != nil {
+			if !ent.IsNotFound(err) {
+				l.Logger.Errorf("query user balance failed:%v\n", err)
+				return
+			}
 		} 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
-				}
+			if creditBalance.Balance > 0 {
+				balance = creditBalance.Balance
+				balanceOwner = "user"
 			} else {
-				if creditBalance.Balance > 0 {
-					balance = creditBalance.Balance
-					balanceOwner = "organization"
+				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"
+					}
 				}
 			}
 		}
-	}
 
-	if balance <= 0 {
-		jsonData := ChatMessage{}
-		jsonData.NeedPay = true
-		jsonDataStr, _ := json.Marshal(jsonData)
-		_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
-		flusher.Flush()
-		return
+		if balance <= 0 {
+			jsonData := ChatMessage{}
+			jsonData.NeedPay = true
+			jsonDataStr, _ := json.Marshal(jsonData)
+			_, _ = fmt.Fprintf(w, "%s", "data: "+string(jsonDataStr)+"\r\n")
+			flusher.Flush()
+			return
+		}
 	}
 
 	employeeItem, err := l.svcCtx.DB.Employee.Query().Where(employee.ID(*req.AgentId)).Only(l.ctx)
@@ -206,57 +210,59 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 			flusher.Flush()
 
 			if finish {
-				// 构造 original_data
-				originalData := custom_types.OriginalData{}
-				originalData.Request = chatReq
-				originalData.Response = chatData
-
-				agentId := strconv.Itoa(int(*req.AgentId))
-				// 记录Token使用信息
-				usageDetailItem, err := l.svcCtx.DB.UsageDetail.Create().
-					SetType(3).            //1-微信 2-名片 3-智能体
-					SetBotID(agentId).     //智能体ID
-					SetReceiverID(userId). //接收者userID
-					SetApp(8).             //8-智能体
-					SetSessionID(0).
-					SetRequest(*req.Content).
-					SetResponse(answer).
-					SetOriginalData(originalData).
-					SetTotalTokens(chatData.Metadata.Usage.TotalTokens).
-					SetPromptTokens(chatData.Metadata.Usage.PromptTokens).
-					SetCompletionTokens(chatData.Metadata.Usage.CompletionTokens).
-					Save(l.ctx)
+				if switcher {
+					// 构造 original_data
+					originalData := custom_types.OriginalData{}
+					originalData.Request = chatReq
+					originalData.Response = chatData
 
-				if err != nil {
-					l.Logger.Errorf("save data to usage_detail error:%v\n", err)
-				}
+					agentId := strconv.Itoa(int(*req.AgentId))
+					// 记录Token使用信息
+					usageDetailItem, err := l.svcCtx.DB.UsageDetail.Create().
+						SetType(3).            //1-微信 2-名片 3-智能体
+						SetBotID(agentId).     //智能体ID
+						SetReceiverID(userId). //接收者userID
+						SetApp(8).             //8-智能体
+						SetSessionID(0).
+						SetRequest(*req.Content).
+						SetResponse(answer).
+						SetOriginalData(originalData).
+						SetTotalTokens(chatData.Metadata.Usage.TotalTokens).
+						SetPromptTokens(chatData.Metadata.Usage.PromptTokens).
+						SetCompletionTokens(chatData.Metadata.Usage.CompletionTokens).
+						Save(l.ctx)
 
-				// 积分明细表记录使用量
-				// 根据1:1000 根据Token换算积分使用量
-				change := float64(chatData.Metadata.Usage.TotalTokens) / float64(1000)
-				number := float32(math.Round(change*1000) / 1000)
-				_, err = l.svcCtx.DB.CreditUsage.Create().
-					SetUserID(userId).
-					SetNumber(number).
-					SetNtype(1).
-					SetNid(usageDetailItem.ID).
-					SetTable("usage_detail").
-					SetOrganizationID(*userInfo.DepartmentId).
-					Save(l.ctx)
-				if err != nil {
-					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)
+						l.Logger.Errorf("save data to usage_detail error:%v\n", err)
 					}
-				} else {
-					_, err := l.svcCtx.DB.CreditBalance.Update().Where(creditbalance.OrganizationID(*userInfo.DepartmentId)).SetBalance(creditBalance.Balance - number).Save(l.ctx)
+
+					// 积分明细表记录使用量
+					// 根据1:1000 根据Token换算积分使用量
+					change := float64(chatData.Metadata.Usage.TotalTokens) / float64(1000)
+					number := float32(math.Round(change*1000) / 1000)
+					_, err = l.svcCtx.DB.CreditUsage.Create().
+						SetUserID(userId).
+						SetNumber(number).
+						SetNtype(1).
+						SetNid(usageDetailItem.ID).
+						SetTable("usage_detail").
+						SetOrganizationID(*userInfo.DepartmentId).
+						Save(l.ctx)
 					if err != nil {
-						l.Logger.Errorf("update organization:%v balance error:%v\n", *userInfo.DepartmentId, err)
+						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)
+						}
 					}
 				}
 				break