xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-09-19 10:06:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: update image request handling to always return one image and improve error logging

Browse Source
This commit is contained in:
Laisky.Cai 2024-11-28 09:15:40 +00:00
parent ae7117beee
commit 36e99cf6ec
3 changed files with 18 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
relay/adaptor/replicate/adaptor.go
View File

@ -9,6 +9,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/relay/adaptor" "github.com/songquanpeng/one-api/relay/adaptor"
"github.com/songquanpeng/one-api/relay/adaptor/openai" "github.com/songquanpeng/one-api/relay/adaptor/openai"
"github.com/songquanpeng/one-api/relay/meta" "github.com/songquanpeng/one-api/relay/meta"
@ -29,7 +30,7 @@ func (*Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) {
Guidance: 3, Guidance: 3,
Seed: int(time.Now().UnixNano()), Seed: int(time.Now().UnixNano()),
SafetyTolerance: 5, SafetyTolerance: 5,
NImages: request.N, NImages: 1, // replicate will always return 1 image
Width: 1440, Width: 1440,
Height: 1440, Height: 1440,
AspectRatio: "1:1", AspectRatio: "1:1",
@ -60,6 +61,7 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *me
} }
func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) { func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) {
logger.Info(c, "send image request to replicate")
return adaptor.DoRequestHelper(a, c, meta, requestBody) return adaptor.DoRequestHelper(a, c, meta, requestBody)
} }

9
relay/adaptor/replicate/image.go
View File

@ -39,6 +39,8 @@ import (
// return nil, nil // return nil, nil
// } // }
var errNextLoop = errors.New("next_loop")
func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) { func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) {
if resp.StatusCode != http.StatusCreated { if resp.StatusCode != http.StatusCreated {
payload, _ := io.ReadAll(resp.Body) payload, _ := io.ReadAll(resp.Body)
@ -67,7 +69,6 @@ func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCo
return errors.Wrap(err, "new request") return errors.Wrap(err, "new request")
} }
logger.Debug(c, "send image request to replicate")
taskReq.Header.Set("Authorization", "Bearer "+meta.GetByContext(c).APIKey) taskReq.Header.Set("Authorization", "Bearer "+meta.GetByContext(c).APIKey)
taskResp, err := http.DefaultClient.Do(taskReq) taskResp, err := http.DefaultClient.Do(taskReq)
if err != nil { if err != nil {
@ -97,7 +98,7 @@ func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCo
return errors.Errorf("task failed: %s", taskData.Status) return errors.Errorf("task failed: %s", taskData.Status)
default: default:
time.Sleep(time.Second * 3) time.Sleep(time.Second * 3)
return nil return errNextLoop
} }
output, err := taskData.GetOutput() output, err := taskData.GetOutput()
@ -170,6 +171,10 @@ func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCo
return nil return nil
}() }()
if err != nil { if err != nil {
if errors.Is(err, errNextLoop) {
continue
}
return openai.ErrorWrapper(err, "image_task_failed", http.StatusInternalServerError), nil return openai.ErrorWrapper(err, "image_task_failed", http.StatusInternalServerError), nil
} }

9
relay/controller/image.go
View File

@ -175,7 +175,14 @@ func RelayImageHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
ratio := modelRatio * groupRatio ratio := modelRatio * groupRatio
userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId) userQuota, err := model.CacheGetUserQuota(ctx, meta.UserId)
quota := int64(ratio*imageCostRatio*1000) * int64(imageRequest.N) var quota int64
switch meta.ChannelType {
case channeltype.Replicate:
// replicate always return 1 image
quota = int64(ratio * imageCostRatio * 1000)
default:
quota = int64(ratio*imageCostRatio*1000) * int64(imageRequest.N)
}
if userQuota-quota < 0 { if userQuota-quota < 0 {
return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) return openai.ErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden)