-
-
Notifications
You must be signed in to change notification settings - Fork 2.3k
feat: 增加即梦生图功能 #1363
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
feat: 增加即梦生图功能 #1363
Conversation
|
""" WalkthroughThis change introduces full relay support for the Jimeng channel, including new API and channel type constants, a dedicated adaptor with request/response conversion, signing logic, image handling, and integration into the relay adaptor selection. It adds all necessary code to enable Jimeng as a new channel within the relay framework. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant GinContext
participant JimengAdaptor
participant JimengAPI
participant RelayInfo
Client->>GinContext: Sends OpenAI/Image request
GinContext->>JimengAdaptor: Convert request (ConvertOpenAIRequest/ConvertImageRequest)
JimengAdaptor->>JimengAdaptor: Build Jimeng payload
JimengAdaptor->>JimengAPI: DoRequest (signed HTTP request)
JimengAPI-->>JimengAdaptor: Returns response
JimengAdaptor->>JimengAdaptor: Handle response (DoResponse/jimengImageHandler)
JimengAdaptor-->>GinContext: Return OpenAI-compatible response
GinContext-->>Client: Respond with result
Possibly related PRs
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
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.
Actionable comments posted: 7
🧹 Nitpick comments (4)
relay/channel/jimeng/sign.go (2)
21-38: Remove commented-out codeThis commented-out function appears to be an alternative implementation that's no longer needed. Please remove it to improve code maintainability.
-// SignRequestForJimeng 对即梦 API 请求进行签名,支持 http.Request 或 header+url+body 方式 -//func SignRequestForJimeng(req *http.Request, accessKey, secretKey string) error { -// var bodyBytes []byte -// var err error -// -// if req.Body != nil { -// bodyBytes, err = io.ReadAll(req.Body) -// if err != nil { -// return fmt.Errorf("read request body failed: %w", err) -// } -// _ = req.Body.Close() -// req.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) // rewind -// } else { -// bodyBytes = []byte{} -// } -// -// return signJimengHeaders(&req.Header, req.Method, req.URL, bodyBytes, accessKey, secretKey) -//} -
146-147: Consider making region and service name configurableHard-coded values for region and service name limit flexibility. Consider making these configurable through the relay info or environment variables.
- region := "cn-north-1" - serviceName := "cv" + // Consider reading from configuration + region := getConfigValue(c, "JIMENG_REGION", "cn-north-1") + serviceName := getConfigValue(c, "JIMENG_SERVICE", "cv")relay/channel/jimeng/adaptor.go (2)
20-22: Consider implementing a base adaptor with default "not implemented" errorsMultiple methods return "not implemented" errors. Consider creating a base adaptor struct that provides these default implementations to reduce boilerplate.
// In a base package: type BaseAdaptor struct{} func (a *BaseAdaptor) ConvertClaudeRequest(*gin.Context, *relaycommon.RelayInfo, *dto.ClaudeRequest) (any, error) { return nil, errors.New("not implemented") } // ... other methods // Then embed in Jimeng adaptor: type Adaptor struct { channel.BaseAdaptor }Also applies to: 32-33, 82-84, 86-88, 90-92, 94-96
24-25: Document or implement the Init methodThe Init method is empty. If initialization is not required, consider adding a comment explaining why, or implement any necessary initialization logic.
func (a *Adaptor) Init(info *relaycommon.RelayInfo) { + // No initialization required for Jimeng adaptor }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
common/api_type.go(1 hunks)constant/api_type.go(1 hunks)relay/channel/api_request.go(1 hunks)relay/channel/jimeng/adaptor.go(1 hunks)relay/channel/jimeng/constants.go(1 hunks)relay/channel/jimeng/image.go(1 hunks)relay/channel/jimeng/sign.go(1 hunks)relay/relay_adaptor.go(4 hunks)
🧰 Additional context used
🧠 Learnings (1)
relay/channel/api_request.go (2)
Learnt from: feitianbubu
PR: QuantumNous/new-api#1228
File: router/main.go:28-36
Timestamp: 2025-06-15T12:38:11.806Z
Learning: gin.Context implements context.Context interface since Gin v1.8.0, providing the methods Deadline(), Done(), Err(), and Value(). When using Gin v1.8.0 or later, gin.Context can be passed directly to functions expecting context.Context without needing to extract c.Request.Context().
Learnt from: feitianbubu
PR: QuantumNous/new-api#1228
File: router/main.go:28-36
Timestamp: 2025-06-15T12:38:11.806Z
Learning: gin.Context implements context.Context interface since Gin v1.8.0, providing the methods Deadline(), Done(), Err(), and Value(). When using Gin v1.8.0 or later, gin.Context can be passed directly to functions expecting context.Context without needing to extract c.Request.Context().
🧬 Code Graph Analysis (3)
relay/channel/api_request.go (1)
relay/common/relay_info.go (1)
RelayInfo(62-112)
common/api_type.go (2)
constant/channel.go (1)
ChannelTypeJimeng(51-51)constant/api_type.go (1)
APITypeJimeng(33-33)
relay/relay_adaptor.go (4)
constant/api_type.go (1)
APITypeJimeng(33-33)relay/channel/jimeng/adaptor.go (1)
Adaptor(17-18)relay/channel/adapter.go (2)
Adaptor(13-29)TaskAdaptor(31-50)relay/channel/task/jimeng/adaptor.go (1)
TaskAdaptor(70-75)
🔇 Additional comments (8)
constant/api_type.go (1)
33-33: LGTM! Consistent constant addition.The new
APITypeJimengconstant follows the established enumeration pattern and is correctly positioned beforeAPITypeDummy.common/api_type.go (1)
66-67: LGTM! Consistent channel-to-API type mapping.The new mapping from
ChannelTypeJimengtoAPITypeJimengfollows the established pattern and correctly integrates the Jimeng channel into the type system.relay/channel/api_request.go (1)
206-208: LGTM! Clean public interface addition.The new
DoRequestfunction provides a clean public wrapper around the existingdoRequestfunctionality, enabling channel adaptors to use standardized request handling.relay/channel/jimeng/constants.go (1)
1-9: LGTM! Well-structured constants definition.The constants file correctly defines the channel name and supported model list for the Jimeng channel, following established patterns and matching the PR objectives.
relay/relay_adaptor.go (4)
18-18: LGTM! Proper import addition.The import for the Jimeng channel adaptor is correctly added and follows the established import pattern.
27-27: LGTM! Clean aliased import.The aliased import for the task adaptor (
taskjimeng) properly avoids naming conflicts while maintaining clarity.
97-98: LGTM! Consistent adaptor registration.The new case for
APITypeJimengcorrectly returns ajimeng.Adaptorinstance, following the established pattern for adaptor registration.
112-112: LGTM! Proper task adaptor registration.The case for
TaskPlatformJimengcorrectly returns ataskjimeng.TaskAdaptorinstance, maintaining consistency with the task adaptor pattern.
| if err != nil { | ||
| return err | ||
| } | ||
| common.LogInfo(c, fmt.Sprintf("SetPayloadHash body: %s", body)) |
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.
Potential security concern: Avoid logging request body
Logging the entire request body could expose sensitive information like API keys, personal data, or other confidential content.
- common.LogInfo(c, fmt.Sprintf("SetPayloadHash body: %s", body))
+ // Only log non-sensitive metadata if debugging is needed
+ common.LogDebug(c, fmt.Sprintf("SetPayloadHash body size: %d bytes", len(body)))🤖 Prompt for AI Agents
In relay/channel/jimeng/sign.go at line 47, avoid logging the entire request
body to prevent exposing sensitive information such as API keys or personal
data. Instead, remove or redact the body content from the log message, or log
only non-sensitive metadata or a hash of the body if needed for debugging.
| payloadHash := sha256.Sum256(bodyBytes) | ||
| hexPayloadHash := hex.EncodeToString(payloadHash[:]) | ||
|
|
||
| method := c.Request.Method |
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.
Use consistent method source
The function uses c.Request.Method instead of req.Method, which is inconsistent since the req parameter is the actual request being signed.
- method := c.Request.Method
+ method := req.Method📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| method := c.Request.Method | |
| method := req.Method |
🤖 Prompt for AI Agents
In relay/channel/jimeng/sign.go at line 75, replace the usage of
c.Request.Method with req.Method to maintain consistency, as req is the actual
request being signed. This ensures the method is obtained directly from the
request parameter instead of the context.
relay/channel/jimeng/image.go
Outdated
| return service.OpenAIErrorWrapper(err, "write_response_body_failed", http.StatusInternalServerError), nil | ||
| } | ||
|
|
||
| return nil, &dto.Usage{} |
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.
🛠️ Refactor suggestion
Track actual usage for billing purposes
Returning an empty usage struct means image generation usage won't be tracked. Consider calculating actual usage based on the number of images generated.
- return nil, &dto.Usage{}
+ // Track usage based on number of images generated
+ usage := &dto.Usage{
+ PromptTokens: 1, // Image generation counts as 1 request
+ TotalTokens: len(jimengResponse.Data.ImageUrls) + len(jimengResponse.Data.BinaryDataBase64),
+ }
+ return nil, usage🤖 Prompt for AI Agents
In relay/channel/jimeng/image.go at line 90, the function currently returns an
empty usage struct which prevents tracking image generation usage for billing.
Modify the code to calculate the actual usage based on the number of images
generated and return this usage data instead of an empty struct to ensure
accurate billing tracking.
| for _, base64Data := range response.Data.BinaryDataBase64 { | ||
| imageResponse.Data = append(imageResponse.Data, dto.ImageData{ | ||
| B64Json: base64Data, | ||
| }) | ||
| } | ||
| for _, imageUrl := range response.Data.ImageUrls { | ||
| imageResponse.Data = append(imageResponse.Data, dto.ImageData{ | ||
| Url: imageUrl, | ||
| }) | ||
| } |
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.
Potential data duplication in response conversion
When both base64 data and URLs exist in the response, they are appended to the same Data slice, which could lead to unexpected behavior if the client expects either format but not both mixed together.
Consider handling based on the original request format:
- for _, base64Data := range response.Data.BinaryDataBase64 {
- imageResponse.Data = append(imageResponse.Data, dto.ImageData{
- B64Json: base64Data,
- })
- }
- for _, imageUrl := range response.Data.ImageUrls {
- imageResponse.Data = append(imageResponse.Data, dto.ImageData{
- Url: imageUrl,
- })
- }
+ // Prefer URLs if available, otherwise use base64
+ if len(response.Data.ImageUrls) > 0 {
+ for _, imageUrl := range response.Data.ImageUrls {
+ imageResponse.Data = append(imageResponse.Data, dto.ImageData{
+ Url: imageUrl,
+ })
+ }
+ } else if len(response.Data.BinaryDataBase64) > 0 {
+ for _, base64Data := range response.Data.BinaryDataBase64 {
+ imageResponse.Data = append(imageResponse.Data, dto.ImageData{
+ B64Json: base64Data,
+ })
+ }
+ }🤖 Prompt for AI Agents
In relay/channel/jimeng/image.go around lines 35 to 44, the code appends both
base64 data and image URLs into the same Data slice, which may cause data
duplication or confusion if the client expects only one format. Modify the logic
to check the original request format and append only the corresponding data type
(either base64 or URLs) to imageResponse.Data, ensuring the response matches the
requested format and avoids mixing both types.
| } | ||
|
|
||
| // Check if the response indicates an error | ||
| if jimengResponse.Code != 10000 { |
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.
🛠️ Refactor suggestion
Consider making success code configurable
The hard-coded success code 10000 should be defined as a constant for better maintainability.
+const JimengSuccessCode = 10000
+
// Check if the response indicates an error
- if jimengResponse.Code != 10000 {
+ if jimengResponse.Code != JimengSuccessCode {🤖 Prompt for AI Agents
In relay/channel/jimeng/image.go at line 64, the success code 10000 is
hard-coded. Define this success code as a constant at the top of the file or in
a relevant constants section, then replace the hard-coded value with this
constant to improve maintainability and readability.
| if len(request.ExtraFields) > 0 { | ||
| if err := json.Unmarshal(request.ExtraFields, &payload); err != nil { | ||
| return nil, fmt.Errorf("failed to unmarshal extra fields: %w", err) | ||
| } | ||
| } |
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.
Validate ExtraFields unmarshaling to prevent field overwrites
Unmarshaling ExtraFields directly into the payload could overwrite critical fields like ReqKey and Prompt that were already set.
if len(request.ExtraFields) > 0 {
- if err := json.Unmarshal(request.ExtraFields, &payload); err != nil {
+ // Create a temporary struct to validate extra fields
+ var extras imageRequestPayload
+ if err := json.Unmarshal(request.ExtraFields, &extras); err != nil {
return nil, fmt.Errorf("failed to unmarshal extra fields: %w", err)
}
+ // Selectively apply extra fields without overwriting critical ones
+ if extras.Seed != 0 {
+ payload.Seed = extras.Seed
+ }
+ if extras.Width != 0 {
+ payload.Width = extras.Width
+ }
+ if extras.Height != 0 {
+ payload.Height = extras.Height
+ }
+ // Apply other non-critical fields...
+ payload.UsePreLLM = extras.UsePreLLM
+ payload.UseSR = extras.UseSR
+ payload.LogoInfo = extras.LogoInfo
+ payload.ImageUrls = extras.ImageUrls
+ payload.BinaryData = extras.BinaryData
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if len(request.ExtraFields) > 0 { | |
| if err := json.Unmarshal(request.ExtraFields, &payload); err != nil { | |
| return nil, fmt.Errorf("failed to unmarshal extra fields: %w", err) | |
| } | |
| } | |
| if len(request.ExtraFields) > 0 { | |
| // Create a temporary struct to validate extra fields | |
| var extras imageRequestPayload | |
| if err := json.Unmarshal(request.ExtraFields, &extras); err != nil { | |
| return nil, fmt.Errorf("failed to unmarshal extra fields: %w", err) | |
| } | |
| // Selectively apply extra fields without overwriting critical ones | |
| if extras.Seed != 0 { | |
| payload.Seed = extras.Seed | |
| } | |
| if extras.Width != 0 { | |
| payload.Width = extras.Width | |
| } | |
| if extras.Height != 0 { | |
| payload.Height = extras.Height | |
| } | |
| // Apply other non-critical fields... | |
| payload.UsePreLLM = extras.UsePreLLM | |
| payload.UseSR = extras.UseSR | |
| payload.LogoInfo = extras.LogoInfo | |
| payload.ImageUrls = extras.ImageUrls | |
| payload.BinaryData = extras.BinaryData | |
| } |
🤖 Prompt for AI Agents
In relay/channel/jimeng/adaptor.go around lines 73 to 77, unmarshaling
request.ExtraFields directly into the payload risks overwriting important fields
like ReqKey and Prompt. To fix this, unmarshal ExtraFields into a separate
temporary map or struct first, then selectively merge only non-conflicting
fields into the payload to avoid overwriting existing critical data.
| Width int `json:"width,omitempty"` // Image width, default 512, range [256, 768] | ||
| Height int `json:"height,omitempty"` // Image height, default 512, range [256, 768] |
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.
🛠️ Refactor suggestion
Add validation for image dimensions
The comments indicate width and height have valid ranges [256, 768], but there's no validation implemented.
Add validation in the ConvertImageRequest method:
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
payload := imageRequestPayload{
ReqKey: request.Model,
Prompt: request.Prompt,
}
+
+ // Apply and validate dimensions if provided
+ if request.Width != nil {
+ if *request.Width < 256 || *request.Width > 768 {
+ return nil, fmt.Errorf("width must be between 256 and 768, got %d", *request.Width)
+ }
+ payload.Width = *request.Width
+ }
+ if request.Height != nil {
+ if *request.Height < 256 || *request.Height > 768 {
+ return nil, fmt.Errorf("height must be between 256 and 768, got %d", *request.Height)
+ }
+ payload.Height = *request.Height
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Width int `json:"width,omitempty"` // Image width, default 512, range [256, 768] | |
| Height int `json:"height,omitempty"` // Image height, default 512, range [256, 768] | |
| func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) { | |
| payload := imageRequestPayload{ | |
| ReqKey: request.Model, | |
| Prompt: request.Prompt, | |
| } | |
| // Apply and validate dimensions if provided | |
| if request.Width != nil { | |
| if *request.Width < 256 || *request.Width > 768 { | |
| return nil, fmt.Errorf("width must be between 256 and 768, got %d", *request.Width) | |
| } | |
| payload.Width = *request.Width | |
| } | |
| if request.Height != nil { | |
| if *request.Height < 256 || *request.Height > 768 { | |
| return nil, fmt.Errorf("height must be between 256 and 768, got %d", *request.Height) | |
| } | |
| payload.Height = *request.Height | |
| } | |
| // ...rest of your method... | |
| } |
🤖 Prompt for AI Agents
In relay/channel/jimeng/adaptor.go around lines 54 to 55, the Width and Height
fields have specified valid ranges [256, 768] in comments but lack validation.
Add validation logic in the ConvertImageRequest method to check if Width and
Height values fall within the range 256 to 768 inclusive. If values are outside
this range, set them to the default value 512 or return an error as appropriate
to ensure only valid image dimensions are processed.
已转化成OpenAI兼容格式, 即梦生图目前只支持固定
jimeng_high_aes_general_v21_L, 兼容格式将这个做为model即梦的额外官方格式放在extra_fields字段, New API会直接透传到官方接口, 理论上支持官方所有参数
请求:
{
"model": "jimeng_high_aes_general_v21_L",
"prompt": "可爱的中国小女孩在花园里玩耍",
"extra_fields": {
"width": 768,
"height": 512,
"seed": -1,
"use_pre_llm": true,
"use_sr": true,
"logo_info": {
"add_logo": true,
"position": 1,
"opacity": 0.4
}
}
}
返回:
{
"data": [
{
"url": "https://p9-aiop-sign.byteimg.com/tos-cn-i-vuqhorh59i/202507131312116C0111649BC99A3D5F31-0~tplv-vuqhorh59i-image.image?rk3s=7f9e702d&x-expires=1752469939&x-signature=NkDFlYSzS3rWHnTVBI8TxiEuHDA%3D",
"b64_json": "",
"revised_prompt": ""
}
],
"created": 1752383531
}
Summary by CodeRabbit
New Features
Bug Fixes
Chores