Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support return date for billing

Browse Source
This commit is contained in:
ckt1031 2023-07-11 17:40:52 +08:00
parent 12365ccf69
commit 80d5d6edfb
4 changed files with 35 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
controller/billing.go
View File

@ -1,18 +1,24 @@
package controller package controller
import ( import (
"github.com/gin-gonic/gin"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
"time"
"github.com/gin-gonic/gin"
) )
func GetSubscription(c *gin.Context) { func GetSubscription(c *gin.Context) {
var quota int var quota int
var err error var err error
var token *model.Token var expirationDate int64
if common.DisplayTokenStatEnabled {
tokenId := c.GetInt("token_id") tokenId := c.GetInt("token_id")
token, err = model.GetTokenById(tokenId) token, err := model.GetTokenById(tokenId)
expirationDate = token.ExpiredTime
if common.DisplayTokenStatEnabled {
quota = token.RemainQuota quota = token.RemainQuota
} else { } else {
userId := c.GetInt("id") userId := c.GetInt("id")
@ -41,6 +47,7 @@ func GetSubscription(c *gin.Context) {
SoftLimitUSD: amount, SoftLimitUSD: amount,
HardLimitUSD: amount, HardLimitUSD: amount,
SystemHardLimitUSD: amount, SystemHardLimitUSD: amount,
AccessUntil: time.Unix(expirationDate, 0),
} }
c.JSON(200, subscription) c.JSON(200, subscription)
return return

1
controller/channel-billing.go
View File

@ -22,6 +22,7 @@ type OpenAISubscriptionResponse struct {
SoftLimitUSD float64 `json:"soft_limit_usd"` SoftLimitUSD float64 `json:"soft_limit_usd"`
HardLimitUSD float64 `json:"hard_limit_usd"` HardLimitUSD float64 `json:"hard_limit_usd"`
SystemHardLimitUSD float64 `json:"system_hard_limit_usd"` SystemHardLimitUSD float64 `json:"system_hard_limit_usd"`
AccessUntil time.Time `json:"access_until"`
} }
type OpenAIUsageDailyCost struct { type OpenAIUsageDailyCost struct {

4
controller/channel-test.go
View File

@ -6,7 +6,7 @@ import (
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"io/ioutil" "io"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
@ -100,7 +100,7 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
} }
} }
} else { } else {
body, err := ioutil.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
return err return err
} }

17
controller/relay-text.go
View File

@ -6,12 +6,14 @@ import (
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"github.com/gin-gonic/gin"
"io" "io"
"log"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
"strings" "strings"
"github.com/gin-gonic/gin"
) )
func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
@ -157,6 +159,19 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
return errorWrapper(err, "do_request_failed", http.StatusInternalServerError) return errorWrapper(err, "do_request_failed", http.StatusInternalServerError)
} }
if resp.StatusCode != http.StatusOK {
// Print Data if Error
bodyBytes, err := io.ReadAll(resp.Body)
if err != nil {
return errorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
}
bodyString := string(bodyBytes)
log.Printf("Error: %s", bodyString)
return errorWrapper(err, "request_failed", resp.StatusCode)
}
err = req.Body.Close() err = req.Body.Close()
if err != nil { if err != nil {
return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError) return errorWrapper(err, "close_request_body_failed", http.StatusInternalServerError)