xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-10 02:23:43 +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
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -39,6 +39,8 @@ import (
// return nil, nil
// }
var errNextLoop = errors.New("next_loop")
func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCode, *model.Usage) {
if resp.StatusCode != http.StatusCreated {
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")
}
logger.Debug(c, "send image request to replicate")
taskReq.Header.Set("Authorization", "Bearer "+meta.GetByContext(c).APIKey)
taskResp, err := http.DefaultClient.Do(taskReq)
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)
default:
time.Sleep(time.Second * 3)
return nil
return errNextLoop
}
output, err := taskData.GetOutput()
@@ -170,6 +171,10 @@ func ImageHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusCo
return nil
}()
if err != nil {
if errors.Is(err, errNextLoop) {
continue
}
return openai.ErrorWrapper(err, "image_task_failed", http.StatusInternalServerError), nil
}