api2gpt-mid/api.go

package main

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"net/http"

	"github.com/gin-gonic/gin"
)

type BalanceInfo struct {
	ServerAddress string  `json:"server_address"`
	AvailableKey  string  `json:"available_key"`
	UserBalance   float64 `json:"user_balance"`
	TokenRatio    float64 `json:"token_ratio"`
}

type Consumption struct {
	SecretKey        string `json:"secretKey"`
	Model            string `json:"model"`
	MsgId            string `json:"msgId"`
	PromptTokens     int    `json:"promptTokens"`
	CompletionTokens int    `json:"completionTokens"`
	TotalTokens      int    `json:"totalTokens"`
}

func mockBalanceInquiry(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"server_address": "https://gptp.any-door.cn",
		"available_key":  "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",
		"user_balance":   10000,
		"token_ratio":    1000,
	})
}

func mockBalanceConsumption(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"success": "true",
	})
}

// 余额查询api调用
func balanceInquiry(key string, model string) (*BalanceInfo, error) {
	url := "http://localhost:8080/mock1?key=" + key + "&model=" + model
	req, err := http.NewRequest("POST", url, nil)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var balanceInfo BalanceInfo
	if err := json.Unmarshal(body, &balanceInfo); err != nil {
		return nil, err
	}

	return &balanceInfo, nil
}

// 余额消费
func balanceConsumption(key string, model string, prompt_tokens int, completion_tokens int, total_tokens int, msg_id string) (string, error) {
	var data = Consumption{
		SecretKey:        key,
		Model:            model,
		MsgId:            msg_id,
		PromptTokens:     prompt_tokens,
		CompletionTokens: completion_tokens,
		TotalTokens:      total_tokens,
	}

	jsonData, err := json.Marshal(data)
	// 构造post请求的body
	reqBody := bytes.NewBuffer(jsonData)

	url := "http://121.4.100.155:8080/other/usageRecord"
	req2, err := http.NewRequest("POST", url, reqBody)

	// 设置http请求的header
	req2.Header.Set("Content-Type", "application/json")

	client := &http.Client{}
	resp, err := client.Do(req2)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		panic(err)
	}
	return string(body), nil
}