Просмотр исходного кода

refactor: Improve message content handling and quota error responses

1808837298@qq.com 1 год назад
Родитель
Сommit
b9b69b01e5
3 измененных файлов с 27 добавлено и 26 удалено
  1. 20 9
      dto/openai_request.go
  2. 6 16
      relay/relay-image.go
  3. 1 1
      relay/relay-text.go

+ 20 - 9
dto/openai_request.go

@@ -88,15 +88,15 @@ func (r GeneralOpenAIRequest) ParseInput() []string {
 }
 
 type Message struct {
-	Role    string          `json:"role"`
-	Content json.RawMessage `json:"content"`
-	// parsedContent not json field
-	parsedContent    []MediaContent
-	Name             *string         `json:"name,omitempty"`
-	Prefix           *bool           `json:"prefix,omitempty"`
-	ReasoningContent string          `json:"reasoning_content,omitempty"`
-	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
-	ToolCallId       string          `json:"tool_call_id,omitempty"`
+	Role                string          `json:"role"`
+	Content             json.RawMessage `json:"content"`
+	Name                *string         `json:"name,omitempty"`
+	Prefix              *bool           `json:"prefix,omitempty"`
+	ReasoningContent    string          `json:"reasoning_content,omitempty"`
+	ToolCalls           json.RawMessage `json:"tool_calls,omitempty"`
+	ToolCallId          string          `json:"tool_call_id,omitempty"`
+	parsedContent       []MediaContent
+	parsedStringContent *string
 }
 
 type MediaContent struct {
@@ -150,6 +150,9 @@ func (m *Message) SetToolCalls(toolCalls any) {
 }
 
 func (m *Message) StringContent() string {
+	if m.parsedStringContent != nil {
+		return *m.parsedStringContent
+	}
 	var stringContent string
 	if err := json.Unmarshal(m.Content, &stringContent); err == nil {
 		return stringContent
@@ -160,16 +163,24 @@ func (m *Message) StringContent() string {
 func (m *Message) SetStringContent(content string) {
 	jsonContent, _ := json.Marshal(content)
 	m.Content = jsonContent
+	m.parsedStringContent = &content
+	m.parsedContent = nil
 }
 
 func (m *Message) SetMediaContent(content []MediaContent) {
 	jsonContent, _ := json.Marshal(content)
 	m.Content = jsonContent
+	m.parsedContent = nil
+	m.parsedStringContent = nil
 }
 
 func (m *Message) IsStringContent() bool {
+	if m.parsedStringContent != nil {
+		return true
+	}
 	var stringContent string
 	if err := json.Unmarshal(m.Content, &stringContent); err == nil {
+		m.parsedStringContent = &stringContent
 		return true
 	}
 	return false

+ 6 - 16
relay/relay-image.go

@@ -86,15 +86,13 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 
 	imageRequest.Model = relayInfo.UpstreamModelName
 
-	modelPrice, success := common.GetModelPrice(imageRequest.Model, true)
-	if !success {
-		modelRatio := common.GetModelRatio(imageRequest.Model)
+	priceData := helper.ModelPriceHelper(c, relayInfo, 0, 0)
+	if !priceData.UsePrice {
 		// modelRatio 16 = modelPrice $0.04
 		// per 1 modelRatio = $0.04 / 16
-		modelPrice = 0.0025 * modelRatio
+		priceData.ModelPrice = 0.0025 * priceData.ModelRatio
 	}
 
-	groupRatio := setting.GetGroupRatio(relayInfo.Group)
 	userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
 
 	sizeRatio := 1.0
@@ -117,11 +115,11 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 		}
 	}
 
-	imageRatio := modelPrice * sizeRatio * qualityRatio * float64(imageRequest.N)
-	quota := int(imageRatio * groupRatio * common.QuotaPerUnit)
+	imageRatio := priceData.ModelPrice * sizeRatio * qualityRatio * float64(imageRequest.N)
+	quota := int(imageRatio * priceData.GroupRatio * common.QuotaPerUnit)
 
 	if userQuota-quota < 0 {
-		return service.OpenAIErrorWrapperLocal(errors.New(fmt.Sprintf("image pre-consumed quota failed, user quota: %d, need quota: %d", userQuota, quota)), "insufficient_user_quota", http.StatusBadRequest)
+		return service.OpenAIErrorWrapperLocal(fmt.Errorf("image pre-consumed quota failed, user quota: %s, need quota: %s", common.FormatQuota(userQuota), common.FormatQuota(quota)), "insufficient_user_quota", http.StatusForbidden)
 	}
 
 	adaptor := GetAdaptor(relayInfo.ApiType)
@@ -178,14 +176,6 @@ func ImageHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 		quality = "hd"
 	}
 
-	priceData := helper.PriceData{
-		UsePrice:               true,
-		GroupRatio:             groupRatio,
-		ModelPrice:             modelPrice,
-		ModelRatio:             0,
-		ShouldPreConsumedQuota: 0,
-	}
-
 	logContent := fmt.Sprintf("大小 %s, 品质 %s", imageRequest.Size, quality)
 	postConsumeQuota(c, relayInfo, usage, 0, userQuota, priceData, logContent)
 	return nil

+ 1 - 1
relay/relay-text.go

@@ -246,7 +246,7 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo
 		return 0, 0, service.OpenAIErrorWrapperLocal(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)
 	}
 	if userQuota-preConsumedQuota < 0 {
-		return 0, 0, service.OpenAIErrorWrapperLocal(fmt.Errorf("chat pre-consumed quota failed, user quota: %s, need quota: %s", common.FormatQuota(userQuota), common.FormatQuota(preConsumedQuota)), "insufficient_user_quota", http.StatusBadRequest)
+		return 0, 0, service.OpenAIErrorWrapperLocal(fmt.Errorf("chat pre-consumed quota failed, user quota: %s, need quota: %s", common.FormatQuota(userQuota), common.FormatQuota(preConsumedQuota)), "insufficient_user_quota", http.StatusForbidden)
 	}
 	if userQuota > 100*preConsumedQuota {
 		// 用户额度充足,判断令牌额度是否充足