Explorar el Código

feat: implement channel settings configuration

fix #620
CalciumIon hace 1 año
padre
commit
f2809917f8

+ 5 - 0
constant/channel_setting.go

@@ -0,0 +1,5 @@
+package constant
+
+var (
+	ForceFormat = "force_format" // ForceFormat 强制格式化为OpenAI格式
+)

+ 1 - 0
middleware/distributor.go

@@ -213,6 +213,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
 	c.Set("channel_id", channel.Id)
 	c.Set("channel_name", channel.Name)
 	c.Set("channel_type", channel.Type)
+	c.Set("channel_setting", channel.GetSetting())
 	if nil != channel.OpenAIOrganization && "" != *channel.OpenAIOrganization {
 		c.Set("channel_organization", *channel.OpenAIOrganization)
 	}

+ 21 - 0
model/channel.go

@@ -34,6 +34,7 @@ type Channel struct {
 	AutoBan           *int    `json:"auto_ban" gorm:"default:1"`
 	OtherInfo         string  `json:"other_info"`
 	Tag               *string `json:"tag" gorm:"index"`
+	Setting           string  `json:"setting" gorm:"type:text"`
 }
 
 func (channel *Channel) GetModels() []string {
@@ -469,3 +470,23 @@ func SearchTags(keyword string, group string, model string, idSort bool) ([]*str
 
 	return tags, nil
 }
+
+func (channel *Channel) GetSetting() map[string]interface{} {
+	setting := make(map[string]interface{})
+	if channel.Setting != "" {
+		err := json.Unmarshal([]byte(channel.Setting), &setting)
+		if err != nil {
+			common.SysError("failed to unmarshal setting: " + err.Error())
+		}
+	}
+	return setting
+}
+
+func (channel *Channel) SetSetting(setting map[string]interface{}) {
+	settingBytes, err := json.Marshal(setting)
+	if err != nil {
+		common.SysError("failed to marshal setting: " + err.Error())
+		return
+	}
+	channel.Setting = string(settingBytes)
+}

+ 51 - 5
relay/channel/openai/relay-openai.go

@@ -5,9 +5,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"github.com/bytedance/gopkg/util/gopool"
-	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"io"
 	"net/http"
 	"one-api/common"
@@ -19,9 +16,33 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/bytedance/gopkg/util/gopool"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 )
 
+func sendStreamData(c *gin.Context, data string, forceFormat bool) error {
+	if data == "" {
+		return nil
+	}
+
+	if forceFormat {
+		var lastStreamResponse dto.ChatCompletionsStreamResponse
+		if err := json.Unmarshal(common.StringToByteSlice(data), &lastStreamResponse); err != nil {
+			return err
+		}
+		return service.ObjectData(c, lastStreamResponse)
+	}
+	return service.StringData(c, data)
+}
+
 func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
+	if resp == nil || resp.Body == nil {
+		common.LogError(c, "invalid response or response body")
+		return service.OpenAIErrorWrapper(fmt.Errorf("invalid response"), "invalid_response", http.StatusInternalServerError), nil
+	}
+
 	containStreamUsage := false
 	var responseId string
 	var createAt int64 = 0
@@ -31,6 +52,13 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
 	var responseTextBuilder strings.Builder
 	var usage = &dto.Usage{}
 	var streamItems []string // store stream items
+	var forceFormat bool
+
+	if info.ChannelType == common.ChannelTypeCustom {
+		if forceFmt, ok := info.ChannelSetting["force_format"].(bool); ok {
+			forceFormat = forceFmt
+		}
+	}
 
 	toolCount := 0
 	scanner := bufio.NewScanner(resp.Body)
@@ -62,7 +90,7 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
 				if lastStreamData != "" {
-					err := service.StringData(c, lastStreamData)
+					err := sendStreamData(c, lastStreamData, forceFormat)
 					if err != nil {
 						common.LogError(c, "streaming error: "+err.Error())
 					}
@@ -105,7 +133,7 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
 		}
 	}
 	if shouldSendLastResp {
-		service.StringData(c, lastStreamData)
+		sendStreamData(c, lastStreamData, forceFormat)
 	}
 
 	// 计算token
@@ -375,6 +403,10 @@ func getTextFromJSON(body []byte) (string, error) {
 }
 
 func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, *dto.RealtimeUsage) {
+	if info == nil || info.ClientWs == nil || info.TargetWs == nil {
+		return service.OpenAIErrorWrapper(fmt.Errorf("invalid websocket connection"), "invalid_connection", http.StatusBadRequest), nil
+	}
+
 	info.IsStream = true
 	clientConn := info.ClientWs
 	targetConn := info.TargetWs
@@ -390,6 +422,11 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
 	sumUsage := &dto.RealtimeUsage{}
 
 	gopool.Go(func() {
+		defer func() {
+			if r := recover(); r != nil {
+				errChan <- fmt.Errorf("panic in client reader: %v", r)
+			}
+		}()
 		for {
 			select {
 			case <-c.Done():
@@ -445,6 +482,11 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
 	})
 
 	gopool.Go(func() {
+		defer func() {
+			if r := recover(); r != nil {
+				errChan <- fmt.Errorf("panic in target reader: %v", r)
+			}
+		}()
 		for {
 			select {
 			case <-c.Done():
@@ -568,6 +610,10 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
 }
 
 func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) error {
+	if usage == nil || totalUsage == nil {
+		return fmt.Errorf("invalid usage pointer")
+	}
+
 	totalUsage.TotalTokens += usage.TotalTokens
 	totalUsage.InputTokens += usage.InputTokens
 	totalUsage.OutputTokens += usage.OutputTokens

+ 6 - 2
relay/common/relay_info.go

@@ -1,13 +1,14 @@
 package common
 
 import (
-	"github.com/gin-gonic/gin"
-	"github.com/gorilla/websocket"
 	"one-api/common"
 	"one-api/dto"
 	"one-api/relay/constant"
 	"strings"
 	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 )
 
 type RelayInfo struct {
@@ -43,6 +44,7 @@ type RelayInfo struct {
 	RealtimeTools        []dto.RealTimeTool
 	IsFirstRequest       bool
 	AudioUsage           bool
+	ChannelSetting       map[string]interface{}
 }
 
 func GenRelayInfoWs(c *gin.Context, ws *websocket.Conn) *RelayInfo {
@@ -57,6 +59,7 @@ func GenRelayInfoWs(c *gin.Context, ws *websocket.Conn) *RelayInfo {
 func GenRelayInfo(c *gin.Context) *RelayInfo {
 	channelType := c.GetInt("channel_type")
 	channelId := c.GetInt("channel_id")
+	channelSetting := c.GetStringMap("channel_setting")
 
 	tokenId := c.GetInt("token_id")
 	tokenKey := c.GetString("token_key")
@@ -87,6 +90,7 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 		ApiVersion:        c.GetString("api_version"),
 		ApiKey:            strings.TrimPrefix(c.Request.Header.Get("Authorization"), "Bearer "),
 		Organization:      c.GetString("channel_organization"),
+		ChannelSetting:    channelSetting,
 	}
 	if strings.HasPrefix(c.Request.URL.Path, "/pg") {
 		info.IsPlayground = true

+ 36 - 0
web/src/pages/Channel/EditChannel.js

@@ -964,6 +964,42 @@ const EditChannel = (props) => {
             value={inputs.weight}
             autoComplete="new-password"
           />
+          {inputs.type === 8 && (
+          <>
+            <div style={{ marginTop: 10 }}>
+              <Typography.Text strong>
+                {t('渠道额外设置')}:
+              </Typography.Text>
+            </div>
+            <TextArea
+              placeholder={t('此项可选,用于配置渠道特定设置,为一个 JSON 字符串,例如:') + '\n{\n  "force_format": true\n}'}
+              name="setting"
+              onChange={(value) => {
+                handleInputChange('setting', value);
+              }}
+              autosize
+              value={inputs.setting}
+              autoComplete="new-password"
+            />
+            <Typography.Text
+              style={{
+                color: 'rgba(var(--semi-blue-5), 1)', 
+                userSelect: 'none',
+                cursor: 'pointer'
+              }}
+              onClick={() => {
+                handleInputChange(
+                  'setting',
+                  JSON.stringify({
+                    force_format: true
+                  }, null, 2)
+                );
+              }}
+            >
+              {t('填入模板')}
+              </Typography.Text>
+            </>
+          )}  
         </Spin>
       </SideSheet>
     </>