-
Notifications
You must be signed in to change notification settings - Fork 11.3k
feat: add model redirection for task requests #2978
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
7fc25a5
c651727
347ad04
ff11c92
10a4739
f15b85f
797c7ac
a469e58
06cdad5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,8 +1,11 @@ | ||
| package sora | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "fmt" | ||
| "io" | ||
| "mime" | ||
| "mime/multipart" | ||
| "net/http" | ||
| "strings" | ||
|
|
||
|
|
@@ -107,9 +110,112 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn | |
| if err != nil { | ||
| return nil, errors.Wrap(err, "get_request_body_failed") | ||
| } | ||
| bodyBytes, err := storage.Bytes() | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "read_request_body_failed") | ||
| } | ||
|
|
||
| // 检查是否需要模型重定向 | ||
| if !info.IsModelMapped { | ||
| // 如果不需要重定向,直接返回原始请求体 | ||
| return bytes.NewReader(bodyBytes), nil | ||
| } | ||
|
|
||
| contentType := c.Request.Header.Get("Content-Type") | ||
|
|
||
| // 处理multipart/form-data请求 | ||
| if strings.Contains(contentType, "multipart/form-data") { | ||
| return buildRequestBodyWithMappedModel(bodyBytes, contentType, info.UpstreamModelName) | ||
| } | ||
| // 处理JSON请求 | ||
| if strings.Contains(contentType, "application/json") { | ||
| var jsonData map[string]interface{} | ||
| if err := common.Unmarshal(bodyBytes, &jsonData); err != nil { | ||
| return nil, errors.Wrap(err, "unmarshal_json_failed") | ||
| } | ||
|
|
||
| // 暂不更改返回 | ||
| // jsonData["model"] = info.UpstreamModelName | ||
|
|
||
| // 重新编码为JSON | ||
| newBody, err := common.Marshal(jsonData) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "marshal_json_failed") | ||
| } | ||
|
|
||
| return bytes.NewReader(newBody), nil | ||
| } | ||
|
|
||
| return common.ReaderOnly(storage), nil | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Potential stale reader: Line 113 calls This line is only reached when Proposed fix: use the already-read bytes- return common.ReaderOnly(storage), nil
+ return bytes.NewReader(bodyBytes), nil🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| func buildRequestBodyWithMappedModel(originalBody []byte, contentType, redirectedModel string) (io.Reader, error) { | ||
| newBuffer := &bytes.Buffer{} | ||
| writer := multipart.NewWriter(newBuffer) | ||
|
|
||
| _, params, err := mime.ParseMediaType(contentType) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "parse_content_type_failed") | ||
| } | ||
| boundary, ok := params["boundary"] | ||
| if !ok { | ||
| return nil, errors.New("boundary_not_found_in_content_type") | ||
| } | ||
| if err := writer.SetBoundary(boundary); err != nil { | ||
| return nil, errors.Wrap(err, "set_boundary_failed") | ||
| } | ||
| r := multipart.NewReader(bytes.NewReader(originalBody), boundary) | ||
|
|
||
| for { | ||
| part, err := r.NextPart() | ||
| if err == io.EOF { | ||
| break | ||
| } | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "read_multipart_part_failed") | ||
| } | ||
|
|
||
| fieldName := part.FormName() | ||
|
|
||
| if fieldName == "model" { | ||
| // 修改 model 字段为映射后的模型名 | ||
| // 暂不更改返回 | ||
| //if err := writer.WriteField("model", redirectedModel); err != nil { | ||
| // return nil, errors.Wrap(err, "write_model_field_failed") | ||
| //} | ||
| } else { | ||
| // 对于其他字段,保留原始内容 | ||
| if part.FileName() != "" { | ||
| newPart, err := writer.CreatePart(part.Header) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "create_form_file_failed") | ||
| } | ||
| if _, err := io.Copy(newPart, part); err != nil { | ||
| return nil, errors.Wrap(err, "copy_file_content_failed") | ||
| } | ||
| } else { | ||
| newPart, err := writer.CreatePart(part.Header) | ||
| if err != nil { | ||
| return nil, errors.Wrap(err, "create_form_field_failed") | ||
| } | ||
| if _, err := io.Copy(newPart, part); err != nil { | ||
| return nil, errors.Wrap(err, "copy_field_content_failed") | ||
| } | ||
| } | ||
| } | ||
|
|
||
| if err := part.Close(); err != nil { | ||
| return nil, errors.Wrap(err, "close_part_failed") | ||
| } | ||
| } | ||
|
|
||
| if err := writer.Close(); err != nil { | ||
| return nil, errors.Wrap(err, "close_multipart_writer_failed") | ||
| } | ||
|
|
||
| return newBuffer, nil | ||
| } | ||
|
|
||
| // DoRequest delegates to common helper. | ||
| func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { | ||
| return channel.DoTaskApiRequest(a, c, info, requestBody) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Model rewrite is commented out, but the multipart path silently drops the
modelfield — this is a data-loss bug.When
info.IsModelMappedistrueand the content type ismultipart/form-data, the code entersbuildRequestBodyWithMappedModel. Inside that helper, whenfieldName == "model"(Line 180), theWriteFieldcall is commented out (Lines 183–185) and theelsebranch that preserves the part is skipped. The result: themodelfield is silently removed from the rebuilt multipart body sent upstream. The upstream API will receive a request with no model field at all.Additionally, the JSON path (Lines 131–147) unmarshals and re-marshals the body with zero modifications (Line 138 is commented out), adding unnecessary overhead.
If the intent is to not rewrite the model field for now, the multipart path must still preserve the original model field rather than dropping it. The simplest fix: when
IsModelMappedisfalse(or rewrite is disabled), return the body unchanged — which is already done at Line 119–122. The current code reaches Lines 127+ only whenIsModelMappedistrue, making the commented-out rewrite contradictory.🐛 Proposed fix: preserve the model field in the multipart path
Option A — If intent is to defer model rewrite entirely, just return the original body when mapped too:
if !info.IsModelMapped { - // 如果不需要重定向,直接返回原始请求体 return bytes.NewReader(bodyBytes), nil } + // TODO: model rewrite for mapped models is not yet enabled; + // return the original body to avoid dropping the model field. + return bytes.NewReader(bodyBytes), nilOption B — If intent is to rewrite the model, uncomment the write:
if fieldName == "model" { - // 修改 model 字段为映射后的模型名 - // 暂不更改返回 - //if err := writer.WriteField("model", redirectedModel); err != nil { - // return nil, errors.Wrap(err, "write_model_field_failed") - //} - } else { + if err := writer.WriteField("model", redirectedModel); err != nil { + return nil, errors.Wrap(err, "write_model_field_failed") + } + } else {🤖 Prompt for AI Agents