Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion constant/channel.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,9 @@ const (
ChannelTypeJimeng = 51
ChannelTypeVidu = 52
ChannelTypeSubmodel = 53
ChannelTypeDoubaoVideo = 54
ChannelTypeDummy // this one is only for count, do not add any channel after this


)

var ChannelBaseURLs = []string{
Expand Down Expand Up @@ -111,4 +111,5 @@ var ChannelBaseURLs = []string{
"https://visual.volcengineapi.com", //51
"https://api.vidu.cn", //52
"https://llm.submodel.ai", //53
"https://ark.cn-beijing.volces.com", //54
}
6 changes: 6 additions & 0 deletions controller/channel-test.go
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ func testChannel(channel *model.Channel, testModel string, endpointType string)
newAPIError: nil,
}
}
if channel.Type == constant.ChannelTypeDoubaoVideo {
return testResult{
localErr: errors.New("doubao video channel test is not supported"),
newAPIError: nil,
}
}
if channel.Type == constant.ChannelTypeVidu {
return testResult{
localErr: errors.New("vidu channel test is not supported"),
Expand Down
86 changes: 86 additions & 0 deletions controller/task_video.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"one-api/relay"
"one-api/relay/channel"
relaycommon "one-api/relay/common"
"one-api/setting/ratio_setting"
"time"
)

Expand Down Expand Up @@ -120,6 +121,91 @@ func updateVideoSingleTask(ctx context.Context, adaptor channel.TaskAdaptor, cha
if !(len(taskResult.Url) > 5 && taskResult.Url[:5] == "data:") {
task.FailReason = taskResult.Url
}

// 如果返回了 total_tokens 并且配置了模型倍率(非固定价格),则重新计费
if taskResult.TotalTokens > 0 {
// 获取模型名称
var taskData map[string]interface{}
if err := json.Unmarshal(task.Data, &taskData); err == nil {
if modelName, ok := taskData["model"].(string); ok && modelName != "" {
// 获取模型价格和倍率
modelRatio, hasRatioSetting, _ := ratio_setting.GetModelRatio(modelName)

// 只有配置了倍率(非固定价格)时才按 token 重新计费
if hasRatioSetting && modelRatio > 0 {
// 获取用户和组的倍率信息
user, err := model.GetUserById(task.UserId, false)
if err == nil {
groupRatio := ratio_setting.GetGroupRatio(user.Group)
userGroupRatio, hasUserGroupRatio := ratio_setting.GetGroupGroupRatio(user.Group, user.Group)

var finalGroupRatio float64
if hasUserGroupRatio {
finalGroupRatio = userGroupRatio
} else {
finalGroupRatio = groupRatio
}

// 计算实际应扣费额度: totalTokens * modelRatio * groupRatio
actualQuota := int(float64(taskResult.TotalTokens) * modelRatio * finalGroupRatio)

// 计算差额
preConsumedQuota := task.Quota
quotaDelta := actualQuota - preConsumedQuota

if quotaDelta > 0 {
// 需要补扣费
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后补扣费:%s(实际消耗:%s,预扣费:%s,tokens:%d)",
task.TaskID,
logger.LogQuota(quotaDelta),
logger.LogQuota(actualQuota),
logger.LogQuota(preConsumedQuota),
taskResult.TotalTokens,
))
if err := model.DecreaseUserQuota(task.UserId, quotaDelta); err != nil {
logger.LogError(ctx, fmt.Sprintf("补扣费失败: %s", err.Error()))
} else {
model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
Comment on lines +168 to +169

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Stop incrementing request_count during post-settlement charges

UpdateUserUsedQuotaAndRequestCount always bumps the user’s request_count by 1. Calling it here makes every post-settlement adjustment look like an extra request, which corrupts usage statistics. Please switch to a helper that only adjusts the used quota (or add one) so this reconciliation doesn’t double-count requests.

🤖 Prompt for AI Agents
In controller/task_video.go around lines 168-169, the code calls
model.UpdateUserUsedQuotaAndRequestCount which increments the user request_count
and thereby double-counts post-settlement adjustments; replace this call with a
helper that only adjusts the user's used quota (e.g.,
model.UpdateUserUsedQuota(userId, quotaDelta)) or add such a function in the
model layer if it doesn't exist, then call that instead while keeping
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta) as-is so reconciliation
only updates quota and not request_count.

task.Quota = actualQuota // 更新任务记录的实际扣费额度

// 记录消费日志
logContent := fmt.Sprintf("视频任务成功补扣费,模型倍率 %.2f,分组倍率 %.2f,tokens %d,预扣费 %s,实际扣费 %s,补扣费 %s",
modelRatio, finalGroupRatio, taskResult.TotalTokens,
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(quotaDelta))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
}
} else if quotaDelta < 0 {
// 需要退还多扣的费用
refundQuota := -quotaDelta
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费后返还:%s(实际消耗:%s,预扣费:%s,tokens:%d)",
task.TaskID,
logger.LogQuota(refundQuota),
logger.LogQuota(actualQuota),
logger.LogQuota(preConsumedQuota),
taskResult.TotalTokens,
))
if err := model.IncreaseUserQuota(task.UserId, refundQuota, false); err != nil {
logger.LogError(ctx, fmt.Sprintf("退还预扣费失败: %s", err.Error()))
} else {
task.Quota = actualQuota // 更新任务记录的实际扣费额度

// 记录退款日志
logContent := fmt.Sprintf("视频任务成功退还多扣费用,模型倍率 %.2f,分组倍率 %.2f,tokens %d,预扣费 %s,实际扣费 %s,退还 %s",
modelRatio, finalGroupRatio, taskResult.TotalTokens,
logger.LogQuota(preConsumedQuota), logger.LogQuota(actualQuota), logger.LogQuota(refundQuota))
model.RecordLog(task.UserId, model.LogTypeSystem, logContent)
Comment on lines +188 to +197

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Refund must also rewind channel/user usage counters

When we refund refundQuota we only credit the user balance. The channel’s used quota and the user’s recorded used quota stay at the higher pre-charge value, so dashboards and reconciliation drift permanently. Mirror the charge path by subtracting refundQuota from the channel/user usage counters (e.g. UpdateChannelUsedQuota with a negative delta and an equivalent used-quota adjustment) so accounting remains consistent.

🤖 Prompt for AI Agents
In controller/task_video.go around lines 188 to 197, the refund only credits the
user balance but does not rewind the channel and user usage counters; update the
accounting by mirroring the original charge path — subtract refundQuota from the
channel's used quota (e.g., call the same UpdateChannelUsedQuota function with a
negative delta or equivalent) and decrement the user's recorded used quota by
refundQuota as well, perform these updates alongside IncreaseUserQuota so state
remains consistent, handle and log any errors from those updates and ensure the
operations are applied atomically or rolled back on failure.

}
} else {
// quotaDelta == 0, 预扣费刚好准确
logger.LogInfo(ctx, fmt.Sprintf("视频任务 %s 预扣费准确(%s,tokens:%d)",
task.TaskID, logger.LogQuota(actualQuota), taskResult.TotalTokens))
}
}
}
}
}
}
case model.TaskStatusFailure:
task.Status = model.TaskStatusFailure
task.Progress = "100%"
Expand Down
248 changes: 248 additions & 0 deletions relay/channel/task/doubao/adaptor.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,248 @@
package doubao

