wechat-api/hook/credit/credit.go

package credit

import (
	"context"
	"fmt"
	"wechat-api/ent"
	"wechat-api/ent/creditbalance"
	"wechat-api/ent/custom_types"
	"wechat-api/ent/usagetotal"
)

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 *Usage) error {
	// 积分明细表记录使用量
	modelName, price := GetModelPrice()
	number := ComputePrice(price, chatData.TotalTokens)

	// 记录Token使用信息
	usageDetailItem, err := tx.UsageDetail.Create().
		SetType(3).            //1-微信 2-名片 3-智能体
		SetBotID(agentId).     //智能体ID
		SetReceiverID(userId). //接收者userID
		SetApp(8).             //8-智能体
		SetSessionID(0).
		SetRequest(*question).
		SetResponse(*answer).
		SetOriginalData(*originalData).
		SetTotalTokens(chatData.TotalTokens).
		SetPromptTokens(chatData.PromptTokens).
		SetCompletionTokens(chatData.CompletionTokens).
		SetModel(modelName).
		SetCredits(number).
		Save(ctx)

	if err != nil {
		fmt.Printf("create usage_detail failed:%v\n", err)
		return err
	}

	// 记录Token使用总量
	usageTotal, err := tx.UsageTotal.Query().Where(usagetotal.OrganizationID(departmentId)).First(ctx)
	if err != nil {
		if ent.IsNotFound(err) {
			usageTotal, err = tx.UsageTotal.Create().
				SetBotID(agentId).
				SetTotalTokens(chatData.TotalTokens).
				SetEndIndex(usageDetailItem.ID).
				SetOrganizationID(departmentId).
				Save(ctx)
		} else {
			fmt.Printf("create usage_total failed:organization_id:%v err:%v\n", departmentId, err)
			return err
		}
	} else {
		// 更新Token使用总量
		_, err = tx.UsageTotal.Update().
			Where(usagetotal.OrganizationID(departmentId)).
			SetTotalTokens(usageTotal.TotalTokens + chatData.TotalTokens).
			SetEndIndex(usageDetailItem.ID).
			Save(ctx)
		if err != nil {
			fmt.Printf("update usage_total failed:organization_id:%v err:%v\n", departmentId, err)
			return err
		}
	}

	creditBalanceItem, err := tx.CreditBalance.Query().Where(creditbalance.OrganizationID(departmentId)).First(ctx)
	if err != nil {
		if ent.IsNotFound(err) {
			creditBalanceItem, err = tx.CreditBalance.Create().
				SetOrganizationID(departmentId).
				SetBalance(0).
				Save(ctx)
			if err != nil {
				fmt.Printf("create credit_balance failed. organization_id:%v error:%v\n", departmentId, err)
				return err
			}
		} else {
			fmt.Printf("query credit_balance failed: organization_id:%v error:%v\n", departmentId, err)
			return err
		}
	}

	// 积分使用明细记录
	beforeNumber := creditBalanceItem.Balance
	afterNumber := Subtraction(beforeNumber, number)
	_, err = tx.CreditUsage.Create().
		SetUserID(userId).
		SetNumber(number).
		SetBeforeNumber(0).
		SetAfterNumber(0).
		SetNtype(1).
		SetNid(usageDetailItem.ID).
		SetTable("usage_detail").
		SetOrganizationID(departmentId).
		Save(ctx)
	if err != nil {
		fmt.Printf("create credit_usage failed:%v\n", err)
		return err
	}

	// 积分账户扣减积分
	_, err = tx.CreditBalance.Update().Where(creditbalance.OrganizationID(departmentId)).SetBalance(afterNumber).Save(ctx)
	if err != nil {
		fmt.Printf("update credit_balance failed: organization_id:%v error:%v\n", departmentId, err)
		return err
	}

	return nil
}