mirror of
https://github.com/songquanpeng/one-api.git
synced 2025-09-27 05:36:40 +08:00
345 lines
7.8 KiB
Go
345 lines
7.8 KiB
Go
package controller
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"github.com/gin-gonic/gin"
|
||
"net/http"
|
||
"one-api/common"
|
||
"one-api/model"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
func GetAllChannels(c *gin.Context) {
|
||
p, _ := strconv.Atoi(c.Query("p"))
|
||
if p < 0 {
|
||
p = 0
|
||
}
|
||
channels, err := model.GetAllChannels(p*common.ItemsPerPage, common.ItemsPerPage, false)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": channels,
|
||
})
|
||
return
|
||
}
|
||
|
||
func SearchChannels(c *gin.Context) {
|
||
keyword := c.Query("keyword")
|
||
channels, err := model.SearchChannels(keyword)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": channels,
|
||
})
|
||
return
|
||
}
|
||
|
||
func GetChannel(c *gin.Context) {
|
||
id, err := strconv.Atoi(c.Param("id"))
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
channel, err := model.GetChannelById(id, false)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": channel,
|
||
})
|
||
return
|
||
}
|
||
|
||
func AddChannel(c *gin.Context) {
|
||
channel := model.Channel{}
|
||
err := c.ShouldBindJSON(&channel)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
channel.CreatedTime = common.GetTimestamp()
|
||
keys := strings.Split(channel.Key, "\n")
|
||
channels := make([]model.Channel, 0)
|
||
for _, key := range keys {
|
||
if key == "" {
|
||
continue
|
||
}
|
||
localChannel := channel
|
||
localChannel.Key = key
|
||
channels = append(channels, localChannel)
|
||
}
|
||
err = model.BatchInsertChannels(channels)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
})
|
||
return
|
||
}
|
||
|
||
func DeleteChannel(c *gin.Context) {
|
||
id, _ := strconv.Atoi(c.Param("id"))
|
||
channel := model.Channel{Id: id}
|
||
err := channel.Delete()
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
})
|
||
return
|
||
}
|
||
|
||
func UpdateChannel(c *gin.Context) {
|
||
channel := model.Channel{}
|
||
err := c.ShouldBindJSON(&channel)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
err = channel.Update()
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"data": channel,
|
||
})
|
||
return
|
||
}
|
||
|
||
func testChannel(channel *model.Channel, request *ChatRequest) error {
|
||
if request.Model == "" {
|
||
request.Model = "gpt-3.5-turbo"
|
||
if channel.Type == common.ChannelTypeAzure {
|
||
request.Model = "gpt-35-turbo"
|
||
}
|
||
}
|
||
requestURL := common.ChannelBaseURLs[channel.Type]
|
||
if channel.Type == common.ChannelTypeAzure {
|
||
requestURL = fmt.Sprintf("%s/openai/deployments/%s/chat/completions?api-version=2023-03-15-preview", channel.BaseURL, request.Model)
|
||
} else {
|
||
if channel.Type == common.ChannelTypeCustom {
|
||
requestURL = channel.BaseURL
|
||
}
|
||
requestURL += "/v1/chat/completions"
|
||
}
|
||
|
||
jsonData, err := json.Marshal(request)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
req, err := http.NewRequest("POST", requestURL, bytes.NewBuffer(jsonData))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if channel.Type == common.ChannelTypeAzure {
|
||
req.Header.Set("api-key", channel.Key)
|
||
} else {
|
||
req.Header.Set("Authorization", "Bearer "+channel.Key)
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
client := &http.Client{}
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
var response TextResponse
|
||
err = json.NewDecoder(resp.Body).Decode(&response)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if response.Error.Message != "" {
|
||
return errors.New(fmt.Sprintf("type %s, code %s, message %s", response.Error.Type, response.Error.Code, response.Error.Message))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func buildTestRequest(c *gin.Context) *ChatRequest {
|
||
model_ := c.Query("model")
|
||
testRequest := &ChatRequest{
|
||
Model: model_,
|
||
MaxTokens: 1,
|
||
}
|
||
testMessage := Message{
|
||
Role: "user",
|
||
Content: "hi",
|
||
}
|
||
testRequest.Messages = append(testRequest.Messages, testMessage)
|
||
return testRequest
|
||
}
|
||
|
||
func TestChannel(c *gin.Context) {
|
||
id, err := strconv.Atoi(c.Param("id"))
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
channel, err := model.GetChannelById(id, true)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
testRequest := buildTestRequest(c)
|
||
tik := time.Now()
|
||
err = testChannel(channel, testRequest)
|
||
tok := time.Now()
|
||
milliseconds := tok.Sub(tik).Milliseconds()
|
||
go channel.UpdateResponseTime(milliseconds)
|
||
consumedTime := float64(milliseconds) / 1000.0
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
"time": consumedTime,
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
"time": consumedTime,
|
||
})
|
||
return
|
||
}
|
||
|
||
var testAllChannelsLock sync.Mutex
|
||
var testAllChannelsRunning bool = false
|
||
|
||
// disable & notify
|
||
func disableChannel(channelId int, channelName string, reason string) {
|
||
if common.RootUserEmail == "" {
|
||
common.RootUserEmail = model.GetRootUserEmail()
|
||
}
|
||
model.UpdateChannelStatusById(channelId, common.ChannelStatusDisabled)
|
||
subject := fmt.Sprintf("Channel「%s」（#%d）has been disabled", channelName, channelId)
|
||
content := fmt.Sprintf("Channel「%s」（#%d）has been disabled, reason：%s", channelName, channelId, reason)
|
||
err := common.SendEmail(subject, common.RootUserEmail, content)
|
||
if err != nil {
|
||
common.SysError(fmt.Sprintf("Failed to send mail：%s", err.Error()))
|
||
}
|
||
}
|
||
|
||
func testAllChannels(c *gin.Context) error {
|
||
testAllChannelsLock.Lock()
|
||
if testAllChannelsRunning {
|
||
testAllChannelsLock.Unlock()
|
||
return errors.New("Test is already running")
|
||
}
|
||
testAllChannelsRunning = true
|
||
testAllChannelsLock.Unlock()
|
||
channels, err := model.GetAllChannels(0, 0, true)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return err
|
||
}
|
||
testRequest := buildTestRequest(c)
|
||
var disableThreshold = int64(common.ChannelDisableThreshold * 1000)
|
||
if disableThreshold == 0 {
|
||
disableThreshold = 10000000 // a impossible value
|
||
}
|
||
go func() {
|
||
for _, channel := range channels {
|
||
if channel.Status != common.ChannelStatusEnabled {
|
||
continue
|
||
}
|
||
tik := time.Now()
|
||
err := testChannel(channel, testRequest)
|
||
tok := time.Now()
|
||
milliseconds := tok.Sub(tik).Milliseconds()
|
||
if err != nil || milliseconds > disableThreshold {
|
||
if milliseconds > disableThreshold {
|
||
err = errors.New(fmt.Sprintf("Response time %.2fs Threshold exceeded %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0))
|
||
}
|
||
disableChannel(channel.Id, channel.Name, err.Error())
|
||
}
|
||
channel.UpdateResponseTime(milliseconds)
|
||
}
|
||
err := common.SendEmail("Channel test complete", common.RootUserEmail, "The channel test is completed, if no disable notification is received, all channels are normal")
|
||
if err != nil {
|
||
common.SysError(fmt.Sprintf("Failed to send mail：%s", err.Error()))
|
||
}
|
||
testAllChannelsLock.Lock()
|
||
testAllChannelsRunning = false
|
||
testAllChannelsLock.Unlock()
|
||
}()
|
||
return nil
|
||
}
|
||
|
||
func TestAllChannels(c *gin.Context) {
|
||
err := testAllChannels(c)
|
||
if err != nil {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": false,
|
||
"message": err.Error(),
|
||
})
|
||
return
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"success": true,
|
||
"message": "",
|
||
})
|
||
return
|
||
}
|