feat: add model redirection for task requests - #2978
Conversation
# Conflicts: # relay/channel/task/sora/adaptor.go # relay/common/relay_utils.go
WalkthroughThis PR adds model-mapping support to the Sora channel adaptor. It introduces conditional request body handling based on whether a model is mapped, validates task actions based on image presence and channel type, initializes model mapping via a helper function, and persists upstream and origin model names in task properties. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant RelayTask
participant Validator as Validator<br/>(relay_utils)
participant ModelMapper as ModelMapper<br/>(helper)
participant Adaptor as Sora Adaptor
participant Storage
Client->>RelayTask: Request with model
RelayTask->>Validator: ValidateBasicTaskRequest()
Validator->>Validator: Determine action from images
Validator->>Validator: Check IsModelMapped
Validator-->>RelayTask: Validation complete
RelayTask->>ModelMapper: ModelMappedHelper()
ModelMapper->>ModelMapper: Map model if needed
ModelMapper-->>RelayTask: Mapped model info
RelayTask->>RelayTask: Set task.Properties<br/>(UpstreamModelName,<br/>OriginModelName)
RelayTask->>Storage: Insert task
RelayTask->>Adaptor: BuildRequestBody()
Adaptor->>Adaptor: Check IsModelMapped
Adaptor->>Adaptor: Parse content-type
alt multipart/form-data
Adaptor->>Adaptor: buildRequestBodyWithMappedModel()
Adaptor->>Adaptor: Rebuild multipart with<br/>mapped model
else application/json
Adaptor->>Adaptor: Unmarshal, preserve payload,<br/>re-marshal
end
Adaptor-->>RelayTask: Rebuilt request body
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
relay/channel/task/sora/adaptor.go (1)
152-217:buildRequestBodyWithMappedModelhas duplicate logic for file vs non-file parts.Lines 188–204: the
if part.FileName() != ""andelsebranches both do exactly the same thing —writer.CreatePart(part.Header)followed byio.Copy. The file-name check is redundant here sinceCreatePartalready preserves headers (includingContent-Dispositionwith filename).♻️ Simplify by removing the redundant branch
} 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") - } - } + newPart, err := writer.CreatePart(part.Header) + if err != nil { + return nil, errors.Wrap(err, "create_part_failed") + } + if _, err := io.Copy(newPart, part); err != nil { + return nil, errors.Wrap(err, "copy_part_content_failed") + } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@relay/channel/task/sora/adaptor.go` around lines 152 - 217, The function buildRequestBodyWithMappedModel contains redundant branching that checks part.FileName() and executes identical logic in both branches; simplify by removing the file-vs-non-file if/else and always create a new part with writer.CreatePart(part.Header) followed by io.Copy from part, ensuring error wraps remain (use the existing error messages like "create_form_part_failed"/"copy_part_content_failed" or consolidate to the existing ones), which preserves file headers (including filename) while eliminating duplicate code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@relay/channel/task/sora/adaptor.go`:
- Line 149: The returned reader may be exhausted because storage.Bytes() was
already called earlier; replace the final return of common.ReaderOnly(storage)
(in the branch reached when info.IsModelMapped is true) with a reader
constructed from the already-read bytes buffer produced by storage.Bytes() (use
the variable holding those bytes instead of re-wrapping storage), so the
upstream receives the actual body for unrecognized content types instead of an
empty reader.
- Around line 119-149: When info.IsModelMapped is true the multipart path calls
buildRequestBodyWithMappedModel which currently drops the "model" form field
(the WriteField call for fieldName == "model" is commented out), causing data
loss; fix by preserving the original "model" part when not rewriting (or by
performing the actual rewrite) inside buildRequestBodyWithMappedModel so that
the "model" field is written into the new multipart body (i.e., restore or
replace the WriteField behavior for fieldName == "model"); also avoid
unnecessary JSON unmarshal/remarshal in the JSON branch of the handler (the
jsonData["model"] rewrite is commented out) — either perform the model rewrite
there or simply return the original body (bytes.NewReader(bodyBytes)) when no
change is needed.
---
Nitpick comments:
In `@relay/channel/task/sora/adaptor.go`:
- Around line 152-217: The function buildRequestBodyWithMappedModel contains
redundant branching that checks part.FileName() and executes identical logic in
both branches; simplify by removing the file-vs-non-file if/else and always
create a new part with writer.CreatePart(part.Header) followed by io.Copy from
part, ensuring error wraps remain (use the existing error messages like
"create_form_part_failed"/"copy_part_content_failed" or consolidate to the
existing ones), which preserves file headers (including filename) while
eliminating duplicate code.
| 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 |
There was a problem hiding this comment.
Model rewrite is commented out, but the multipart path silently drops the model field — this is a data-loss bug.
When info.IsModelMapped is true and the content type is multipart/form-data, the code enters buildRequestBodyWithMappedModel. Inside that helper, when fieldName == "model" (Line 180), the WriteField call is commented out (Lines 183–185) and the else branch that preserves the part is skipped. The result: the model field 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 IsModelMapped is false (or rewrite is disabled), return the body unchanged — which is already done at Line 119–122. The current code reaches Lines 127+ only when IsModelMapped is true, 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
Verify each finding against the current code and only fix it if needed.
In `@relay/channel/task/sora/adaptor.go` around lines 119 - 149, When
info.IsModelMapped is true the multipart path calls
buildRequestBodyWithMappedModel which currently drops the "model" form field
(the WriteField call for fieldName == "model" is commented out), causing data
loss; fix by preserving the original "model" part when not rewriting (or by
performing the actual rewrite) inside buildRequestBodyWithMappedModel so that
the "model" field is written into the new multipart body (i.e., restore or
replace the WriteField behavior for fieldName == "model"); also avoid
unnecessary JSON unmarshal/remarshal in the JSON branch of the handler (the
jsonData["model"] rewrite is commented out) — either perform the model rewrite
there or simply return the original body (bytes.NewReader(bodyBytes)) when no
change is needed.
| return bytes.NewReader(newBody), nil | ||
| } | ||
|
|
||
| return common.ReaderOnly(storage), nil |
There was a problem hiding this comment.
Potential stale reader: common.ReaderOnly(storage) may return an exhausted reader.
Line 113 calls storage.Bytes() which may consume the underlying reader. Line 149 then returns common.ReaderOnly(storage) for non-multipart/non-JSON content types. If storage doesn't support re-reading after Bytes() is called, this will return an empty body.
This line is only reached when info.IsModelMapped is true (the !IsModelMapped early return is at Line 119), so for mapped models with an unrecognized content type, the upstream would receive an empty body.
Proposed fix: use the already-read bytes
- return common.ReaderOnly(storage), nil
+ return bytes.NewReader(bodyBytes), nil🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@relay/channel/task/sora/adaptor.go` at line 149, The returned reader may be
exhausted because storage.Bytes() was already called earlier; replace the final
return of common.ReaderOnly(storage) (in the branch reached when
info.IsModelMapped is true) with a reader constructed from the already-read
bytes buffer produced by storage.Bytes() (use the variable holding those bytes
instead of re-wrapping storage), so the upstream receives the actual body for
unrecognized content types instead of an empty reader.
3 participants
fix #2130
Summary by CodeRabbit