Skip to content

Commit 7a71841

Browse files
zijiren233zijiren233
committed
refactor: do response usage filed
1 parent 42b210a commit 7a71841

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

81 files changed

+833
-775
lines changed

‎core/controller/relay-controller.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -126,11 +126,8 @@ func (w *wrapAdaptor) DoResponse(
126126
store adaptor.Store,
127127
c *gin.Context,
128128
resp *http.Response,
129-
) (*model.Usage, adaptor.Error) {
129+
) (model.Usage, adaptor.Error) {
130130
usage, relayErr := w.Adaptor.DoResponse(meta, store, c, resp)
131-
if usage == nil {
132-
return nil, relayErr
133-
}
134131

135132
if usage.TotalTokens > 0 {
136133
count, overLimitCount, secondCount := reqlimit.PushChannelModelTokensRequest(

‎core/relay/adaptor/ali/adaptor.go‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,44 +29,44 @@ func (a *Adaptor) DefaultBaseURL() string {
2929
return baseURL
3030
}
3131

32-
func (a *Adaptor) GetRequestURL(meta *meta.Meta, _ adaptor.Store) (*adaptor.RequestURL, error) {
32+
func (a *Adaptor) GetRequestURL(meta *meta.Meta, _ adaptor.Store) (adaptor.RequestURL, error) {
3333
u := meta.Channel.BaseURL
3434
if u == "" {
3535
u = baseURL
3636
}
3737
switch meta.Mode {
3838
case mode.ImagesGenerations:
39-
return &adaptor.RequestURL{
39+
return adaptor.RequestURL{
4040
Method: http.MethodPost,
4141
URL: u + "/api/v1/services/aigc/text2image/image-synthesis",
4242
}, nil
4343
case mode.ChatCompletions:
44-
return &adaptor.RequestURL{
44+
return adaptor.RequestURL{
4545
Method: http.MethodPost,
4646
URL: u + "/compatible-mode/v1/chat/completions",
4747
}, nil
4848
case mode.Completions:
49-
return &adaptor.RequestURL{
49+
return adaptor.RequestURL{
5050
Method: http.MethodPost,
5151
URL: u + "/compatible-mode/v1/completions",
5252
}, nil
5353
case mode.Embeddings:
54-
return &adaptor.RequestURL{
54+
return adaptor.RequestURL{
5555
Method: http.MethodPost,
5656
URL: u + "/compatible-mode/v1/embeddings",
5757
}, nil
5858
case mode.AudioSpeech, mode.AudioTranscription:
59-
return &adaptor.RequestURL{
59+
return adaptor.RequestURL{
6060
Method: http.MethodPost,
6161
URL: u + "/api-ws/v1/inference",
6262
}, nil
6363
case mode.Rerank:
64-
return &adaptor.RequestURL{
64+
return adaptor.RequestURL{
6565
Method: http.MethodPost,
6666
URL: u + "/api/v1/services/rerank/text-rerank/text-rerank",
6767
}, nil
6868
default:
69-
return nil, fmt.Errorf("unsupported mode: %s", meta.Mode)
69+
return adaptor.RequestURL{}, fmt.Errorf("unsupported mode: %s", meta.Mode)
7070
}
7171
}
7272

@@ -114,7 +114,7 @@ func (a *Adaptor) ConvertRequest(
114114
meta *meta.Meta,
115115
store adaptor.Store,
116116
req *http.Request,
117-
) (*adaptor.ConvertRequestResult, error) {
117+
) (adaptor.ConvertResult, error) {
118118
switch meta.Mode {
119119
case mode.ImagesGenerations:
120120
return ConvertImageRequest(meta, req)
@@ -143,7 +143,7 @@ func (a *Adaptor) ConvertRequest(
143143
case mode.AudioTranscription:
144144
return ConvertSTTRequest(meta, req)
145145
default:
146-
return nil, fmt.Errorf("unsupported mode: %s", meta.Mode)
146+
return adaptor.ConvertResult{}, fmt.Errorf("unsupported mode: %s", meta.Mode)
147147
}
148148
}
149149

@@ -170,7 +170,7 @@ func (a *Adaptor) DoResponse(
170170
store adaptor.Store,
171171
c *gin.Context,
172172
resp *http.Response,
173-
) (*model.Usage, adaptor.Error) {
173+
) (model.Usage, adaptor.Error) {
174174
switch meta.Mode {
175175
case mode.ImagesGenerations:
176176
return ImageHandler(meta, c, resp)
@@ -179,23 +179,23 @@ func (a *Adaptor) DoResponse(
179179
case mode.ChatCompletions:
180180
reqBody, err := common.GetRequestBody(c.Request)
181181
if err != nil {
182-
return nil, relaymodel.WrapperOpenAIErrorWithMessage(
182+
return model.Usage{}, relaymodel.WrapperOpenAIErrorWithMessage(
183183
fmt.Sprintf("get request body failed: %s", err),
184184
"get_request_body_failed",
185185
http.StatusInternalServerError,
186186
)
187187
}
188188
enableSearch, err := getEnableSearch(reqBody)
189189
if err != nil {
190-
return nil, relaymodel.WrapperOpenAIErrorWithMessage(
190+
return model.Usage{}, relaymodel.WrapperOpenAIErrorWithMessage(
191191
fmt.Sprintf("get enable_search failed: %s", err),
192192
"get_enable_search_failed",
193193
http.StatusInternalServerError,
194194
)
195195
}
196196
u, e := openai.DoResponse(meta, store, c, resp)
197197
if e != nil {
198-
return nil, e
198+
return model.Usage{}, e
199199
}
200200
if enableSearch {
201201
u.WebSearchCount++
@@ -208,7 +208,7 @@ func (a *Adaptor) DoResponse(
208208
case mode.AudioTranscription:
209209
return STTDoResponse(meta, c, resp)
210210
default:
211-
return nil, relaymodel.WrapperOpenAIErrorWithMessage(
211+
return model.Usage{}, relaymodel.WrapperOpenAIErrorWithMessage(
212212
fmt.Sprintf("unsupported mode: %s", meta.Mode),
213213
"unsupported_mode",
214214
http.StatusBadRequest,

‎core/relay/adaptor/ali/embeddings.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,19 +88,23 @@ func EmbeddingsHandler(
8888
meta *meta.Meta,
8989
c *gin.Context,
9090
resp *http.Response,
91-
) (*model.Usage, adaptor.Error) {
91+
) (model.Usage, adaptor.Error) {
9292
defer resp.Body.Close()
9393

9494
log := middleware.GetLogger(c)
9595

9696
responseBody, err := io.ReadAll(resp.Body)
9797
if err != nil {
98-
return nil, relaymodel.WrapperOpenAIError(err, "read_response_body_failed", resp.StatusCode)
98+
return model.Usage{}, relaymodel.WrapperOpenAIError(
99+
err,
100+
"read_response_body_failed",
101+
resp.StatusCode,
102+
)
99103
}
100104
var respBody EmbeddingResponse
101105
err = sonic.Unmarshal(responseBody, &respBody)
102106
if err != nil {
103-
return nil, relaymodel.WrapperOpenAIError(
107+
return model.Usage{}, relaymodel.WrapperOpenAIError(
104108
err,
105109
"unmarshal_response_body_failed",
106110
resp.StatusCode,

‎core/relay/adaptor/ali/image.go‎

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,10 @@ const MetaResponseFormat = "response_format"
2828
func ConvertImageRequest(
2929
meta *meta.Meta,
3030
req *http.Request,
31-
) (*adaptor.ConvertRequestResult, error) {
31+
) (adaptor.ConvertResult, error) {
3232
request, err := utils.UnmarshalImageRequest(req)
3333
if err != nil {
34-
return nil, err
34+
return adaptor.ConvertResult{}, err
3535
}
3636
request.Model = meta.ActualModel
3737

@@ -46,11 +46,12 @@ func ConvertImageRequest(
4646

4747
data, err := sonic.Marshal(&imageRequest)
4848
if err != nil {
49-
return nil, err
49+
return adaptor.ConvertResult{}, err
5050
}
51-
return &adaptor.ConvertRequestResult{
51+
return adaptor.ConvertResult{
5252
Header: http.Header{
5353
"X-Dashscope-Async": {"enable"},
54+
"Content-Type": {"application/json"},
5455
},
5556
Body: bytes.NewReader(data),
5657
}, nil
@@ -60,9 +61,9 @@ func ImageHandler(
6061
meta *meta.Meta,
6162
c *gin.Context,
6263
resp *http.Response,
63-
) (*model.Usage, adaptor.Error) {
64+
) (model.Usage, adaptor.Error) {
6465
if resp.StatusCode != http.StatusOK {
65-
return nil, openai.ErrorHanlder(resp)
66+
return model.Usage{}, openai.ErrorHanlder(resp)
6667
}
6768

6869
defer resp.Body.Close()
@@ -74,15 +75,15 @@ func ImageHandler(
7475
var aliTaskResponse TaskResponse
7576
responseBody, err := io.ReadAll(resp.Body)
7677
if err != nil {
77-
return nil, relaymodel.WrapperOpenAIError(
78+
return model.Usage{}, relaymodel.WrapperOpenAIError(
7879
err,
7980
"read_response_body_failed",
8081
http.StatusInternalServerError,
8182
)
8283
}
8384
err = sonic.Unmarshal(responseBody, &aliTaskResponse)
8485
if err != nil {
85-
return nil, relaymodel.WrapperOpenAIError(
86+
return model.Usage{}, relaymodel.WrapperOpenAIError(
8687
err,
8788
"unmarshal_response_body_failed",
8889
http.StatusInternalServerError,
@@ -91,7 +92,7 @@ func ImageHandler(
9192

9293
if aliTaskResponse.Message != "" {
9394
log.Error("aliAsyncTask err: " + aliTaskResponse.Message)
94-
return nil, relaymodel.WrapperOpenAIError(
95+
return model.Usage{}, relaymodel.WrapperOpenAIError(
9596
errors.New(aliTaskResponse.Message),
9697
"ali_async_task_failed",
9798
http.StatusInternalServerError,
@@ -100,15 +101,15 @@ func ImageHandler(
100101

101102
aliResponse, err := asyncTaskWait(c, aliTaskResponse.Output.TaskID, meta.Channel.Key)
102103
if err != nil {
103-
return nil, relaymodel.WrapperOpenAIError(
104+
return model.Usage{}, relaymodel.WrapperOpenAIError(
104105
err,
105106
"ali_async_task_wait_failed",
106107
http.StatusInternalServerError,
107108
)
108109
}
109110

110111
if aliResponse.Output.TaskStatus != "SUCCEEDED" {
111-
return nil, relaymodel.WrapperOpenAIErrorWithMessage(
112+
return model.Usage{}, relaymodel.WrapperOpenAIErrorWithMessage(
112113
aliResponse.Output.Message,
113114
"ali_error",
114115
resp.StatusCode,
@@ -118,7 +119,7 @@ func ImageHandler(
118119
fullTextResponse := responseAli2OpenAIImage(c.Request.Context(), aliResponse, responseFormat)
119120
jsonResponse, err := sonic.Marshal(fullTextResponse)
120121
if err != nil {
121-
return nil, relaymodel.WrapperOpenAIError(
122+
return fullTextResponse.Usage.ToModelUsage(), relaymodel.WrapperOpenAIError(
122123
err,
123124
"marshal_response_body_failed",
124125
http.StatusInternalServerError,
@@ -127,10 +128,7 @@ func ImageHandler(
127128
c.Writer.Header().Set("Content-Type", "application/json")
128129
c.Writer.Header().Set("Content-Length", strconv.Itoa(len(jsonResponse)))
129130
_, _ = c.Writer.Write(jsonResponse)
130-
return &model.Usage{
131-
OutputTokens: model.ZeroNullInt64(len(jsonResponse)),
132-
TotalTokens: model.ZeroNullInt64(len(jsonResponse)),
133-
}, nil
131+
return fullTextResponse.Usage.ToModelUsage(), nil
134132
}
135133

136134
func asyncTask(ctx context.Context, taskID, key string) (*TaskResponse, error) {
@@ -229,5 +227,9 @@ func responseAli2OpenAIImage(
229227
RevisedPrompt: "",
230228
})
231229
}
230+
imageResponse.Usage = &relaymodel.ImageUsage{
231+
OutputTokens: int64(len(imageResponse.Data)),
232+
TotalTokens: int64(len(imageResponse.Data)),
233+
}
232234
return &imageResponse
233235
}

‎core/relay/adaptor/ali/rerank.go‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,11 @@ type RerankUsage struct {
3232
func ConvertRerankRequest(
3333
meta *meta.Meta,
3434
req *http.Request,
35-
) (*adaptor.ConvertRequestResult, error) {
35+
) (adaptor.ConvertResult, error) {
3636
reqMap := make(map[string]any)
3737
err := common.UnmarshalBodyReusable(req, &reqMap)
3838
if err != nil {
39-
return nil, err
39+
return adaptor.ConvertResult{}, err
4040
}
4141
reqMap["model"] = meta.ActualModel
4242
reqMap["input"] = map[string]any{
@@ -56,21 +56,23 @@ func ConvertRerankRequest(
5656
reqMap["parameters"] = parameters
5757
jsonData, err := sonic.Marshal(reqMap)
5858
if err != nil {
59-
return nil, err
59+
return adaptor.ConvertResult{}, err
6060
}
61-
return &adaptor.ConvertRequestResult{
62-
Header: nil,
63-
Body: bytes.NewReader(jsonData),
61+
return adaptor.ConvertResult{
62+
Header: http.Header{
63+
"Content-Type": {"application/json"},
64+
},
65+
Body: bytes.NewReader(jsonData),
6466
}, nil
6567
}
6668

6769
func RerankHandler(
6870
meta *meta.Meta,
6971
c *gin.Context,
7072
resp *http.Response,
71-
) (*model.Usage, adaptor.Error) {
73+
) (model.Usage, adaptor.Error) {
7274
if resp.StatusCode != http.StatusOK {
73-
return nil, openai.ErrorHanlder(resp)
75+
return model.Usage{}, openai.ErrorHanlder(resp)
7476
}
7577

7678
defer resp.Body.Close()
@@ -79,7 +81,7 @@ func RerankHandler(
7981

8082
responseBody, err := io.ReadAll(resp.Body)
8183
if err != nil {
82-
return nil, relaymodel.WrapperOpenAIError(
84+
return model.Usage{}, relaymodel.WrapperOpenAIError(
8385
err,
8486
"read_response_body_failed",
8587
http.StatusInternalServerError,
@@ -88,7 +90,7 @@ func RerankHandler(
8890
var rerankResponse RerankResponse
8991
err = sonic.Unmarshal(responseBody, &rerankResponse)
9092
if err != nil {
91-
return nil, relaymodel.WrapperOpenAIError(
93+
return model.Usage{}, relaymodel.WrapperOpenAIError(
9294
err,
9395
"unmarshal_response_body_failed",
9496
http.StatusInternalServerError,
@@ -106,14 +108,14 @@ func RerankHandler(
106108
ID: rerankResponse.RequestID,
107109
}
108110

109-
var usage *model.Usage
111+
var usage model.Usage
110112
if rerankResponse.Usage == nil {
111-
usage = &model.Usage{
113+
usage = model.Usage{
112114
InputTokens: meta.RequestUsage.InputTokens,
113115
TotalTokens: meta.RequestUsage.InputTokens,
114116
}
115117
} else {
116-
usage = &model.Usage{
118+
usage = model.Usage{
117119
InputTokens: model.ZeroNullInt64(rerankResponse.Usage.TotalTokens),
118120
TotalTokens: model.ZeroNullInt64(rerankResponse.Usage.TotalTokens),
119121
}

0 commit comments

Comments
 (0)