import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"one-api/constant"
"one-api/dto"
"one-api/model"
"one-api/relay/channel"
relaycommon "one-api/relay/common"
"one-api/service"

"github.com/gin-gonic/gin"
"github.com/pkg/errors"
)

// ============================
// Request / Response structures
// ============================

type ContentItem struct {
Type string `json:"type"` // "text" or "image_url"
Text string `json:"text,omitempty"` // for text type
ImageURL *ImageURL `json:"image_url,omitempty"` // for image_url type
}

type ImageURL struct {
URL string `json:"url"`
}

type requestPayload struct {
Model string `json:"model"`
Content []ContentItem `json:"content"`
}

type responsePayload struct {
ID string `json:"id"` // task_id
}

type responseTask struct {
ID string `json:"id"`
Model string `json:"model"`
Status string `json:"status"`
Content struct {
VideoURL string `json:"video_url"`
} `json:"content"`
Seed int `json:"seed"`
Resolution string `json:"resolution"`
Duration int `json:"duration"`
Ratio string `json:"ratio"`
FramesPerSecond int `json:"framespersecond"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
}

// ============================
// Adaptor implementation
// ============================

type TaskAdaptor struct {
ChannelType int
apiKey string
baseURL string
}

func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
a.ChannelType = info.ChannelType
a.baseURL = info.ChannelBaseUrl
a.apiKey = info.ApiKey
}

// ValidateRequestAndSetAction parses body, validates fields and sets default action.
func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
// Accept only POST /v1/video/generations as "generate" action.
return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionGenerate)
}

// BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
return fmt.Sprintf("%s/api/v3/contents/generations/tasks", a.baseURL), nil
}

// BuildRequestHeader sets required headers.
func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
req.Header.Set("Authorization", "Bearer "+a.apiKey)
return nil
}

// BuildRequestBody converts request into Doubao specific format.
func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) {
v, exists := c.Get("task_request")
if !exists {
return nil, fmt.Errorf("request not found in context")
}
req := v.(relaycommon.TaskSubmitReq)

body, err := a.convertToRequestPayload(&req)
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
data, err := json.Marshal(body)
if err != nil {
return nil, err
}
return bytes.NewReader(data), 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)
}

// DoResponse handles upstream response, returns taskID etc.
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
return
}
_ = resp.Body.Close()

// Parse Doubao response
var dResp responsePayload
if err := json.Unmarshal(responseBody, &dResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}

if dResp.ID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}

c.JSON(http.StatusOK, gin.H{"task_id": dResp.ID})
return dResp.ID, responseBody, nil
}
Comment on lines +123 to +145

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Handle non-2xx Doubao responses before parsing success payload

When Doubao returns a 4xx/5xx (e.g., auth failure or validation error), we still unmarshal into responsePayload. Because id is empty, this path always bubbles up as our own 500/"task_id is empty", masking the original failure and misreporting the status. Please gate the success flow on resp.StatusCode so we surface upstream errors accurately and keep retries/backoffs sane.

@@
-	// Parse Doubao response
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		taskErr = service.TaskErrorWrapper(
+			fmt.Errorf("upstream status %d: %s", resp.StatusCode, string(responseBody)),
+			"upstream_request_failed",
+			resp.StatusCode,
+		)
+		return
+	}
+
+	// Parse Doubao response
 	var dResp responsePayload
📝 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.

Suggested change
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
return
}
_ = resp.Body.Close()
// Parse Doubao response
var dResp responsePayload
if err := json.Unmarshal(responseBody, &dResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
if dResp.ID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}
c.JSON(http.StatusOK, gin.H{"task_id": dResp.ID})
return dResp.ID, responseBody, nil
}
func (a *TaskAdaptor) DoResponse(
c *gin.Context,
resp *http.Response,
info *relaycommon.RelayInfo,
) (taskID string, taskData []byte, taskErr *dto.TaskError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
return
}
_ = resp.Body.Close()
// Handle non-2xx upstream responses
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
taskErr = service.TaskErrorWrapper(
fmt.Errorf("upstream status %d: %s", resp.StatusCode, string(responseBody)),
"upstream_request_failed",
resp.StatusCode,
)
return
}
// Parse Doubao response
var dResp responsePayload
if err := json.Unmarshal(responseBody, &dResp); err != nil {
taskErr = service.TaskErrorWrapper(
errors.Wrapf(err, "body: %s", responseBody),
"unmarshal_response_body_failed",
http.StatusInternalServerError,
)
return
}
if dResp.ID == "" {
taskErr = service.TaskErrorWrapper(
fmt.Errorf("task_id is empty"),
"invalid_response",
http.StatusInternalServerError,
)
return
}
c.JSON(http.StatusOK, gin.H{"task_id": dResp.ID})
return dResp.ID, responseBody, nil
}
🤖 Prompt for AI Agents
In relay/channel/task/doubao/adaptor.go around lines 123 to 145, the code
currently unmarshals the response body and treats missing dResp.ID as an
internal 500 error even when Doubao returned a non-2xx (4xx/5xx) error; change
the flow to first check resp.StatusCode and, for non-2xx responses, wrap and
return the original status and body as a TaskError (including status code and
response body text) without treating it as a missing task_id; only on 2xx should
you unmarshal into responsePayload and proceed to return the task ID and body,
ensuring resp.Body is closed and errors from ReadAll/Unmarshal are still wrapped
appropriately.


// FetchTask fetch task status
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any) (*http.Response, error) {
taskID, ok := body["task_id"].(string)
if !ok {
return nil, fmt.Errorf("invalid task_id")
}

uri := fmt.Sprintf("%s/api/v3/contents/generations/tasks/%s", baseUrl, taskID)

req, err := http.NewRequest(http.MethodGet, uri, nil)
if err != nil {
return nil, err
}

req.Header.Set("Accept", "application/json")
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+key)

return service.GetHttpClient().Do(req)
}

func (a *TaskAdaptor) GetModelList() []string {
return ModelList
}

func (a *TaskAdaptor) GetChannelName() string {
return ChannelName
}

func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
r := requestPayload{
Model: req.Model,
Content: []ContentItem{},
}

// Add text prompt
if req.Prompt != "" {
r.Content = append(r.Content, ContentItem{
Type: "text",
Text: req.Prompt,
})
}

// Add images if present
if req.HasImage() {
for _, imgURL := range req.Images {
r.Content = append(r.Content, ContentItem{
Type: "image_url",
ImageURL: &ImageURL{
URL: imgURL,
},
})
}
}

// TODO: Add support for additional parameters from metadata
// such as ratio, duration, seed, etc.
// metadata := req.Metadata
// if metadata != nil {
// // Parse and apply metadata parameters
// }

return &r, nil
}

func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}

taskResult := relaycommon.TaskInfo{
Code: 0,
}

// Map Doubao status to internal status
switch resTask.Status {
case "pending", "queued":
taskResult.Status = model.TaskStatusQueued
taskResult.Progress = "10%"
case "processing":
taskResult.Status = model.TaskStatusInProgress
taskResult.Progress = "50%"
case "succeeded":
taskResult.Status = model.TaskStatusSuccess
taskResult.Progress = "100%"
taskResult.Url = resTask.Content.VideoURL
// 解析 usage 信息用于按倍率计费
taskResult.CompletionTokens = resTask.Usage.CompletionTokens
taskResult.TotalTokens = resTask.Usage.TotalTokens
case "failed":
taskResult.Status = model.TaskStatusFailure
taskResult.Progress = "100%"
taskResult.Reason = "task failed"
default:
// Unknown status, treat as processing
taskResult.Status = model.TaskStatusInProgress
taskResult.Progress = "30%"
}

return &taskResult, nil
}
9 changes: 9 additions & 0 deletions relay/channel/task/doubao/constants.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
package doubao

var ModelList = []string{
"doubao-seedance-1-0-pro-250528",
"doubao-seedance-1-0-lite-t2v",
"doubao-seedance-1-0-lite-i2v",
}

var ChannelName = "doubao-video"
Loading