统一 token 和积分记录

hook/credit/credit.go

@@ -7,30 +7,35 @@ import (
 	"wechat-api/ent/creditbalance"
 	"wechat-api/ent/custom_types"
 	"wechat-api/ent/usagetotal"
-	"wechat-api/hook/dify"
 )
 
+type Usage struct {
+	PromptTokens     uint64 `json:"prompt_tokens"`
+	CompletionTokens uint64 `json:"completion_tokens"`
+	TotalTokens      uint64 `json:"total_tokens"`
+}
+
 func AddCreditUsage(tx *ent.Tx, ctx context.Context,
 	agentId string, userId string, departmentId uint64,
 	question *string, answer *string,
-	originalData *custom_types.OriginalData, chatData *dify.ChatResp) error {
+	originalData *custom_types.OriginalData, chatData *Usage) error {
 	// 积分明细表记录使用量
 	modelName, price := GetModelPrice()
-	number := ComputePrice(price, chatData.Metadata.Usage.TotalTokens)
+	number := ComputePrice(price, chatData.TotalTokens)
 
 	// 记录Token使用信息
 	usageDetailItem, err := tx.UsageDetail.Create().
-		SetType(3).            //1-微信 2-名片 3-智能体
-		SetBotID(agentId).     //智能体ID
+		SetType(3). //1-微信 2-名片 3-智能体
+		SetBotID(agentId). //智能体ID
 		SetReceiverID(userId). //接收者userID
-		SetApp(8).             //8-智能体
+		SetApp(8). //8-智能体
 		SetSessionID(0).
 		SetRequest(*question).
 		SetResponse(*answer).
 		SetOriginalData(*originalData).
-		SetTotalTokens(chatData.Metadata.Usage.TotalTokens).
-		SetPromptTokens(chatData.Metadata.Usage.PromptTokens).
-		SetCompletionTokens(chatData.Metadata.Usage.CompletionTokens).
+		SetTotalTokens(chatData.TotalTokens).
+		SetPromptTokens(chatData.PromptTokens).
+		SetCompletionTokens(chatData.CompletionTokens).
 		SetModel(modelName).
 		SetCredits(number).
 		Save(ctx)
@@ -46,7 +51,7 @@ func AddCreditUsage(tx *ent.Tx, ctx context.Context,
 		if ent.IsNotFound(err) {
 			usageTotal, err = tx.UsageTotal.Create().
 				SetBotID(agentId).
-				SetTotalTokens(chatData.Metadata.Usage.TotalTokens).
+				SetTotalTokens(chatData.TotalTokens).
 				SetEndIndex(usageDetailItem.ID).
 				SetOrganizationID(departmentId).
 				Save(ctx)
@@ -58,7 +63,7 @@ func AddCreditUsage(tx *ent.Tx, ctx context.Context,
 		// 更新Token使用总量
 		_, err = tx.UsageTotal.Update().
 			Where(usagetotal.OrganizationID(departmentId)).
-			SetTotalTokens(usageTotal.TotalTokens + chatData.Metadata.Usage.TotalTokens).
+			SetTotalTokens(usageTotal.TotalTokens + chatData.TotalTokens).
 			SetEndIndex(usageDetailItem.ID).
 			Save(ctx)
 		if err != nil {

internal/logic/chat/chat_completions_logic.go

@@ -10,8 +10,8 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
-
 	"wechat-api/ent"
+	"wechat-api/hook/credit"
 	"wechat-api/internal/svc"
 	"wechat-api/internal/types"
 	"wechat-api/internal/utils/compapi"
@@ -184,44 +184,76 @@ func (l *ChatCompletionsLogic) AppendAsyncRequest(apiKeyObj *ent.ApiKey, req *ty
 }
 
 func (l *ChatCompletionsLogic) AppendUsageDetailLog(authToken string, req *types.CompApiReq, resp *types.CompOpenApiResp) error {
-
-	logType := 5
 	rawReqResp := custom_types.OriginalData{Request: req, Response: resp}
-	tmpId := 0
-	tmpId, _ = strconv.Atoi(resp.ID)
-	sessionId := uint64(tmpId)
 	orgId := uint64(0)
 	apiKeyObj, ok := contextkey.AuthTokenInfoKey.GetValue(l.ctx)
 	if ok {
 		orgId = apiKeyObj.OrganizationID
 	}
-	promptTokens := uint64(resp.Usage.PromptTokens)
-	completionToken := uint64(resp.Usage.CompletionTokens)
-	totalTokens := promptTokens + completionToken
-
 	msgContent := getMessageContentStr(req.Messages[0].Content)
-
-	_, _, _ = logType, sessionId, totalTokens
-	res, err := l.svcCtx.DB.UsageDetail.Create().
-		SetNotNilType(&logType).
-		SetNotNilBotID(&authToken).
-		SetNotNilReceiverID(&req.EventType).
-		SetNotNilSessionID(&sessionId).
-		SetNillableRequest(&msgContent).
-		SetNillableResponse(&resp.Choices[0].Message.Content).
-		SetNillableOrganizationID(&orgId).
-		SetOriginalData(rawReqResp).
-		SetNillablePromptTokens(&promptTokens).
-		SetNillableCompletionTokens(&completionToken).
-		SetNillableTotalTokens(&totalTokens).
-		Save(l.ctx)
-
-	if err == nil { //插入UsageDetai之后再统计UsageTotal
-		l.updateUsageTotal(authToken, res.ID, orgId)
+	tx, err := l.svcCtx.DB.Tx(context.Background())
+	if err != nil {
+		l.Logger.Errorf("start transaction error:%v\n", err)
+	} else {
+		usage := credit.Usage{
+			CompletionTokens: uint64(resp.Usage.CompletionTokens),
+			PromptTokens:     uint64(resp.Usage.PromptTokens),
+			TotalTokens:      uint64(resp.Usage.TotalTokens),
+		}
+		err = credit.AddCreditUsage(tx, l.ctx,
+			authToken, req.EventType, orgId,
+			&msgContent, &resp.Choices[0].Message.Content,
+			&rawReqResp, &usage,
+		)
+		if err != nil {
+			_ = tx.Rollback()
+			l.Logger.Errorf("save credits info failed:%v\n", err)
+		} else {
+			_ = tx.Commit()
+		}
 	}
 	return err
 }
 
+//func (l *ChatCompletionsLogic) AppendUsageDetailLog(authToken string, req *types.CompApiReq, resp *types.CompOpenApiResp) error {
+//
+//	logType := 5
+//	rawReqResp := custom_types.OriginalData{Request: req, Response: resp}
+//	tmpId := 0
+//	tmpId, _ = strconv.Atoi(resp.ID)
+//	sessionId := uint64(tmpId)
+//	orgId := uint64(0)
+//	apiKeyObj, ok := contextkey.AuthTokenInfoKey.GetValue(l.ctx)
+//	if ok {
+//		orgId = apiKeyObj.OrganizationID
+//	}
+//	promptTokens := uint64(resp.Usage.PromptTokens)
+//	completionToken := uint64(resp.Usage.CompletionTokens)
+//	totalTokens := promptTokens + completionToken
+//
+//	msgContent := getMessageContentStr(req.Messages[0].Content)
+//
+//	_, _, _ = logType, sessionId, totalTokens
+//	res, err := l.svcCtx.DB.UsageDetail.Create().
+//		SetNotNilType(&logType).
+//		SetNotNilBotID(&authToken).
+//		SetNotNilReceiverID(&req.EventType).
+//		SetNotNilSessionID(&sessionId).
+//		SetNillableRequest(&msgContent).
+//		SetNillableResponse(&resp.Choices[0].Message.Content).
+//		SetNillableOrganizationID(&orgId).
+//		SetOriginalData(rawReqResp).
+//		SetNillablePromptTokens(&promptTokens).
+//		SetNillableCompletionTokens(&completionToken).
+//		SetNillableTotalTokens(&totalTokens).
+//		Save(l.ctx)
+//
+//	if err == nil { //插入UsageDetai之后再统计UsageTotal
+//		l.updateUsageTotal(authToken, res.ID, orgId)
+//	}
+//	return err
+//}
+
 func (l *ChatCompletionsLogic) getUsagetotalIdByToken(authToken string) (uint64, error) {
 
 	var predicates []predicate.UsageTotal

internal/logic/chatrecords/gpts_submit_api_chat_logic.go

@@ -204,14 +204,18 @@ func (l *GptsSubmitApiChatLogic) GptsSubmitApiChat(tokenStr string, req *types.G
 					originalData := custom_types.OriginalData{}
 					originalData.Request = chatReq
 					originalData.Response = chatData
-
+					usage := credit.Usage{
+						CompletionTokens: chatData.Metadata.Usage.CompletionTokens,
+						PromptTokens:     chatData.Metadata.Usage.PromptTokens,
+						TotalTokens:      chatData.Metadata.Usage.TotalTokens,
+					}
 					if err != nil {
 						l.Logger.Errorf("start transaction error:%v\n", err)
 					} else {
 						err = credit.AddCreditUsage(tx, l.ctx,
 							agentId, userId, *userInfo.DepartmentId,
 							req.Content, &answer,
-							&originalData, &chatData,
+							&originalData, &usage,
 						)
 						if err != nil {
 							_ = tx.Rollback()